blob: 8dd3a65f43de9ced48caa18054a8d644de437223 [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
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000201 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000202 if (igo == NULL)
203 return NULL;
204 igo->parent = (PyObject *)parent;
205 Py_INCREF(parent);
206 igo->tgtkey = tgtkey;
207 Py_INCREF(tgtkey);
208
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000209 PyObject_GC_Track(igo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000210 return (PyObject *)igo;
211}
212
213static void
214_grouper_dealloc(_grouperobject *igo)
215{
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000216 PyObject_GC_UnTrack(igo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000217 Py_DECREF(igo->parent);
218 Py_DECREF(igo->tgtkey);
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000219 PyObject_GC_Del(igo);
220}
221
222static int
223_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
224{
225 Py_VISIT(igo->parent);
226 Py_VISIT(igo->tgtkey);
227 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000228}
229
230static PyObject *
231_grouper_next(_grouperobject *igo)
232{
233 groupbyobject *gbo = (groupbyobject *)igo->parent;
234 PyObject *newvalue, *newkey, *r;
235 int rcmp;
236
237 if (gbo->currvalue == NULL) {
238 newvalue = PyIter_Next(gbo->it);
239 if (newvalue == NULL)
240 return NULL;
241
242 if (gbo->keyfunc == Py_None) {
243 newkey = newvalue;
244 Py_INCREF(newvalue);
245 } else {
246 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
247 newvalue, NULL);
248 if (newkey == NULL) {
249 Py_DECREF(newvalue);
250 return NULL;
251 }
252 }
253
254 assert(gbo->currkey == NULL);
255 gbo->currkey = newkey;
256 gbo->currvalue = newvalue;
257 }
258
259 assert(gbo->currkey != NULL);
260 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
261 if (rcmp <= 0)
262 /* got any error or current group is end */
263 return NULL;
264
265 r = gbo->currvalue;
266 gbo->currvalue = NULL;
Raymond Hettinger75ccea32004-09-01 07:02:44 +0000267 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000268
269 return r;
270}
271
272static PyTypeObject _grouper_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000273 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000274 "itertools._grouper", /* tp_name */
275 sizeof(_grouperobject), /* tp_basicsize */
276 0, /* tp_itemsize */
277 /* methods */
278 (destructor)_grouper_dealloc, /* tp_dealloc */
279 0, /* tp_print */
280 0, /* tp_getattr */
281 0, /* tp_setattr */
282 0, /* tp_compare */
283 0, /* tp_repr */
284 0, /* tp_as_number */
285 0, /* tp_as_sequence */
286 0, /* tp_as_mapping */
287 0, /* tp_hash */
288 0, /* tp_call */
289 0, /* tp_str */
290 PyObject_GenericGetAttr, /* tp_getattro */
291 0, /* tp_setattro */
292 0, /* tp_as_buffer */
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000293 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000294 0, /* tp_doc */
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000295 (traverseproc)_grouper_traverse,/* tp_traverse */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000296 0, /* tp_clear */
297 0, /* tp_richcompare */
298 0, /* tp_weaklistoffset */
299 PyObject_SelfIter, /* tp_iter */
300 (iternextfunc)_grouper_next, /* tp_iternext */
301 0, /* tp_methods */
302 0, /* tp_members */
303 0, /* tp_getset */
304 0, /* tp_base */
305 0, /* tp_dict */
306 0, /* tp_descr_get */
307 0, /* tp_descr_set */
308 0, /* tp_dictoffset */
309 0, /* tp_init */
310 0, /* tp_alloc */
311 0, /* tp_new */
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000312 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000313};
314
315
316
Raymond Hettingerad983e72003-11-12 14:32:26 +0000317/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000318
Raymond Hettingerad983e72003-11-12 14:32:26 +0000319/* The teedataobject pre-allocates space for LINKCELLS number of objects.
320 To help the object fit neatly inside cache lines (space for 16 to 32
321 pointers), the value should be a multiple of 16 minus space for
322 the other structure members including PyHEAD overhead. The larger the
323 value, the less memory overhead per object and the less time spent
324 allocating/deallocating new links. The smaller the number, the less
325 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000326*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000327#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000328
329typedef struct {
330 PyObject_HEAD
331 PyObject *it;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000332 int numread;
333 PyObject *nextlink;
334 PyObject *(values[LINKCELLS]);
335} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000336
337typedef struct {
338 PyObject_HEAD
Raymond Hettingerad983e72003-11-12 14:32:26 +0000339 teedataobject *dataobj;
340 int index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000341 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000342} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000343
Raymond Hettingerad983e72003-11-12 14:32:26 +0000344static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000345
346static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000347teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000348{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000349 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000350
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000351 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000352 if (tdo == NULL)
Raymond Hettinger45143692003-10-25 06:37:47 +0000353 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000354
355 tdo->numread = 0;
356 tdo->nextlink = NULL;
357 Py_INCREF(it);
358 tdo->it = it;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000359 PyObject_GC_Track(tdo);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000360 return (PyObject *)tdo;
361}
362
363static PyObject *
364teedataobject_jumplink(teedataobject *tdo)
365{
366 if (tdo->nextlink == NULL)
367 tdo->nextlink = teedataobject_new(tdo->it);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000368 Py_XINCREF(tdo->nextlink);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000369 return tdo->nextlink;
370}
371
372static PyObject *
373teedataobject_getitem(teedataobject *tdo, int i)
374{
375 PyObject *value;
376
377 assert(i < LINKCELLS);
378 if (i < tdo->numread)
379 value = tdo->values[i];
380 else {
381 /* this is the lead iterator, so fetch more data */
382 assert(i == tdo->numread);
383 value = PyIter_Next(tdo->it);
384 if (value == NULL)
385 return NULL;
386 tdo->numread++;
387 tdo->values[i] = value;
Raymond Hettinger45143692003-10-25 06:37:47 +0000388 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000389 Py_INCREF(value);
390 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000391}
392
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000393static int
394teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
395{
396 int i;
397 Py_VISIT(tdo->it);
398 for (i = 0; i < tdo->numread; i++)
399 Py_VISIT(tdo->values[i]);
400 Py_VISIT(tdo->nextlink);
401 return 0;
402}
403
404static int
405teedataobject_clear(teedataobject *tdo)
406{
407 int i;
408 Py_CLEAR(tdo->it);
409 for (i=0 ; i<tdo->numread ; i++)
410 Py_CLEAR(tdo->values[i]);
411 Py_CLEAR(tdo->nextlink);
412 return 0;
413}
414
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000416teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417{
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000418 PyObject_GC_UnTrack(tdo);
419 teedataobject_clear(tdo);
420 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000421}
422
Raymond Hettingerad983e72003-11-12 14:32:26 +0000423PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000424
Raymond Hettingerad983e72003-11-12 14:32:26 +0000425static PyTypeObject teedataobject_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000426 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000427 "itertools.tee_dataobject", /* tp_name */
428 sizeof(teedataobject), /* tp_basicsize */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000429 0, /* tp_itemsize */
430 /* methods */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000431 (destructor)teedataobject_dealloc, /* tp_dealloc */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000432 0, /* tp_print */
433 0, /* tp_getattr */
434 0, /* tp_setattr */
435 0, /* tp_compare */
436 0, /* tp_repr */
437 0, /* tp_as_number */
438 0, /* tp_as_sequence */
439 0, /* tp_as_mapping */
440 0, /* tp_hash */
441 0, /* tp_call */
442 0, /* tp_str */
443 PyObject_GenericGetAttr, /* tp_getattro */
444 0, /* tp_setattro */
445 0, /* tp_as_buffer */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000446 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000447 teedataobject_doc, /* tp_doc */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000448 (traverseproc)teedataobject_traverse, /* tp_traverse */
449 (inquiry)teedataobject_clear, /* tp_clear */
450 0, /* tp_richcompare */
451 0, /* tp_weaklistoffset */
452 0, /* tp_iter */
453 0, /* tp_iternext */
454 0, /* tp_methods */
455 0, /* tp_members */
456 0, /* tp_getset */
457 0, /* tp_base */
458 0, /* tp_dict */
459 0, /* tp_descr_get */
460 0, /* tp_descr_set */
461 0, /* tp_dictoffset */
462 0, /* tp_init */
463 0, /* tp_alloc */
464 0, /* tp_new */
465 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000466};
467
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000468
469static PyTypeObject tee_type;
470
471static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000472tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000473{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000474 PyObject *value, *link;
475
476 if (to->index >= LINKCELLS) {
477 link = teedataobject_jumplink(to->dataobj);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000478 Py_DECREF(to->dataobj);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000479 to->dataobj = (teedataobject *)link;
480 to->index = 0;
481 }
482 value = teedataobject_getitem(to->dataobj, to->index);
483 if (value == NULL)
484 return NULL;
485 to->index++;
486 return value;
487}
488
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000489static int
490tee_traverse(teeobject *to, visitproc visit, void *arg)
491{
492 Py_VISIT((PyObject *)to->dataobj);
493 return 0;
494}
495
Raymond Hettingerad983e72003-11-12 14:32:26 +0000496static PyObject *
497tee_copy(teeobject *to)
498{
499 teeobject *newto;
500
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000501 newto = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000502 if (newto == NULL)
503 return NULL;
504 Py_INCREF(to->dataobj);
505 newto->dataobj = to->dataobj;
506 newto->index = to->index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000507 newto->weakreflist = NULL;
Thomas Woutersb3deb942006-04-15 22:33:13 +0000508 PyObject_GC_Track(newto);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000509 return (PyObject *)newto;
510}
511
512PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
513
514static PyObject *
515tee_fromiterable(PyObject *iterable)
516{
517 teeobject *to;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000518 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000519
520 it = PyObject_GetIter(iterable);
521 if (it == NULL)
522 return NULL;
523 if (PyObject_TypeCheck(it, &tee_type)) {
524 to = (teeobject *)tee_copy((teeobject *)it);
525 goto done;
526 }
527
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000528 to = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000529 if (to == NULL)
530 goto done;
531 to->dataobj = (teedataobject *)teedataobject_new(it);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000532 if (!to->dataobj) {
533 PyObject_GC_Del(to);
534 to = NULL;
535 goto done;
536 }
537
Raymond Hettingerad983e72003-11-12 14:32:26 +0000538 to->index = 0;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000539 to->weakreflist = NULL;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000540 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000541done:
542 Py_XDECREF(it);
543 return (PyObject *)to;
544}
545
546static PyObject *
547tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
548{
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000549 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000550
551 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
552 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000553 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000554}
555
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000556static int
557tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000558{
Raymond Hettingera9f60922004-10-17 16:40:14 +0000559 if (to->weakreflist != NULL)
560 PyObject_ClearWeakRefs((PyObject *) to);
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000561 Py_CLEAR(to->dataobj);
562 return 0;
563}
564
565static void
566tee_dealloc(teeobject *to)
567{
568 PyObject_GC_UnTrack(to);
569 tee_clear(to);
570 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000571}
572
Raymond Hettingerad983e72003-11-12 14:32:26 +0000573PyDoc_STRVAR(teeobject_doc,
574"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000575
Raymond Hettingerad983e72003-11-12 14:32:26 +0000576static PyMethodDef tee_methods[] = {
577 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
578 {NULL, NULL} /* sentinel */
579};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000580
581static PyTypeObject tee_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000582 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000583 "itertools.tee", /* tp_name */
584 sizeof(teeobject), /* tp_basicsize */
585 0, /* tp_itemsize */
586 /* methods */
587 (destructor)tee_dealloc, /* tp_dealloc */
588 0, /* tp_print */
589 0, /* tp_getattr */
590 0, /* tp_setattr */
591 0, /* tp_compare */
592 0, /* tp_repr */
593 0, /* tp_as_number */
594 0, /* tp_as_sequence */
595 0, /* tp_as_mapping */
596 0, /* tp_hash */
597 0, /* tp_call */
598 0, /* tp_str */
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000599 0, /* tp_getattro */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000600 0, /* tp_setattro */
601 0, /* tp_as_buffer */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000602 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000603 teeobject_doc, /* tp_doc */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000604 (traverseproc)tee_traverse, /* tp_traverse */
605 (inquiry)tee_clear, /* tp_clear */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000606 0, /* tp_richcompare */
Raymond Hettingera9f60922004-10-17 16:40:14 +0000607 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000608 PyObject_SelfIter, /* tp_iter */
609 (iternextfunc)tee_next, /* tp_iternext */
610 tee_methods, /* tp_methods */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000611 0, /* tp_members */
612 0, /* tp_getset */
613 0, /* tp_base */
614 0, /* tp_dict */
615 0, /* tp_descr_get */
616 0, /* tp_descr_set */
617 0, /* tp_dictoffset */
618 0, /* tp_init */
619 0, /* tp_alloc */
620 tee_new, /* tp_new */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000621 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000622};
623
Raymond Hettingerad983e72003-11-12 14:32:26 +0000624static PyObject *
625tee(PyObject *self, PyObject *args)
626{
Neal Norwitz69e88972006-09-02 02:58:13 +0000627 Py_ssize_t i, n=2;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000628 PyObject *it, *iterable, *copyable, *result;
629
Neal Norwitz69e88972006-09-02 02:58:13 +0000630 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000631 return NULL;
Neal Norwitz69e88972006-09-02 02:58:13 +0000632 if (n < 0) {
633 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
634 return NULL;
635 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000636 result = PyTuple_New(n);
637 if (result == NULL)
638 return NULL;
639 if (n == 0)
640 return result;
641 it = PyObject_GetIter(iterable);
642 if (it == NULL) {
643 Py_DECREF(result);
644 return NULL;
645 }
646 if (!PyObject_HasAttrString(it, "__copy__")) {
647 copyable = tee_fromiterable(it);
648 Py_DECREF(it);
649 if (copyable == NULL) {
650 Py_DECREF(result);
651 return NULL;
652 }
653 } else
654 copyable = it;
655 PyTuple_SET_ITEM(result, 0, copyable);
656 for (i=1 ; i<n ; i++) {
657 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
658 if (copyable == NULL) {
659 Py_DECREF(result);
660 return NULL;
661 }
662 PyTuple_SET_ITEM(result, i, copyable);
663 }
664 return result;
665}
666
667PyDoc_STRVAR(tee_doc,
668"tee(iterable, n=2) --> tuple of n independent iterators.");
669
670
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000671/* cycle object **********************************************************/
672
673typedef struct {
674 PyObject_HEAD
675 PyObject *it;
676 PyObject *saved;
677 int firstpass;
678} cycleobject;
679
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000680static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000681
682static PyObject *
683cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
684{
685 PyObject *it;
686 PyObject *iterable;
687 PyObject *saved;
688 cycleobject *lz;
689
Georg Brandlb84c1372007-01-21 10:28:43 +0000690 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000691 return NULL;
692
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000693 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
694 return NULL;
695
696 /* Get iterator. */
697 it = PyObject_GetIter(iterable);
698 if (it == NULL)
699 return NULL;
700
701 saved = PyList_New(0);
702 if (saved == NULL) {
703 Py_DECREF(it);
704 return NULL;
705 }
706
707 /* create cycleobject structure */
708 lz = (cycleobject *)type->tp_alloc(type, 0);
709 if (lz == NULL) {
710 Py_DECREF(it);
711 Py_DECREF(saved);
712 return NULL;
713 }
714 lz->it = it;
715 lz->saved = saved;
716 lz->firstpass = 0;
717
718 return (PyObject *)lz;
719}
720
721static void
722cycle_dealloc(cycleobject *lz)
723{
724 PyObject_GC_UnTrack(lz);
725 Py_XDECREF(lz->saved);
726 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +0000727 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000728}
729
730static int
731cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
732{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000733 Py_VISIT(lz->it);
734 Py_VISIT(lz->saved);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000735 return 0;
736}
737
738static PyObject *
739cycle_next(cycleobject *lz)
740{
741 PyObject *item;
742 PyObject *it;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000743 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000744
745 while (1) {
746 item = PyIter_Next(lz->it);
747 if (item != NULL) {
748 if (!lz->firstpass)
749 PyList_Append(lz->saved, item);
750 return item;
751 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000752 if (PyErr_Occurred()) {
753 if (PyErr_ExceptionMatches(PyExc_StopIteration))
754 PyErr_Clear();
755 else
756 return NULL;
757 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000758 if (PyList_Size(lz->saved) == 0)
759 return NULL;
760 it = PyObject_GetIter(lz->saved);
761 if (it == NULL)
762 return NULL;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000763 tmp = lz->it;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000764 lz->it = it;
765 lz->firstpass = 1;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000766 Py_DECREF(tmp);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000767 }
768}
769
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000770PyDoc_STRVAR(cycle_doc,
771"cycle(iterable) --> cycle object\n\
772\n\
773Return elements from the iterable until it is exhausted.\n\
774Then repeat the sequence indefinitely.");
775
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000776static PyTypeObject cycle_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000777 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000778 "itertools.cycle", /* tp_name */
779 sizeof(cycleobject), /* tp_basicsize */
780 0, /* tp_itemsize */
781 /* methods */
782 (destructor)cycle_dealloc, /* tp_dealloc */
783 0, /* tp_print */
784 0, /* tp_getattr */
785 0, /* tp_setattr */
786 0, /* tp_compare */
787 0, /* tp_repr */
788 0, /* tp_as_number */
789 0, /* tp_as_sequence */
790 0, /* tp_as_mapping */
791 0, /* tp_hash */
792 0, /* tp_call */
793 0, /* tp_str */
794 PyObject_GenericGetAttr, /* tp_getattro */
795 0, /* tp_setattro */
796 0, /* tp_as_buffer */
797 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
798 Py_TPFLAGS_BASETYPE, /* tp_flags */
799 cycle_doc, /* tp_doc */
800 (traverseproc)cycle_traverse, /* tp_traverse */
801 0, /* tp_clear */
802 0, /* tp_richcompare */
803 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000804 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000805 (iternextfunc)cycle_next, /* tp_iternext */
806 0, /* tp_methods */
807 0, /* tp_members */
808 0, /* tp_getset */
809 0, /* tp_base */
810 0, /* tp_dict */
811 0, /* tp_descr_get */
812 0, /* tp_descr_set */
813 0, /* tp_dictoffset */
814 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000815 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000816 cycle_new, /* tp_new */
817 PyObject_GC_Del, /* tp_free */
818};
819
820
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000821/* dropwhile object **********************************************************/
822
823typedef struct {
824 PyObject_HEAD
825 PyObject *func;
826 PyObject *it;
827 long start;
828} dropwhileobject;
829
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000830static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000831
832static PyObject *
833dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
834{
835 PyObject *func, *seq;
836 PyObject *it;
837 dropwhileobject *lz;
838
Georg Brandlb84c1372007-01-21 10:28:43 +0000839 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000840 return NULL;
841
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000842 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
843 return NULL;
844
845 /* Get iterator. */
846 it = PyObject_GetIter(seq);
847 if (it == NULL)
848 return NULL;
849
850 /* create dropwhileobject structure */
851 lz = (dropwhileobject *)type->tp_alloc(type, 0);
852 if (lz == NULL) {
853 Py_DECREF(it);
854 return NULL;
855 }
856 Py_INCREF(func);
857 lz->func = func;
858 lz->it = it;
859 lz->start = 0;
860
861 return (PyObject *)lz;
862}
863
864static void
865dropwhile_dealloc(dropwhileobject *lz)
866{
867 PyObject_GC_UnTrack(lz);
868 Py_XDECREF(lz->func);
869 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +0000870 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000871}
872
873static int
874dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
875{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000876 Py_VISIT(lz->it);
877 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000878 return 0;
879}
880
881static PyObject *
882dropwhile_next(dropwhileobject *lz)
883{
884 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000885 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000886 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000887 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000888
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000889 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +0000890 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000891 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000892 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000893 if (item == NULL)
894 return NULL;
895 if (lz->start == 1)
896 return item;
897
898 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
899 if (good == NULL) {
900 Py_DECREF(item);
901 return NULL;
902 }
903 ok = PyObject_IsTrue(good);
904 Py_DECREF(good);
905 if (!ok) {
906 lz->start = 1;
907 return item;
908 }
909 Py_DECREF(item);
910 }
911}
912
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000913PyDoc_STRVAR(dropwhile_doc,
914"dropwhile(predicate, iterable) --> dropwhile object\n\
915\n\
916Drop items from the iterable while predicate(item) is true.\n\
917Afterwards, return every element until the iterable is exhausted.");
918
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000919static PyTypeObject dropwhile_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000920 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000921 "itertools.dropwhile", /* tp_name */
922 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000923 0, /* tp_itemsize */
924 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000925 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000926 0, /* tp_print */
927 0, /* tp_getattr */
928 0, /* tp_setattr */
929 0, /* tp_compare */
930 0, /* tp_repr */
931 0, /* tp_as_number */
932 0, /* tp_as_sequence */
933 0, /* tp_as_mapping */
934 0, /* tp_hash */
935 0, /* tp_call */
936 0, /* tp_str */
937 PyObject_GenericGetAttr, /* tp_getattro */
938 0, /* tp_setattro */
939 0, /* tp_as_buffer */
940 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
941 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000942 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000943 (traverseproc)dropwhile_traverse, /* tp_traverse */
944 0, /* tp_clear */
945 0, /* tp_richcompare */
946 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000947 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000948 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000949 0, /* tp_methods */
950 0, /* tp_members */
951 0, /* tp_getset */
952 0, /* tp_base */
953 0, /* tp_dict */
954 0, /* tp_descr_get */
955 0, /* tp_descr_set */
956 0, /* tp_dictoffset */
957 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000958 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000959 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000960 PyObject_GC_Del, /* tp_free */
961};
962
963
964/* takewhile object **********************************************************/
965
966typedef struct {
967 PyObject_HEAD
968 PyObject *func;
969 PyObject *it;
970 long stop;
971} takewhileobject;
972
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000973static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000974
975static PyObject *
976takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
977{
978 PyObject *func, *seq;
979 PyObject *it;
980 takewhileobject *lz;
981
Georg Brandlb84c1372007-01-21 10:28:43 +0000982 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000983 return NULL;
984
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000985 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
986 return NULL;
987
988 /* Get iterator. */
989 it = PyObject_GetIter(seq);
990 if (it == NULL)
991 return NULL;
992
993 /* create takewhileobject structure */
994 lz = (takewhileobject *)type->tp_alloc(type, 0);
995 if (lz == NULL) {
996 Py_DECREF(it);
997 return NULL;
998 }
999 Py_INCREF(func);
1000 lz->func = func;
1001 lz->it = it;
1002 lz->stop = 0;
1003
1004 return (PyObject *)lz;
1005}
1006
1007static void
1008takewhile_dealloc(takewhileobject *lz)
1009{
1010 PyObject_GC_UnTrack(lz);
1011 Py_XDECREF(lz->func);
1012 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00001013 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001014}
1015
1016static int
1017takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1018{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001019 Py_VISIT(lz->it);
1020 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001021 return 0;
1022}
1023
1024static PyObject *
1025takewhile_next(takewhileobject *lz)
1026{
1027 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001028 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001029 long ok;
1030
1031 if (lz->stop == 1)
1032 return NULL;
1033
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001034 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00001035 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001036 if (item == NULL)
1037 return NULL;
1038
1039 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1040 if (good == NULL) {
1041 Py_DECREF(item);
1042 return NULL;
1043 }
1044 ok = PyObject_IsTrue(good);
1045 Py_DECREF(good);
1046 if (ok)
1047 return item;
1048 Py_DECREF(item);
1049 lz->stop = 1;
1050 return NULL;
1051}
1052
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053PyDoc_STRVAR(takewhile_doc,
1054"takewhile(predicate, iterable) --> takewhile object\n\
1055\n\
1056Return successive entries from an iterable as long as the \n\
1057predicate evaluates to true for each entry.");
1058
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001059static PyTypeObject takewhile_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001060 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001061 "itertools.takewhile", /* tp_name */
1062 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001063 0, /* tp_itemsize */
1064 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001065 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001066 0, /* tp_print */
1067 0, /* tp_getattr */
1068 0, /* tp_setattr */
1069 0, /* tp_compare */
1070 0, /* tp_repr */
1071 0, /* tp_as_number */
1072 0, /* tp_as_sequence */
1073 0, /* tp_as_mapping */
1074 0, /* tp_hash */
1075 0, /* tp_call */
1076 0, /* tp_str */
1077 PyObject_GenericGetAttr, /* tp_getattro */
1078 0, /* tp_setattro */
1079 0, /* tp_as_buffer */
1080 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1081 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001082 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001083 (traverseproc)takewhile_traverse, /* tp_traverse */
1084 0, /* tp_clear */
1085 0, /* tp_richcompare */
1086 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001087 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001088 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001089 0, /* tp_methods */
1090 0, /* tp_members */
1091 0, /* tp_getset */
1092 0, /* tp_base */
1093 0, /* tp_dict */
1094 0, /* tp_descr_get */
1095 0, /* tp_descr_set */
1096 0, /* tp_dictoffset */
1097 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001098 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001099 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001100 PyObject_GC_Del, /* tp_free */
1101};
1102
1103
1104/* islice object ************************************************************/
1105
1106typedef struct {
1107 PyObject_HEAD
1108 PyObject *it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001109 Py_ssize_t next;
1110 Py_ssize_t stop;
1111 Py_ssize_t step;
1112 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001113} isliceobject;
1114
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001115static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116
1117static PyObject *
1118islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1119{
1120 PyObject *seq;
Jack Diederich6c433a92006-05-26 11:15:17 +00001121 Py_ssize_t start=0, stop=-1, step=1;
Raymond Hettingerb2594052004-12-05 09:25:51 +00001122 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001123 Py_ssize_t numargs;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001124 isliceobject *lz;
1125
Georg Brandlb84c1372007-01-21 10:28:43 +00001126 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001127 return NULL;
1128
Raymond Hettingerb2594052004-12-05 09:25:51 +00001129 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001130 return NULL;
1131
Raymond Hettingerb2594052004-12-05 09:25:51 +00001132 numargs = PyTuple_Size(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001133 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001134 if (a1 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001135 stop = PyInt_AsSsize_t(a1);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001136 if (stop == -1) {
1137 if (PyErr_Occurred())
1138 PyErr_Clear();
1139 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001140 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001141 return NULL;
1142 }
1143 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001144 } else {
Raymond Hettingerb2594052004-12-05 09:25:51 +00001145 if (a1 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001146 start = PyInt_AsSsize_t(a1);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001147 if (start == -1 && PyErr_Occurred())
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001148 PyErr_Clear();
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001149 if (a2 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001150 stop = PyInt_AsSsize_t(a2);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001151 if (stop == -1) {
1152 if (PyErr_Occurred())
1153 PyErr_Clear();
1154 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001155 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001156 return NULL;
1157 }
1158 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001159 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001160 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001161 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001162 "Indices for islice() must be non-negative integers or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001163 return NULL;
1164 }
1165
Raymond Hettingerb2594052004-12-05 09:25:51 +00001166 if (a3 != NULL) {
1167 if (a3 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001168 step = PyInt_AsSsize_t(a3);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001169 if (step == -1 && PyErr_Occurred())
1170 PyErr_Clear();
1171 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001172 if (step<1) {
1173 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001174 "Step for islice() must be a positive integer or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001175 return NULL;
1176 }
1177
1178 /* Get iterator. */
1179 it = PyObject_GetIter(seq);
1180 if (it == NULL)
1181 return NULL;
1182
1183 /* create isliceobject structure */
1184 lz = (isliceobject *)type->tp_alloc(type, 0);
1185 if (lz == NULL) {
1186 Py_DECREF(it);
1187 return NULL;
1188 }
1189 lz->it = it;
1190 lz->next = start;
1191 lz->stop = stop;
1192 lz->step = step;
1193 lz->cnt = 0L;
1194
1195 return (PyObject *)lz;
1196}
1197
1198static void
1199islice_dealloc(isliceobject *lz)
1200{
1201 PyObject_GC_UnTrack(lz);
1202 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00001203 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204}
1205
1206static int
1207islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1208{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001209 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001210 return 0;
1211}
1212
1213static PyObject *
1214islice_next(isliceobject *lz)
1215{
1216 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001217 PyObject *it = lz->it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001218 Py_ssize_t oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001219 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001220
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001221 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00001222 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001223 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001224 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001225 if (item == NULL)
1226 return NULL;
1227 Py_DECREF(item);
1228 lz->cnt++;
1229 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001230 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001231 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001232 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001233 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001234 if (item == NULL)
1235 return NULL;
1236 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001237 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001238 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001239 if (lz->next < oldnext) /* Check for overflow */
1240 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001241 return item;
1242}
1243
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001244PyDoc_STRVAR(islice_doc,
1245"islice(iterable, [start,] stop [, step]) --> islice object\n\
1246\n\
1247Return an iterator whose next() method returns selected values from an\n\
1248iterable. If start is specified, will skip all preceding elements;\n\
1249otherwise, start defaults to zero. Step defaults to one. If\n\
1250specified as another value, step determines how many values are \n\
1251skipped between successive calls. Works like a slice() on a list\n\
1252but returns an iterator.");
1253
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001254static PyTypeObject islice_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001255 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001256 "itertools.islice", /* tp_name */
1257 sizeof(isliceobject), /* tp_basicsize */
1258 0, /* tp_itemsize */
1259 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001260 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001261 0, /* tp_print */
1262 0, /* tp_getattr */
1263 0, /* tp_setattr */
1264 0, /* tp_compare */
1265 0, /* tp_repr */
1266 0, /* tp_as_number */
1267 0, /* tp_as_sequence */
1268 0, /* tp_as_mapping */
1269 0, /* tp_hash */
1270 0, /* tp_call */
1271 0, /* tp_str */
1272 PyObject_GenericGetAttr, /* tp_getattro */
1273 0, /* tp_setattro */
1274 0, /* tp_as_buffer */
1275 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1276 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001277 islice_doc, /* tp_doc */
1278 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001279 0, /* tp_clear */
1280 0, /* tp_richcompare */
1281 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001282 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001283 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001284 0, /* tp_methods */
1285 0, /* tp_members */
1286 0, /* tp_getset */
1287 0, /* tp_base */
1288 0, /* tp_dict */
1289 0, /* tp_descr_get */
1290 0, /* tp_descr_set */
1291 0, /* tp_dictoffset */
1292 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001293 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001294 islice_new, /* tp_new */
1295 PyObject_GC_Del, /* tp_free */
1296};
1297
1298
1299/* starmap object ************************************************************/
1300
1301typedef struct {
1302 PyObject_HEAD
1303 PyObject *func;
1304 PyObject *it;
1305} starmapobject;
1306
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001307static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001308
1309static PyObject *
1310starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1311{
1312 PyObject *func, *seq;
1313 PyObject *it;
1314 starmapobject *lz;
1315
Georg Brandlb84c1372007-01-21 10:28:43 +00001316 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001317 return NULL;
1318
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001319 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1320 return NULL;
1321
1322 /* Get iterator. */
1323 it = PyObject_GetIter(seq);
1324 if (it == NULL)
1325 return NULL;
1326
1327 /* create starmapobject structure */
1328 lz = (starmapobject *)type->tp_alloc(type, 0);
1329 if (lz == NULL) {
1330 Py_DECREF(it);
1331 return NULL;
1332 }
1333 Py_INCREF(func);
1334 lz->func = func;
1335 lz->it = it;
1336
1337 return (PyObject *)lz;
1338}
1339
1340static void
1341starmap_dealloc(starmapobject *lz)
1342{
1343 PyObject_GC_UnTrack(lz);
1344 Py_XDECREF(lz->func);
1345 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00001346 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001347}
1348
1349static int
1350starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1351{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001352 Py_VISIT(lz->it);
1353 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001354 return 0;
1355}
1356
1357static PyObject *
1358starmap_next(starmapobject *lz)
1359{
1360 PyObject *args;
1361 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001362 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001363
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001364 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00001365 args = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001366 if (args == NULL)
1367 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001368 if (!PyTuple_CheckExact(args)) {
Raymond Hettinger47317092008-01-17 03:02:14 +00001369 PyObject *newargs = PySequence_Tuple(args);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001370 Py_DECREF(args);
Raymond Hettinger47317092008-01-17 03:02:14 +00001371 if (newargs == NULL)
1372 return NULL;
1373 args = newargs;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001374 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001375 result = PyObject_Call(lz->func, args, NULL);
1376 Py_DECREF(args);
1377 return result;
1378}
1379
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001380PyDoc_STRVAR(starmap_doc,
1381"starmap(function, sequence) --> starmap object\n\
1382\n\
1383Return an iterator whose values are returned from the function evaluated\n\
1384with a argument tuple taken from the given sequence.");
1385
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001386static PyTypeObject starmap_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001387 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001388 "itertools.starmap", /* tp_name */
1389 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001390 0, /* tp_itemsize */
1391 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001392 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001393 0, /* tp_print */
1394 0, /* tp_getattr */
1395 0, /* tp_setattr */
1396 0, /* tp_compare */
1397 0, /* tp_repr */
1398 0, /* tp_as_number */
1399 0, /* tp_as_sequence */
1400 0, /* tp_as_mapping */
1401 0, /* tp_hash */
1402 0, /* tp_call */
1403 0, /* tp_str */
1404 PyObject_GenericGetAttr, /* tp_getattro */
1405 0, /* tp_setattro */
1406 0, /* tp_as_buffer */
1407 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1408 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001409 starmap_doc, /* tp_doc */
1410 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001411 0, /* tp_clear */
1412 0, /* tp_richcompare */
1413 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001414 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001415 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001416 0, /* tp_methods */
1417 0, /* tp_members */
1418 0, /* tp_getset */
1419 0, /* tp_base */
1420 0, /* tp_dict */
1421 0, /* tp_descr_get */
1422 0, /* tp_descr_set */
1423 0, /* tp_dictoffset */
1424 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001425 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001426 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001427 PyObject_GC_Del, /* tp_free */
1428};
1429
1430
1431/* imap object ************************************************************/
1432
1433typedef struct {
1434 PyObject_HEAD
1435 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001436 PyObject *func;
1437} imapobject;
1438
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001439static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001440
1441static PyObject *
1442imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1443{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001444 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001445 imapobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001446 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001447
Raymond Hettinger10f40a62008-03-13 16:43:17 +00001448 if (Py_Py3kWarningFlag &&
1449 PyErr_Warn(PyExc_DeprecationWarning,
1450 "In 3.x, itertools.imap() was moved to builtin map()") < 0)
1451 return NULL;
1452
Georg Brandlb84c1372007-01-21 10:28:43 +00001453 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001454 return NULL;
1455
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001456 numargs = PyTuple_Size(args);
1457 if (numargs < 2) {
1458 PyErr_SetString(PyExc_TypeError,
1459 "imap() must have at least two arguments.");
1460 return NULL;
1461 }
1462
1463 iters = PyTuple_New(numargs-1);
1464 if (iters == NULL)
1465 return NULL;
1466
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001467 for (i=1 ; i<numargs ; i++) {
1468 /* Get iterator. */
1469 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1470 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001471 Py_DECREF(iters);
1472 return NULL;
1473 }
1474 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001475 }
1476
1477 /* create imapobject structure */
1478 lz = (imapobject *)type->tp_alloc(type, 0);
1479 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001480 Py_DECREF(iters);
1481 return NULL;
1482 }
1483 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001484 func = PyTuple_GET_ITEM(args, 0);
1485 Py_INCREF(func);
1486 lz->func = func;
1487
1488 return (PyObject *)lz;
1489}
1490
1491static void
1492imap_dealloc(imapobject *lz)
1493{
1494 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001495 Py_XDECREF(lz->iters);
1496 Py_XDECREF(lz->func);
Christian Heimese93237d2007-12-19 02:37:44 +00001497 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001498}
1499
1500static int
1501imap_traverse(imapobject *lz, visitproc visit, void *arg)
1502{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001503 Py_VISIT(lz->iters);
1504 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001505 return 0;
1506}
1507
Raymond Hettinger2012f172003-02-07 05:32:58 +00001508/*
1509imap() is an iterator version of __builtins__.map() except that it does
1510not have the None fill-in feature. That was intentionally left out for
1511the following reasons:
1512
1513 1) Itertools are designed to be easily combined and chained together.
1514 Having all tools stop with the shortest input is a unifying principle
1515 that makes it easier to combine finite iterators (supplying data) with
1516 infinite iterators like count() and repeat() (for supplying sequential
1517 or constant arguments to a function).
1518
1519 2) In typical use cases for combining itertools, having one finite data
1520 supplier run out before another is likely to be an error condition which
1521 should not pass silently by automatically supplying None.
1522
1523 3) The use cases for automatic None fill-in are rare -- not many functions
1524 do something useful when a parameter suddenly switches type and becomes
1525 None.
1526
1527 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001528 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001529
1530 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1531*/
1532
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001533static PyObject *
1534imap_next(imapobject *lz)
1535{
1536 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001537 PyObject *argtuple;
1538 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001539 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001540
1541 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001542 argtuple = PyTuple_New(numargs);
1543 if (argtuple == NULL)
1544 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001545
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001546 for (i=0 ; i<numargs ; i++) {
1547 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1548 if (val == NULL) {
1549 Py_DECREF(argtuple);
1550 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001551 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001552 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001553 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001554 if (lz->func == Py_None)
1555 return argtuple;
1556 result = PyObject_Call(lz->func, argtuple, NULL);
1557 Py_DECREF(argtuple);
1558 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001559}
1560
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001561PyDoc_STRVAR(imap_doc,
1562"imap(func, *iterables) --> imap object\n\
1563\n\
1564Make an iterator that computes the function using arguments from\n\
1565each of the iterables. Like map() except that it returns\n\
1566an iterator instead of a list and that it stops when the shortest\n\
1567iterable is exhausted instead of filling in None for shorter\n\
1568iterables.");
1569
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001570static PyTypeObject imap_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001571 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001572 "itertools.imap", /* tp_name */
1573 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001574 0, /* tp_itemsize */
1575 /* methods */
1576 (destructor)imap_dealloc, /* tp_dealloc */
1577 0, /* tp_print */
1578 0, /* tp_getattr */
1579 0, /* tp_setattr */
1580 0, /* tp_compare */
1581 0, /* tp_repr */
1582 0, /* tp_as_number */
1583 0, /* tp_as_sequence */
1584 0, /* tp_as_mapping */
1585 0, /* tp_hash */
1586 0, /* tp_call */
1587 0, /* tp_str */
1588 PyObject_GenericGetAttr, /* tp_getattro */
1589 0, /* tp_setattro */
1590 0, /* tp_as_buffer */
1591 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1592 Py_TPFLAGS_BASETYPE, /* tp_flags */
1593 imap_doc, /* tp_doc */
1594 (traverseproc)imap_traverse, /* tp_traverse */
1595 0, /* tp_clear */
1596 0, /* tp_richcompare */
1597 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001598 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001599 (iternextfunc)imap_next, /* tp_iternext */
1600 0, /* tp_methods */
1601 0, /* tp_members */
1602 0, /* tp_getset */
1603 0, /* tp_base */
1604 0, /* tp_dict */
1605 0, /* tp_descr_get */
1606 0, /* tp_descr_set */
1607 0, /* tp_dictoffset */
1608 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001609 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001610 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001611 PyObject_GC_Del, /* tp_free */
1612};
1613
1614
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001615/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001616
1617typedef struct {
1618 PyObject_HEAD
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001619 PyObject *source; /* Iterator over input iterables */
1620 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001621} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001622
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001623static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001624
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001625static PyObject *
1626chain_new_internal(PyTypeObject *type, PyObject *source)
1627{
1628 chainobject *lz;
1629
1630 lz = (chainobject *)type->tp_alloc(type, 0);
1631 if (lz == NULL) {
1632 Py_DECREF(source);
1633 return NULL;
1634 }
1635
1636 lz->source = source;
1637 lz->active = NULL;
1638 return (PyObject *)lz;
1639}
1640
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001641static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001642chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001643{
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001644 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001645
Georg Brandlb84c1372007-01-21 10:28:43 +00001646 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001647 return NULL;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001648
1649 source = PyObject_GetIter(args);
1650 if (source == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001651 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001652
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001653 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001654}
1655
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001656static PyObject *
1657chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1658{
1659 PyObject *source;
1660
1661 source = PyObject_GetIter(arg);
1662 if (source == NULL)
1663 return NULL;
1664
1665 return chain_new_internal(type, source);
1666}
1667
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001668static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001669chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001670{
1671 PyObject_GC_UnTrack(lz);
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001672 Py_XDECREF(lz->active);
1673 Py_XDECREF(lz->source);
Christian Heimese93237d2007-12-19 02:37:44 +00001674 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001675}
1676
Raymond Hettinger2012f172003-02-07 05:32:58 +00001677static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001678chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001679{
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001680 Py_VISIT(lz->source);
1681 Py_VISIT(lz->active);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001682 return 0;
1683}
1684
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001685static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001686chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001687{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001688 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001689
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001690 if (lz->source == NULL)
1691 return NULL; /* already stopped */
1692
1693 if (lz->active == NULL) {
1694 PyObject *iterable = PyIter_Next(lz->source);
1695 if (iterable == NULL) {
1696 Py_CLEAR(lz->source);
1697 return NULL; /* no more input sources */
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001698 }
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001699 lz->active = PyObject_GetIter(iterable);
Raymond Hettingerf1cca2b2008-03-04 22:29:44 +00001700 Py_DECREF(iterable);
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001701 if (lz->active == NULL) {
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001702 Py_CLEAR(lz->source);
1703 return NULL; /* input not iterable */
1704 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001705 }
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001706 item = PyIter_Next(lz->active);
1707 if (item != NULL)
1708 return item;
1709 if (PyErr_Occurred()) {
1710 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1711 PyErr_Clear();
1712 else
1713 return NULL; /* input raised an exception */
1714 }
1715 Py_CLEAR(lz->active);
1716 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001717}
1718
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001719PyDoc_STRVAR(chain_doc,
1720"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001721\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001722Return a chain object whose .next() method returns elements from the\n\
1723first iterable until it is exhausted, then elements from the next\n\
1724iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001725
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001726PyDoc_STRVAR(chain_from_iterable_doc,
1727"chain.from_iterable(iterable) --> chain object\n\
1728\n\
1729Alternate chain() contructor taking a single iterable argument\n\
1730that evaluates lazily.");
1731
1732static PyMethodDef chain_methods[] = {
1733 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1734 chain_from_iterable_doc},
1735 {NULL, NULL} /* sentinel */
1736};
1737
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001738static PyTypeObject chain_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001739 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001740 "itertools.chain", /* tp_name */
1741 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001742 0, /* tp_itemsize */
1743 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001744 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001745 0, /* tp_print */
1746 0, /* tp_getattr */
1747 0, /* tp_setattr */
1748 0, /* tp_compare */
1749 0, /* tp_repr */
1750 0, /* tp_as_number */
1751 0, /* tp_as_sequence */
1752 0, /* tp_as_mapping */
1753 0, /* tp_hash */
1754 0, /* tp_call */
1755 0, /* tp_str */
1756 PyObject_GenericGetAttr, /* tp_getattro */
1757 0, /* tp_setattro */
1758 0, /* tp_as_buffer */
1759 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1760 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001761 chain_doc, /* tp_doc */
1762 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001763 0, /* tp_clear */
1764 0, /* tp_richcompare */
1765 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001766 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001767 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001768 chain_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001769 0, /* tp_members */
1770 0, /* tp_getset */
1771 0, /* tp_base */
1772 0, /* tp_dict */
1773 0, /* tp_descr_get */
1774 0, /* tp_descr_set */
1775 0, /* tp_dictoffset */
1776 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001777 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001778 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001779 PyObject_GC_Del, /* tp_free */
1780};
1781
1782
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001783/* product object ************************************************************/
1784
1785typedef struct {
1786 PyObject_HEAD
1787 PyObject *pools; /* tuple of pool tuples */
Raymond Hettinger532316d2008-02-23 04:03:50 +00001788 Py_ssize_t *indices; /* one index per pool */
1789 PyObject *result; /* most recently returned result tuple */
1790 int stopped; /* set to 1 when the product iterator is exhausted */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001791} productobject;
1792
1793static PyTypeObject product_type;
1794
1795static PyObject *
1796product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1797{
1798 productobject *lz;
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001799 Py_ssize_t nargs, npools, repeat=1;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001800 PyObject *pools = NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001801 Py_ssize_t *indices = NULL;
1802 Py_ssize_t i;
1803
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001804 if (kwds != NULL) {
1805 char *kwlist[] = {"repeat", 0};
1806 PyObject *tmpargs = PyTuple_New(0);
1807 if (tmpargs == NULL)
1808 return NULL;
1809 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1810 Py_DECREF(tmpargs);
1811 return NULL;
1812 }
1813 Py_DECREF(tmpargs);
1814 if (repeat < 0) {
1815 PyErr_SetString(PyExc_ValueError,
1816 "repeat argument cannot be negative");
1817 return NULL;
1818 }
1819 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001820
1821 assert(PyTuple_Check(args));
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001822 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1823 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001824
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001825 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
Raymond Hettinger61024b92008-03-02 11:57:16 +00001826 if (indices == NULL) {
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001827 PyErr_NoMemory();
1828 goto error;
1829 }
1830
1831 pools = PyTuple_New(npools);
1832 if (pools == NULL)
1833 goto error;
1834
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001835 for (i=0; i < nargs ; ++i) {
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001836 PyObject *item = PyTuple_GET_ITEM(args, i);
1837 PyObject *pool = PySequence_Tuple(item);
1838 if (pool == NULL)
1839 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001840 PyTuple_SET_ITEM(pools, i, pool);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001841 indices[i] = 0;
1842 }
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001843 for ( ; i < npools; ++i) {
1844 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1845 Py_INCREF(pool);
1846 PyTuple_SET_ITEM(pools, i, pool);
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001847 indices[i] = 0;
1848 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001849
1850 /* create productobject structure */
1851 lz = (productobject *)type->tp_alloc(type, 0);
Raymond Hettinger3bd77122008-02-27 01:08:04 +00001852 if (lz == NULL)
Raymond Hettinger73d79632008-02-23 02:20:41 +00001853 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001854
1855 lz->pools = pools;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001856 lz->indices = indices;
1857 lz->result = NULL;
1858 lz->stopped = 0;
1859
1860 return (PyObject *)lz;
1861
1862error:
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001863 if (indices != NULL)
1864 PyMem_Free(indices);
1865 Py_XDECREF(pools);
1866 return NULL;
1867}
1868
1869static void
1870product_dealloc(productobject *lz)
1871{
1872 PyObject_GC_UnTrack(lz);
1873 Py_XDECREF(lz->pools);
1874 Py_XDECREF(lz->result);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001875 PyMem_Free(lz->indices);
1876 Py_TYPE(lz)->tp_free(lz);
1877}
1878
1879static int
1880product_traverse(productobject *lz, visitproc visit, void *arg)
1881{
1882 Py_VISIT(lz->pools);
1883 Py_VISIT(lz->result);
1884 return 0;
1885}
1886
1887static PyObject *
1888product_next(productobject *lz)
1889{
1890 PyObject *pool;
1891 PyObject *elem;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001892 PyObject *oldelem;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001893 PyObject *pools = lz->pools;
1894 PyObject *result = lz->result;
1895 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1896 Py_ssize_t i;
1897
1898 if (lz->stopped)
1899 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001900
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001901 if (result == NULL) {
Raymond Hettinger73d79632008-02-23 02:20:41 +00001902 /* On the first pass, return an initial tuple filled with the
Raymond Hettingerd553d852008-03-04 04:17:08 +00001903 first element from each pool. */
Raymond Hettinger73d79632008-02-23 02:20:41 +00001904 result = PyTuple_New(npools);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001905 if (result == NULL)
1906 goto empty;
1907 lz->result = result;
1908 for (i=0; i < npools; i++) {
1909 pool = PyTuple_GET_ITEM(pools, i);
1910 if (PyTuple_GET_SIZE(pool) == 0)
1911 goto empty;
1912 elem = PyTuple_GET_ITEM(pool, 0);
1913 Py_INCREF(elem);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001914 PyTuple_SET_ITEM(result, i, elem);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001915 }
1916 } else {
1917 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001918
1919 /* Copy the previous result tuple or re-use it if available */
1920 if (Py_REFCNT(result) > 1) {
1921 PyObject *old_result = result;
1922 result = PyTuple_New(npools);
1923 if (result == NULL)
1924 goto empty;
1925 lz->result = result;
1926 for (i=0; i < npools; i++) {
1927 elem = PyTuple_GET_ITEM(old_result, i);
1928 Py_INCREF(elem);
1929 PyTuple_SET_ITEM(result, i, elem);
1930 }
1931 Py_DECREF(old_result);
1932 }
1933 /* Now, we've got the only copy so we can update it in-place */
Raymond Hettingere3fabd12008-03-02 12:02:19 +00001934 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001935
1936 /* Update the pool indices right-to-left. Only advance to the
1937 next pool when the previous one rolls-over */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001938 for (i=npools-1 ; i >= 0 ; i--) {
1939 pool = PyTuple_GET_ITEM(pools, i);
1940 indices[i]++;
Raymond Hettinger61024b92008-03-02 11:57:16 +00001941 if (indices[i] == PyTuple_GET_SIZE(pool)) {
Raymond Hettinger73d79632008-02-23 02:20:41 +00001942 /* Roll-over and advance to next pool */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001943 indices[i] = 0;
1944 elem = PyTuple_GET_ITEM(pool, 0);
1945 Py_INCREF(elem);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001946 oldelem = PyTuple_GET_ITEM(result, i);
1947 PyTuple_SET_ITEM(result, i, elem);
1948 Py_DECREF(oldelem);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001949 } else {
Raymond Hettinger73d79632008-02-23 02:20:41 +00001950 /* No rollover. Just increment and stop here. */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001951 elem = PyTuple_GET_ITEM(pool, indices[i]);
1952 Py_INCREF(elem);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001953 oldelem = PyTuple_GET_ITEM(result, i);
1954 PyTuple_SET_ITEM(result, i, elem);
1955 Py_DECREF(oldelem);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001956 break;
1957 }
1958 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00001959
1960 /* If i is negative, then the indices have all rolled-over
1961 and we're done. */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001962 if (i < 0)
Raymond Hettinger73d79632008-02-23 02:20:41 +00001963 goto empty;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001964 }
1965
Raymond Hettinger73d79632008-02-23 02:20:41 +00001966 Py_INCREF(result);
1967 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001968
1969empty:
1970 lz->stopped = 1;
1971 return NULL;
1972}
1973
1974PyDoc_STRVAR(product_doc,
1975"product(*iterables) --> product object\n\
1976\n\
Raymond Hettinger73d79632008-02-23 02:20:41 +00001977Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001978For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1979The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1980cycle in a manner similar to an odometer (with the rightmost element changing\n\
1981on every iteration).\n\n\
1982product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1983product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1984
1985static PyTypeObject product_type = {
1986 PyVarObject_HEAD_INIT(NULL, 0)
1987 "itertools.product", /* tp_name */
1988 sizeof(productobject), /* tp_basicsize */
1989 0, /* tp_itemsize */
1990 /* methods */
1991 (destructor)product_dealloc, /* tp_dealloc */
1992 0, /* tp_print */
1993 0, /* tp_getattr */
1994 0, /* tp_setattr */
1995 0, /* tp_compare */
1996 0, /* tp_repr */
1997 0, /* tp_as_number */
1998 0, /* tp_as_sequence */
1999 0, /* tp_as_mapping */
2000 0, /* tp_hash */
2001 0, /* tp_call */
2002 0, /* tp_str */
2003 PyObject_GenericGetAttr, /* tp_getattro */
2004 0, /* tp_setattro */
2005 0, /* tp_as_buffer */
2006 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2007 Py_TPFLAGS_BASETYPE, /* tp_flags */
2008 product_doc, /* tp_doc */
2009 (traverseproc)product_traverse, /* tp_traverse */
2010 0, /* tp_clear */
2011 0, /* tp_richcompare */
2012 0, /* tp_weaklistoffset */
2013 PyObject_SelfIter, /* tp_iter */
2014 (iternextfunc)product_next, /* tp_iternext */
2015 0, /* tp_methods */
2016 0, /* tp_members */
2017 0, /* tp_getset */
2018 0, /* tp_base */
2019 0, /* tp_dict */
2020 0, /* tp_descr_get */
2021 0, /* tp_descr_set */
2022 0, /* tp_dictoffset */
2023 0, /* tp_init */
2024 0, /* tp_alloc */
2025 product_new, /* tp_new */
2026 PyObject_GC_Del, /* tp_free */
2027};
2028
2029
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002030/* combinations object ************************************************************/
2031
2032typedef struct {
2033 PyObject_HEAD
2034 PyObject *pool; /* input converted to a tuple */
2035 Py_ssize_t *indices; /* one index per result element */
2036 PyObject *result; /* most recently returned result tuple */
2037 Py_ssize_t r; /* size of result tuple */
2038 int stopped; /* set to 1 when the combinations iterator is exhausted */
2039} combinationsobject;
2040
2041static PyTypeObject combinations_type;
2042
2043static PyObject *
2044combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2045{
2046 combinationsobject *co;
2047 Py_ssize_t n;
2048 Py_ssize_t r;
2049 PyObject *pool = NULL;
2050 PyObject *iterable = NULL;
2051 Py_ssize_t *indices = NULL;
2052 Py_ssize_t i;
2053 static char *kwargs[] = {"iterable", "r", NULL};
2054
2055 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2056 &iterable, &r))
2057 return NULL;
2058
2059 pool = PySequence_Tuple(iterable);
2060 if (pool == NULL)
2061 goto error;
2062 n = PyTuple_GET_SIZE(pool);
2063 if (r < 0) {
2064 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2065 goto error;
2066 }
2067 if (r > n) {
2068 PyErr_SetString(PyExc_ValueError, "r cannot be bigger than the iterable");
2069 goto error;
2070 }
2071
2072 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2073 if (indices == NULL) {
2074 PyErr_NoMemory();
2075 goto error;
2076 }
2077
2078 for (i=0 ; i<r ; i++)
2079 indices[i] = i;
2080
2081 /* create combinationsobject structure */
2082 co = (combinationsobject *)type->tp_alloc(type, 0);
2083 if (co == NULL)
2084 goto error;
2085
2086 co->pool = pool;
2087 co->indices = indices;
2088 co->result = NULL;
2089 co->r = r;
2090 co->stopped = 0;
2091
2092 return (PyObject *)co;
2093
2094error:
2095 if (indices != NULL)
2096 PyMem_Free(indices);
2097 Py_XDECREF(pool);
2098 return NULL;
2099}
2100
2101static void
2102combinations_dealloc(combinationsobject *co)
2103{
2104 PyObject_GC_UnTrack(co);
2105 Py_XDECREF(co->pool);
2106 Py_XDECREF(co->result);
2107 PyMem_Free(co->indices);
2108 Py_TYPE(co)->tp_free(co);
2109}
2110
2111static int
2112combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2113{
2114 Py_VISIT(co->pool);
2115 Py_VISIT(co->result);
2116 return 0;
2117}
2118
2119static PyObject *
2120combinations_next(combinationsobject *co)
2121{
2122 PyObject *elem;
2123 PyObject *oldelem;
2124 PyObject *pool = co->pool;
2125 Py_ssize_t *indices = co->indices;
2126 PyObject *result = co->result;
2127 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2128 Py_ssize_t r = co->r;
2129 Py_ssize_t i, j, index;
2130
2131 if (co->stopped)
2132 return NULL;
2133
2134 if (result == NULL) {
2135 /* On the first pass, initialize result tuple using the indices */
2136 result = PyTuple_New(r);
2137 if (result == NULL)
2138 goto empty;
2139 co->result = result;
2140 for (i=0; i<r ; i++) {
2141 index = indices[i];
2142 elem = PyTuple_GET_ITEM(pool, index);
2143 Py_INCREF(elem);
2144 PyTuple_SET_ITEM(result, i, elem);
2145 }
2146 } else {
2147 /* Copy the previous result tuple or re-use it if available */
2148 if (Py_REFCNT(result) > 1) {
2149 PyObject *old_result = result;
2150 result = PyTuple_New(r);
2151 if (result == NULL)
2152 goto empty;
2153 co->result = result;
2154 for (i=0; i<r ; i++) {
2155 elem = PyTuple_GET_ITEM(old_result, i);
2156 Py_INCREF(elem);
2157 PyTuple_SET_ITEM(result, i, elem);
2158 }
2159 Py_DECREF(old_result);
2160 }
Christian Heimescdddf182008-02-28 11:18:49 +00002161 /* Now, we've got the only copy so we can update it in-place
2162 * CPython's empty tuple is a singleton and cached in
2163 * PyTuple's freelist.
2164 */
2165 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002166
2167 /* Scan indices right-to-left until finding one that is not
2168 at its maximum (i + n - r). */
2169 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2170 ;
2171
2172 /* If i is negative, then the indices are all at
2173 their maximum value and we're done. */
2174 if (i < 0)
2175 goto empty;
2176
2177 /* Increment the current index which we know is not at its
2178 maximum. Then move back to the right setting each index
2179 to its lowest possible value (one higher than the index
2180 to its left -- this maintains the sort order invariant). */
2181 indices[i]++;
2182 for (j=i+1 ; j<r ; j++)
2183 indices[j] = indices[j-1] + 1;
2184
2185 /* Update the result tuple for the new indices
2186 starting with i, the leftmost index that changed */
2187 for ( ; i<r ; i++) {
2188 index = indices[i];
2189 elem = PyTuple_GET_ITEM(pool, index);
2190 Py_INCREF(elem);
2191 oldelem = PyTuple_GET_ITEM(result, i);
2192 PyTuple_SET_ITEM(result, i, elem);
2193 Py_DECREF(oldelem);
2194 }
2195 }
2196
2197 Py_INCREF(result);
2198 return result;
2199
2200empty:
2201 co->stopped = 1;
2202 return NULL;
2203}
2204
2205PyDoc_STRVAR(combinations_doc,
2206"combinations(iterables) --> combinations object\n\
2207\n\
2208Return successive r-length combinations of elements in the iterable.\n\n\
2209combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2210
2211static PyTypeObject combinations_type = {
2212 PyVarObject_HEAD_INIT(NULL, 0)
2213 "itertools.combinations", /* tp_name */
2214 sizeof(combinationsobject), /* tp_basicsize */
2215 0, /* tp_itemsize */
2216 /* methods */
2217 (destructor)combinations_dealloc, /* tp_dealloc */
2218 0, /* tp_print */
2219 0, /* tp_getattr */
2220 0, /* tp_setattr */
2221 0, /* tp_compare */
2222 0, /* tp_repr */
2223 0, /* tp_as_number */
2224 0, /* tp_as_sequence */
2225 0, /* tp_as_mapping */
2226 0, /* tp_hash */
2227 0, /* tp_call */
2228 0, /* tp_str */
2229 PyObject_GenericGetAttr, /* tp_getattro */
2230 0, /* tp_setattro */
2231 0, /* tp_as_buffer */
2232 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2233 Py_TPFLAGS_BASETYPE, /* tp_flags */
2234 combinations_doc, /* tp_doc */
2235 (traverseproc)combinations_traverse, /* tp_traverse */
2236 0, /* tp_clear */
2237 0, /* tp_richcompare */
2238 0, /* tp_weaklistoffset */
2239 PyObject_SelfIter, /* tp_iter */
2240 (iternextfunc)combinations_next, /* tp_iternext */
2241 0, /* tp_methods */
2242 0, /* tp_members */
2243 0, /* tp_getset */
2244 0, /* tp_base */
2245 0, /* tp_dict */
2246 0, /* tp_descr_get */
2247 0, /* tp_descr_set */
2248 0, /* tp_dictoffset */
2249 0, /* tp_init */
2250 0, /* tp_alloc */
2251 combinations_new, /* tp_new */
2252 PyObject_GC_Del, /* tp_free */
2253};
2254
2255
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002256/* permutations object ************************************************************
2257
2258def permutations(iterable, r=None):
2259 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2260 pool = tuple(iterable)
2261 n = len(pool)
2262 r = n if r is None else r
2263 indices = range(n)
2264 cycles = range(n-r+1, n+1)[::-1]
2265 yield tuple(pool[i] for i in indices[:r])
2266 while n:
2267 for i in reversed(range(r)):
2268 cycles[i] -= 1
2269 if cycles[i] == 0:
2270 indices[i:] = indices[i+1:] + indices[i:i+1]
2271 cycles[i] = n - i
2272 else:
2273 j = cycles[i]
2274 indices[i], indices[-j] = indices[-j], indices[i]
2275 yield tuple(pool[i] for i in indices[:r])
2276 break
2277 else:
2278 return
2279*/
2280
2281typedef struct {
2282 PyObject_HEAD
2283 PyObject *pool; /* input converted to a tuple */
2284 Py_ssize_t *indices; /* one index per element in the pool */
2285 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2286 PyObject *result; /* most recently returned result tuple */
2287 Py_ssize_t r; /* size of result tuple */
2288 int stopped; /* set to 1 when the permutations iterator is exhausted */
2289} permutationsobject;
2290
2291static PyTypeObject permutations_type;
2292
2293static PyObject *
2294permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2295{
2296 permutationsobject *po;
2297 Py_ssize_t n;
2298 Py_ssize_t r;
2299 PyObject *robj = Py_None;
2300 PyObject *pool = NULL;
2301 PyObject *iterable = NULL;
2302 Py_ssize_t *indices = NULL;
2303 Py_ssize_t *cycles = NULL;
2304 Py_ssize_t i;
2305 static char *kwargs[] = {"iterable", "r", NULL};
2306
2307 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2308 &iterable, &robj))
2309 return NULL;
2310
2311 pool = PySequence_Tuple(iterable);
2312 if (pool == NULL)
2313 goto error;
2314 n = PyTuple_GET_SIZE(pool);
2315
2316 r = n;
2317 if (robj != Py_None) {
2318 r = PyInt_AsSsize_t(robj);
2319 if (r == -1 && PyErr_Occurred())
2320 goto error;
2321 }
2322 if (r < 0) {
2323 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2324 goto error;
2325 }
2326 if (r > n) {
2327 PyErr_SetString(PyExc_ValueError, "r cannot be bigger than the iterable");
2328 goto error;
2329 }
2330
2331 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2332 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2333 if (indices == NULL || cycles == NULL) {
2334 PyErr_NoMemory();
2335 goto error;
2336 }
2337
2338 for (i=0 ; i<n ; i++)
2339 indices[i] = i;
2340 for (i=0 ; i<r ; i++)
2341 cycles[i] = n - i;
2342
2343 /* create permutationsobject structure */
2344 po = (permutationsobject *)type->tp_alloc(type, 0);
2345 if (po == NULL)
2346 goto error;
2347
2348 po->pool = pool;
2349 po->indices = indices;
2350 po->cycles = cycles;
2351 po->result = NULL;
2352 po->r = r;
2353 po->stopped = 0;
2354
2355 return (PyObject *)po;
2356
2357error:
2358 if (indices != NULL)
2359 PyMem_Free(indices);
2360 if (cycles != NULL)
2361 PyMem_Free(cycles);
2362 Py_XDECREF(pool);
2363 return NULL;
2364}
2365
2366static void
2367permutations_dealloc(permutationsobject *po)
2368{
2369 PyObject_GC_UnTrack(po);
2370 Py_XDECREF(po->pool);
2371 Py_XDECREF(po->result);
2372 PyMem_Free(po->indices);
2373 PyMem_Free(po->cycles);
2374 Py_TYPE(po)->tp_free(po);
2375}
2376
2377static int
2378permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2379{
Raymond Hettinger6e3e4152008-03-05 21:04:32 +00002380 Py_VISIT(po->pool);
2381 Py_VISIT(po->result);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002382 return 0;
2383}
2384
2385static PyObject *
2386permutations_next(permutationsobject *po)
2387{
2388 PyObject *elem;
2389 PyObject *oldelem;
2390 PyObject *pool = po->pool;
2391 Py_ssize_t *indices = po->indices;
2392 Py_ssize_t *cycles = po->cycles;
2393 PyObject *result = po->result;
2394 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2395 Py_ssize_t r = po->r;
2396 Py_ssize_t i, j, k, index;
2397
2398 if (po->stopped)
2399 return NULL;
2400
2401 if (result == NULL) {
2402 /* On the first pass, initialize result tuple using the indices */
2403 result = PyTuple_New(r);
2404 if (result == NULL)
2405 goto empty;
2406 po->result = result;
2407 for (i=0; i<r ; i++) {
2408 index = indices[i];
2409 elem = PyTuple_GET_ITEM(pool, index);
2410 Py_INCREF(elem);
2411 PyTuple_SET_ITEM(result, i, elem);
2412 }
2413 } else {
2414 if (n == 0)
2415 goto empty;
2416
2417 /* Copy the previous result tuple or re-use it if available */
2418 if (Py_REFCNT(result) > 1) {
2419 PyObject *old_result = result;
2420 result = PyTuple_New(r);
2421 if (result == NULL)
2422 goto empty;
2423 po->result = result;
2424 for (i=0; i<r ; i++) {
2425 elem = PyTuple_GET_ITEM(old_result, i);
2426 Py_INCREF(elem);
2427 PyTuple_SET_ITEM(result, i, elem);
2428 }
2429 Py_DECREF(old_result);
2430 }
2431 /* Now, we've got the only copy so we can update it in-place */
2432 assert(r == 0 || Py_REFCNT(result) == 1);
2433
2434 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2435 for (i=r-1 ; i>=0 ; i--) {
2436 cycles[i] -= 1;
2437 if (cycles[i] == 0) {
2438 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2439 index = indices[i];
2440 for (j=i ; j<n-1 ; j++)
2441 indices[j] = indices[j+1];
2442 indices[n-1] = index;
2443 cycles[i] = n - i;
2444 } else {
2445 j = cycles[i];
2446 index = indices[i];
2447 indices[i] = indices[n-j];
2448 indices[n-j] = index;
2449
2450 for (k=i; k<r ; k++) {
2451 /* start with i, the leftmost element that changed */
2452 /* yield tuple(pool[k] for k in indices[:r]) */
2453 index = indices[k];
2454 elem = PyTuple_GET_ITEM(pool, index);
2455 Py_INCREF(elem);
2456 oldelem = PyTuple_GET_ITEM(result, k);
2457 PyTuple_SET_ITEM(result, k, elem);
2458 Py_DECREF(oldelem);
2459 }
2460 break;
2461 }
2462 }
2463 /* If i is negative, then the cycles have all
2464 rolled-over and we're done. */
2465 if (i < 0)
2466 goto empty;
2467 }
2468 Py_INCREF(result);
2469 return result;
2470
2471empty:
2472 po->stopped = 1;
2473 return NULL;
2474}
2475
2476PyDoc_STRVAR(permutations_doc,
2477"permutations(iterables[, r]) --> permutations object\n\
2478\n\
2479Return successive r-length permutations of elements in the iterable.\n\n\
2480permutations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2481
2482static PyTypeObject permutations_type = {
2483 PyVarObject_HEAD_INIT(NULL, 0)
2484 "itertools.permutations", /* tp_name */
2485 sizeof(permutationsobject), /* tp_basicsize */
2486 0, /* tp_itemsize */
2487 /* methods */
2488 (destructor)permutations_dealloc, /* tp_dealloc */
2489 0, /* tp_print */
2490 0, /* tp_getattr */
2491 0, /* tp_setattr */
2492 0, /* tp_compare */
2493 0, /* tp_repr */
2494 0, /* tp_as_number */
2495 0, /* tp_as_sequence */
2496 0, /* tp_as_mapping */
2497 0, /* tp_hash */
2498 0, /* tp_call */
2499 0, /* tp_str */
2500 PyObject_GenericGetAttr, /* tp_getattro */
2501 0, /* tp_setattro */
2502 0, /* tp_as_buffer */
2503 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2504 Py_TPFLAGS_BASETYPE, /* tp_flags */
2505 permutations_doc, /* tp_doc */
2506 (traverseproc)permutations_traverse, /* tp_traverse */
2507 0, /* tp_clear */
2508 0, /* tp_richcompare */
2509 0, /* tp_weaklistoffset */
2510 PyObject_SelfIter, /* tp_iter */
2511 (iternextfunc)permutations_next, /* tp_iternext */
2512 0, /* tp_methods */
2513 0, /* tp_members */
2514 0, /* tp_getset */
2515 0, /* tp_base */
2516 0, /* tp_dict */
2517 0, /* tp_descr_get */
2518 0, /* tp_descr_set */
2519 0, /* tp_dictoffset */
2520 0, /* tp_init */
2521 0, /* tp_alloc */
2522 permutations_new, /* tp_new */
2523 PyObject_GC_Del, /* tp_free */
2524};
2525
2526
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002527/* ifilter object ************************************************************/
2528
2529typedef struct {
2530 PyObject_HEAD
2531 PyObject *func;
2532 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002533} ifilterobject;
2534
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002535static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002536
2537static PyObject *
2538ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2539{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002540 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002541 PyObject *it;
2542 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002543
Raymond Hettinger10f40a62008-03-13 16:43:17 +00002544 if (Py_Py3kWarningFlag &&
2545 PyErr_Warn(PyExc_DeprecationWarning,
2546 "In 3.x, itertools.ifilter() was moved to builtin filter().") < 0)
2547 return NULL;
2548
Georg Brandlb84c1372007-01-21 10:28:43 +00002549 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002550 return NULL;
2551
Raymond Hettinger60eca932003-02-09 06:40:58 +00002552 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002553 return NULL;
2554
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002555 /* Get iterator. */
2556 it = PyObject_GetIter(seq);
2557 if (it == NULL)
2558 return NULL;
2559
2560 /* create ifilterobject structure */
2561 lz = (ifilterobject *)type->tp_alloc(type, 0);
2562 if (lz == NULL) {
2563 Py_DECREF(it);
2564 return NULL;
2565 }
2566 Py_INCREF(func);
2567 lz->func = func;
2568 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002569
2570 return (PyObject *)lz;
2571}
2572
2573static void
2574ifilter_dealloc(ifilterobject *lz)
2575{
2576 PyObject_GC_UnTrack(lz);
2577 Py_XDECREF(lz->func);
2578 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00002579 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002580}
2581
2582static int
2583ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
2584{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002585 Py_VISIT(lz->it);
2586 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002587 return 0;
2588}
2589
2590static PyObject *
2591ifilter_next(ifilterobject *lz)
2592{
2593 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002594 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002595 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002596 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002597
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002598 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00002599 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002600 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002601 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002602 if (item == NULL)
2603 return NULL;
2604
Facundo Batistab1d70e22008-02-25 23:46:02 +00002605 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002606 ok = PyObject_IsTrue(item);
2607 } else {
2608 PyObject *good;
2609 good = PyObject_CallFunctionObjArgs(lz->func,
2610 item, NULL);
2611 if (good == NULL) {
2612 Py_DECREF(item);
2613 return NULL;
2614 }
2615 ok = PyObject_IsTrue(good);
2616 Py_DECREF(good);
2617 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00002618 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002619 return item;
2620 Py_DECREF(item);
2621 }
2622}
2623
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002624PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002625"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002626\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002627Return those items of sequence for which function(item) is true.\n\
2628If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002629
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002630static PyTypeObject ifilter_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002631 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002632 "itertools.ifilter", /* tp_name */
2633 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002634 0, /* tp_itemsize */
2635 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002636 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002637 0, /* tp_print */
2638 0, /* tp_getattr */
2639 0, /* tp_setattr */
2640 0, /* tp_compare */
2641 0, /* tp_repr */
2642 0, /* tp_as_number */
2643 0, /* tp_as_sequence */
2644 0, /* tp_as_mapping */
2645 0, /* tp_hash */
2646 0, /* tp_call */
2647 0, /* tp_str */
2648 PyObject_GenericGetAttr, /* tp_getattro */
2649 0, /* tp_setattro */
2650 0, /* tp_as_buffer */
2651 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2652 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002653 ifilter_doc, /* tp_doc */
2654 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002655 0, /* tp_clear */
2656 0, /* tp_richcompare */
2657 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002658 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002659 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002660 0, /* tp_methods */
2661 0, /* tp_members */
2662 0, /* tp_getset */
2663 0, /* tp_base */
2664 0, /* tp_dict */
2665 0, /* tp_descr_get */
2666 0, /* tp_descr_set */
2667 0, /* tp_dictoffset */
2668 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002669 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002670 ifilter_new, /* tp_new */
2671 PyObject_GC_Del, /* tp_free */
2672};
2673
2674
2675/* ifilterfalse object ************************************************************/
2676
2677typedef struct {
2678 PyObject_HEAD
2679 PyObject *func;
2680 PyObject *it;
2681} ifilterfalseobject;
2682
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002683static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002684
2685static PyObject *
2686ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2687{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00002688 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002689 PyObject *it;
2690 ifilterfalseobject *lz;
2691
Raymond Hettinger10f40a62008-03-13 16:43:17 +00002692 if (Py_Py3kWarningFlag &&
2693 PyErr_Warn(PyExc_DeprecationWarning,
2694 "In 3.x, ifilterfalse() was renamed to filterfalse().") < 0)
2695 return NULL;
2696
Georg Brandlb84c1372007-01-21 10:28:43 +00002697 if (type == &ifilterfalse_type &&
2698 !_PyArg_NoKeywords("ifilterfalse()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002699 return NULL;
2700
Raymond Hettinger60eca932003-02-09 06:40:58 +00002701 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
2702 return NULL;
2703
2704 /* Get iterator. */
2705 it = PyObject_GetIter(seq);
2706 if (it == NULL)
2707 return NULL;
2708
2709 /* create ifilterfalseobject structure */
2710 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
2711 if (lz == NULL) {
2712 Py_DECREF(it);
2713 return NULL;
2714 }
2715 Py_INCREF(func);
2716 lz->func = func;
2717 lz->it = it;
2718
2719 return (PyObject *)lz;
2720}
2721
2722static void
2723ifilterfalse_dealloc(ifilterfalseobject *lz)
2724{
2725 PyObject_GC_UnTrack(lz);
2726 Py_XDECREF(lz->func);
2727 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00002728 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002729}
2730
2731static int
2732ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
2733{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002734 Py_VISIT(lz->it);
2735 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002736 return 0;
2737}
2738
2739static PyObject *
2740ifilterfalse_next(ifilterfalseobject *lz)
2741{
2742 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002743 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002744 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002745 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002746
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002747 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00002748 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002749 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002750 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002751 if (item == NULL)
2752 return NULL;
2753
Facundo Batistab1d70e22008-02-25 23:46:02 +00002754 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
Raymond Hettinger60eca932003-02-09 06:40:58 +00002755 ok = PyObject_IsTrue(item);
2756 } else {
2757 PyObject *good;
2758 good = PyObject_CallFunctionObjArgs(lz->func,
2759 item, NULL);
2760 if (good == NULL) {
2761 Py_DECREF(item);
2762 return NULL;
2763 }
2764 ok = PyObject_IsTrue(good);
2765 Py_DECREF(good);
2766 }
2767 if (!ok)
2768 return item;
2769 Py_DECREF(item);
2770 }
2771}
2772
Raymond Hettinger60eca932003-02-09 06:40:58 +00002773PyDoc_STRVAR(ifilterfalse_doc,
2774"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
2775\n\
2776Return those items of sequence for which function(item) is false.\n\
2777If function is None, return the items that are false.");
2778
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002779static PyTypeObject ifilterfalse_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002780 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002781 "itertools.ifilterfalse", /* tp_name */
2782 sizeof(ifilterfalseobject), /* tp_basicsize */
2783 0, /* tp_itemsize */
2784 /* methods */
2785 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
2786 0, /* tp_print */
2787 0, /* tp_getattr */
2788 0, /* tp_setattr */
2789 0, /* tp_compare */
2790 0, /* tp_repr */
2791 0, /* tp_as_number */
2792 0, /* tp_as_sequence */
2793 0, /* tp_as_mapping */
2794 0, /* tp_hash */
2795 0, /* tp_call */
2796 0, /* tp_str */
2797 PyObject_GenericGetAttr, /* tp_getattro */
2798 0, /* tp_setattro */
2799 0, /* tp_as_buffer */
2800 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2801 Py_TPFLAGS_BASETYPE, /* tp_flags */
2802 ifilterfalse_doc, /* tp_doc */
2803 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
2804 0, /* tp_clear */
2805 0, /* tp_richcompare */
2806 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002807 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002808 (iternextfunc)ifilterfalse_next, /* tp_iternext */
2809 0, /* tp_methods */
2810 0, /* tp_members */
2811 0, /* tp_getset */
2812 0, /* tp_base */
2813 0, /* tp_dict */
2814 0, /* tp_descr_get */
2815 0, /* tp_descr_set */
2816 0, /* tp_dictoffset */
2817 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002818 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002819 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002820 PyObject_GC_Del, /* tp_free */
2821};
2822
2823
2824/* count object ************************************************************/
2825
2826typedef struct {
2827 PyObject_HEAD
Jack Diederich6c433a92006-05-26 11:15:17 +00002828 Py_ssize_t cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002829 PyObject *long_cnt; /* Arbitrarily large count when cnt >= PY_SSIZE_T_MAX */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002830} countobject;
2831
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002832static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002833
2834static PyObject *
2835count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2836{
2837 countobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00002838 Py_ssize_t cnt = 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002839 PyObject *cnt_arg = NULL;
2840 PyObject *long_cnt = NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002841
Georg Brandlb84c1372007-01-21 10:28:43 +00002842 if (type == &count_type && !_PyArg_NoKeywords("count()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002843 return NULL;
2844
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002845 if (!PyArg_UnpackTuple(args, "count", 0, 1, &cnt_arg))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002846 return NULL;
2847
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002848 if (cnt_arg != NULL) {
2849 cnt = PyInt_AsSsize_t(cnt_arg);
2850 if (cnt == -1 && PyErr_Occurred()) {
2851 PyErr_Clear();
2852 if (!PyLong_Check(cnt_arg)) {
2853 PyErr_SetString(PyExc_TypeError, "an integer is required");
2854 return NULL;
2855 }
2856 long_cnt = cnt_arg;
2857 Py_INCREF(long_cnt);
2858 cnt = PY_SSIZE_T_MAX;
2859 }
2860 }
2861
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002862 /* create countobject structure */
2863 lz = (countobject *)PyObject_New(countobject, &count_type);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002864 if (lz == NULL) {
2865 Py_XDECREF(long_cnt);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002866 return NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002867 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002868 lz->cnt = cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002869 lz->long_cnt = long_cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002870
2871 return (PyObject *)lz;
2872}
2873
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002874static void
2875count_dealloc(countobject *lz)
2876{
2877 Py_XDECREF(lz->long_cnt);
2878 PyObject_Del(lz);
2879}
2880
2881static PyObject *
2882count_nextlong(countobject *lz)
2883{
2884 static PyObject *one = NULL;
2885 PyObject *cnt;
2886 PyObject *stepped_up;
2887
2888 if (lz->long_cnt == NULL) {
2889 lz->long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
2890 if (lz->long_cnt == NULL)
2891 return NULL;
2892 }
2893 if (one == NULL) {
2894 one = PyInt_FromLong(1);
2895 if (one == NULL)
2896 return NULL;
2897 }
2898 cnt = lz->long_cnt;
2899 assert(cnt != NULL);
2900 stepped_up = PyNumber_Add(cnt, one);
2901 if (stepped_up == NULL)
2902 return NULL;
2903 lz->long_cnt = stepped_up;
2904 return cnt;
2905}
2906
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002907static PyObject *
2908count_next(countobject *lz)
2909{
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002910 if (lz->cnt == PY_SSIZE_T_MAX)
2911 return count_nextlong(lz);
Jack Diederich36234e82006-09-21 17:50:26 +00002912 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002913}
2914
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002915static PyObject *
2916count_repr(countobject *lz)
2917{
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002918 PyObject *cnt_repr;
2919 PyObject *result;
2920
2921 if (lz->cnt != PY_SSIZE_T_MAX)
2922 return PyString_FromFormat("count(%zd)", lz->cnt);
2923
2924 cnt_repr = PyObject_Repr(lz->long_cnt);
2925 if (cnt_repr == NULL)
2926 return NULL;
2927 result = PyString_FromFormat("count(%s)", PyString_AS_STRING(cnt_repr));
2928 Py_DECREF(cnt_repr);
2929 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002930}
2931
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002932PyDoc_STRVAR(count_doc,
2933"count([firstval]) --> count object\n\
2934\n\
2935Return a count object whose .next() method returns consecutive\n\
2936integers starting from zero or, if specified, from firstval.");
2937
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002938static PyTypeObject count_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002939 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002940 "itertools.count", /* tp_name */
2941 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002942 0, /* tp_itemsize */
2943 /* methods */
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002944 (destructor)count_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002945 0, /* tp_print */
2946 0, /* tp_getattr */
2947 0, /* tp_setattr */
2948 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002949 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002950 0, /* tp_as_number */
2951 0, /* tp_as_sequence */
2952 0, /* tp_as_mapping */
2953 0, /* tp_hash */
2954 0, /* tp_call */
2955 0, /* tp_str */
2956 PyObject_GenericGetAttr, /* tp_getattro */
2957 0, /* tp_setattro */
2958 0, /* tp_as_buffer */
2959 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002960 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002961 0, /* tp_traverse */
2962 0, /* tp_clear */
2963 0, /* tp_richcompare */
2964 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002965 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002966 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002967 0, /* tp_methods */
2968 0, /* tp_members */
2969 0, /* tp_getset */
2970 0, /* tp_base */
2971 0, /* tp_dict */
2972 0, /* tp_descr_get */
2973 0, /* tp_descr_set */
2974 0, /* tp_dictoffset */
2975 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002976 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002977 count_new, /* tp_new */
2978};
2979
2980
2981/* izip object ************************************************************/
2982
2983#include "Python.h"
2984
2985typedef struct {
2986 PyObject_HEAD
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002987 Py_ssize_t tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002988 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002989 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002990} izipobject;
2991
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002992static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002993
2994static PyObject *
2995izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2996{
2997 izipobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00002998 Py_ssize_t i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002999 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00003000 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00003001 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003002
Raymond Hettinger10f40a62008-03-13 16:43:17 +00003003 if (Py_Py3kWarningFlag &&
3004 PyErr_Warn(PyExc_DeprecationWarning,
3005 "In 3.x, itertools.izip() was moved to builtin zip()") < 0)
3006 return NULL;
3007
Georg Brandlb84c1372007-01-21 10:28:43 +00003008 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00003009 return NULL;
3010
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003011 /* args must be a tuple */
3012 assert(PyTuple_Check(args));
3013
3014 /* obtain iterators */
3015 ittuple = PyTuple_New(tuplesize);
Neal Norwitz2f3136b2006-05-27 05:18:57 +00003016 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003017 return NULL;
3018 for (i=0; i < tuplesize; ++i) {
3019 PyObject *item = PyTuple_GET_ITEM(args, i);
3020 PyObject *it = PyObject_GetIter(item);
3021 if (it == NULL) {
3022 if (PyErr_ExceptionMatches(PyExc_TypeError))
3023 PyErr_Format(PyExc_TypeError,
Neal Norwitz2f3136b2006-05-27 05:18:57 +00003024 "izip argument #%zd must support iteration",
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003025 i+1);
3026 Py_DECREF(ittuple);
3027 return NULL;
3028 }
3029 PyTuple_SET_ITEM(ittuple, i, it);
3030 }
3031
Raymond Hettinger2012f172003-02-07 05:32:58 +00003032 /* create a result holder */
3033 result = PyTuple_New(tuplesize);
3034 if (result == NULL) {
3035 Py_DECREF(ittuple);
3036 return NULL;
3037 }
3038 for (i=0 ; i < tuplesize ; i++) {
3039 Py_INCREF(Py_None);
3040 PyTuple_SET_ITEM(result, i, Py_None);
3041 }
3042
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003043 /* create izipobject structure */
3044 lz = (izipobject *)type->tp_alloc(type, 0);
3045 if (lz == NULL) {
3046 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003047 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003048 return NULL;
3049 }
3050 lz->ittuple = ittuple;
3051 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00003052 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003053
3054 return (PyObject *)lz;
3055}
3056
3057static void
3058izip_dealloc(izipobject *lz)
3059{
3060 PyObject_GC_UnTrack(lz);
3061 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003062 Py_XDECREF(lz->result);
Christian Heimese93237d2007-12-19 02:37:44 +00003063 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003064}
3065
3066static int
3067izip_traverse(izipobject *lz, visitproc visit, void *arg)
3068{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00003069 Py_VISIT(lz->ittuple);
3070 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003071 return 0;
3072}
3073
3074static PyObject *
3075izip_next(izipobject *lz)
3076{
Jack Diederich6c433a92006-05-26 11:15:17 +00003077 Py_ssize_t i;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00003078 Py_ssize_t tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00003079 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003080 PyObject *it;
3081 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00003082 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003083
Raymond Hettingerb5a42082003-08-08 05:10:41 +00003084 if (tuplesize == 0)
3085 return NULL;
Christian Heimese93237d2007-12-19 02:37:44 +00003086 if (Py_REFCNT(result) == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00003087 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003088 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003089 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00003090 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00003091 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00003092 if (item == NULL) {
3093 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003094 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00003095 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00003096 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003097 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00003098 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003099 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003100 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00003101 result = PyTuple_New(tuplesize);
3102 if (result == NULL)
3103 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003104 for (i=0 ; i < tuplesize ; i++) {
3105 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00003106 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00003107 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003108 if (item == NULL) {
3109 Py_DECREF(result);
3110 return NULL;
3111 }
3112 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003113 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003114 }
3115 return result;
3116}
3117
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003118PyDoc_STRVAR(izip_doc,
3119"izip(iter1 [,iter2 [...]]) --> izip object\n\
3120\n\
3121Return a izip object whose .next() method returns a tuple where\n\
3122the i-th element comes from the i-th iterable argument. The .next()\n\
3123method continues until the shortest iterable in the argument sequence\n\
3124is exhausted and then it raises StopIteration. Works like the zip()\n\
3125function but consumes less memory by returning an iterator instead of\n\
3126a list.");
3127
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003128static PyTypeObject izip_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003129 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003130 "itertools.izip", /* tp_name */
3131 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003132 0, /* tp_itemsize */
3133 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003134 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003135 0, /* tp_print */
3136 0, /* tp_getattr */
3137 0, /* tp_setattr */
3138 0, /* tp_compare */
3139 0, /* tp_repr */
3140 0, /* tp_as_number */
3141 0, /* tp_as_sequence */
3142 0, /* tp_as_mapping */
3143 0, /* tp_hash */
3144 0, /* tp_call */
3145 0, /* tp_str */
3146 PyObject_GenericGetAttr, /* tp_getattro */
3147 0, /* tp_setattro */
3148 0, /* tp_as_buffer */
3149 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3150 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003151 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003152 (traverseproc)izip_traverse, /* tp_traverse */
3153 0, /* tp_clear */
3154 0, /* tp_richcompare */
3155 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003156 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003157 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003158 0, /* tp_methods */
3159 0, /* tp_members */
3160 0, /* tp_getset */
3161 0, /* tp_base */
3162 0, /* tp_dict */
3163 0, /* tp_descr_get */
3164 0, /* tp_descr_set */
3165 0, /* tp_dictoffset */
3166 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003167 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003168 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003169 PyObject_GC_Del, /* tp_free */
3170};
3171
3172
3173/* repeat object ************************************************************/
3174
3175typedef struct {
3176 PyObject_HEAD
3177 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00003178 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003179} repeatobject;
3180
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003181static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003182
3183static PyObject *
3184repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3185{
3186 repeatobject *ro;
3187 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00003188 Py_ssize_t cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003189
Georg Brandlb84c1372007-01-21 10:28:43 +00003190 if (type == &repeat_type && !_PyArg_NoKeywords("repeat()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00003191 return NULL;
3192
Jack Diederich6c433a92006-05-26 11:15:17 +00003193 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003194 return NULL;
3195
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003196 if (PyTuple_Size(args) == 2 && cnt < 0)
3197 cnt = 0;
3198
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003199 ro = (repeatobject *)type->tp_alloc(type, 0);
3200 if (ro == NULL)
3201 return NULL;
3202 Py_INCREF(element);
3203 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003204 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003205 return (PyObject *)ro;
3206}
3207
3208static void
3209repeat_dealloc(repeatobject *ro)
3210{
3211 PyObject_GC_UnTrack(ro);
3212 Py_XDECREF(ro->element);
Christian Heimese93237d2007-12-19 02:37:44 +00003213 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003214}
3215
3216static int
3217repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3218{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00003219 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003220 return 0;
3221}
3222
3223static PyObject *
3224repeat_next(repeatobject *ro)
3225{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003226 if (ro->cnt == 0)
3227 return NULL;
3228 if (ro->cnt > 0)
3229 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003230 Py_INCREF(ro->element);
3231 return ro->element;
3232}
3233
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003234static PyObject *
3235repeat_repr(repeatobject *ro)
3236{
3237 PyObject *result, *objrepr;
3238
3239 objrepr = PyObject_Repr(ro->element);
3240 if (objrepr == NULL)
3241 return NULL;
3242
3243 if (ro->cnt == -1)
3244 result = PyString_FromFormat("repeat(%s)",
3245 PyString_AS_STRING(objrepr));
3246 else
Jack Diederich6c433a92006-05-26 11:15:17 +00003247 result = PyString_FromFormat("repeat(%s, %zd)",
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003248 PyString_AS_STRING(objrepr), ro->cnt);
3249 Py_DECREF(objrepr);
3250 return result;
3251}
3252
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003253static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003254repeat_len(repeatobject *ro)
3255{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003256 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003257 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003258 return NULL;
3259 }
Jack Diederich6c433a92006-05-26 11:15:17 +00003260 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003261}
3262
Armin Rigof5b3e362006-02-11 21:32:43 +00003263PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003264
3265static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00003266 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003267 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003268};
3269
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003270PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003271"repeat(element [,times]) -> create an iterator which returns the element\n\
3272for the specified number of times. If not specified, returns the element\n\
3273endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003274
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003275static PyTypeObject repeat_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003276 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003277 "itertools.repeat", /* tp_name */
3278 sizeof(repeatobject), /* tp_basicsize */
3279 0, /* tp_itemsize */
3280 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003281 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003282 0, /* tp_print */
3283 0, /* tp_getattr */
3284 0, /* tp_setattr */
3285 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003286 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003287 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003288 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003289 0, /* tp_as_mapping */
3290 0, /* tp_hash */
3291 0, /* tp_call */
3292 0, /* tp_str */
3293 PyObject_GenericGetAttr, /* tp_getattro */
3294 0, /* tp_setattro */
3295 0, /* tp_as_buffer */
3296 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3297 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003298 repeat_doc, /* tp_doc */
3299 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003300 0, /* tp_clear */
3301 0, /* tp_richcompare */
3302 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003303 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003304 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003305 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003306 0, /* tp_members */
3307 0, /* tp_getset */
3308 0, /* tp_base */
3309 0, /* tp_dict */
3310 0, /* tp_descr_get */
3311 0, /* tp_descr_set */
3312 0, /* tp_dictoffset */
3313 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003314 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003315 repeat_new, /* tp_new */
3316 PyObject_GC_Del, /* tp_free */
3317};
3318
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003319/* iziplongest object ************************************************************/
3320
3321#include "Python.h"
3322
3323typedef struct {
3324 PyObject_HEAD
3325 Py_ssize_t tuplesize;
3326 Py_ssize_t numactive;
3327 PyObject *ittuple; /* tuple of iterators */
3328 PyObject *result;
3329 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003330} iziplongestobject;
3331
3332static PyTypeObject iziplongest_type;
3333
3334static PyObject *
3335izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3336{
3337 iziplongestobject *lz;
3338 Py_ssize_t i;
3339 PyObject *ittuple; /* tuple of iterators */
3340 PyObject *result;
3341 PyObject *fillvalue = Py_None;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003342 Py_ssize_t tuplesize = PySequence_Length(args);
3343
Raymond Hettinger10f40a62008-03-13 16:43:17 +00003344 if (Py_Py3kWarningFlag &&
3345 PyErr_Warn(PyExc_DeprecationWarning,
3346 "In 3.x, izip_longest() is renamed to zip_longest().") < 0)
3347 return NULL;
3348
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003349 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3350 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3351 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3352 PyErr_SetString(PyExc_TypeError,
3353 "izip_longest() got an unexpected keyword argument");
3354 return NULL;
3355 }
3356 }
3357
3358 /* args must be a tuple */
3359 assert(PyTuple_Check(args));
3360
3361 /* obtain iterators */
3362 ittuple = PyTuple_New(tuplesize);
3363 if (ittuple == NULL)
3364 return NULL;
3365 for (i=0; i < tuplesize; ++i) {
3366 PyObject *item = PyTuple_GET_ITEM(args, i);
3367 PyObject *it = PyObject_GetIter(item);
3368 if (it == NULL) {
3369 if (PyErr_ExceptionMatches(PyExc_TypeError))
3370 PyErr_Format(PyExc_TypeError,
3371 "izip_longest argument #%zd must support iteration",
3372 i+1);
3373 Py_DECREF(ittuple);
3374 return NULL;
3375 }
3376 PyTuple_SET_ITEM(ittuple, i, it);
3377 }
3378
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003379 /* create a result holder */
3380 result = PyTuple_New(tuplesize);
3381 if (result == NULL) {
3382 Py_DECREF(ittuple);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003383 return NULL;
3384 }
3385 for (i=0 ; i < tuplesize ; i++) {
3386 Py_INCREF(Py_None);
3387 PyTuple_SET_ITEM(result, i, Py_None);
3388 }
3389
3390 /* create iziplongestobject structure */
3391 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3392 if (lz == NULL) {
3393 Py_DECREF(ittuple);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003394 Py_DECREF(result);
3395 return NULL;
3396 }
3397 lz->ittuple = ittuple;
3398 lz->tuplesize = tuplesize;
3399 lz->numactive = tuplesize;
3400 lz->result = result;
3401 Py_INCREF(fillvalue);
3402 lz->fillvalue = fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003403 return (PyObject *)lz;
3404}
3405
3406static void
3407izip_longest_dealloc(iziplongestobject *lz)
3408{
3409 PyObject_GC_UnTrack(lz);
3410 Py_XDECREF(lz->ittuple);
3411 Py_XDECREF(lz->result);
3412 Py_XDECREF(lz->fillvalue);
Christian Heimese93237d2007-12-19 02:37:44 +00003413 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003414}
3415
3416static int
3417izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3418{
3419 Py_VISIT(lz->ittuple);
3420 Py_VISIT(lz->result);
3421 Py_VISIT(lz->fillvalue);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003422 return 0;
3423}
3424
3425static PyObject *
3426izip_longest_next(iziplongestobject *lz)
3427{
3428 Py_ssize_t i;
3429 Py_ssize_t tuplesize = lz->tuplesize;
3430 PyObject *result = lz->result;
3431 PyObject *it;
3432 PyObject *item;
3433 PyObject *olditem;
3434
3435 if (tuplesize == 0)
3436 return NULL;
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003437 if (lz->numactive == 0)
3438 return NULL;
Christian Heimese93237d2007-12-19 02:37:44 +00003439 if (Py_REFCNT(result) == 1) {
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003440 Py_INCREF(result);
3441 for (i=0 ; i < tuplesize ; i++) {
3442 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003443 if (it == NULL) {
3444 Py_INCREF(lz->fillvalue);
3445 item = lz->fillvalue;
3446 } else {
3447 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00003448 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003449 if (item == NULL) {
3450 lz->numactive -= 1;
3451 if (lz->numactive == 0) {
3452 Py_DECREF(result);
3453 return NULL;
3454 } else {
3455 Py_INCREF(lz->fillvalue);
3456 item = lz->fillvalue;
3457 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3458 Py_DECREF(it);
3459 }
3460 }
3461 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003462 olditem = PyTuple_GET_ITEM(result, i);
3463 PyTuple_SET_ITEM(result, i, item);
3464 Py_DECREF(olditem);
3465 }
3466 } else {
3467 result = PyTuple_New(tuplesize);
3468 if (result == NULL)
3469 return NULL;
3470 for (i=0 ; i < tuplesize ; i++) {
3471 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003472 if (it == NULL) {
3473 Py_INCREF(lz->fillvalue);
3474 item = lz->fillvalue;
3475 } else {
3476 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00003477 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003478 if (item == NULL) {
3479 lz->numactive -= 1;
3480 if (lz->numactive == 0) {
3481 Py_DECREF(result);
3482 return NULL;
3483 } else {
3484 Py_INCREF(lz->fillvalue);
3485 item = lz->fillvalue;
3486 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3487 Py_DECREF(it);
3488 }
3489 }
3490 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003491 PyTuple_SET_ITEM(result, i, item);
3492 }
3493 }
3494 return result;
3495}
3496
3497PyDoc_STRVAR(izip_longest_doc,
3498"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3499\n\
3500Return an izip_longest object whose .next() method returns a tuple where\n\
3501the i-th element comes from the i-th iterable argument. The .next()\n\
3502method continues until the longest iterable in the argument sequence\n\
3503is exhausted and then it raises StopIteration. When the shorter iterables\n\
3504are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3505defaults to None or can be specified by a keyword argument.\n\
3506");
3507
3508static PyTypeObject iziplongest_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003509 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003510 "itertools.izip_longest", /* tp_name */
3511 sizeof(iziplongestobject), /* tp_basicsize */
3512 0, /* tp_itemsize */
3513 /* methods */
3514 (destructor)izip_longest_dealloc, /* tp_dealloc */
3515 0, /* tp_print */
3516 0, /* tp_getattr */
3517 0, /* tp_setattr */
3518 0, /* tp_compare */
3519 0, /* tp_repr */
3520 0, /* tp_as_number */
3521 0, /* tp_as_sequence */
3522 0, /* tp_as_mapping */
3523 0, /* tp_hash */
3524 0, /* tp_call */
3525 0, /* tp_str */
3526 PyObject_GenericGetAttr, /* tp_getattro */
3527 0, /* tp_setattro */
3528 0, /* tp_as_buffer */
3529 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3530 Py_TPFLAGS_BASETYPE, /* tp_flags */
3531 izip_longest_doc, /* tp_doc */
3532 (traverseproc)izip_longest_traverse, /* tp_traverse */
3533 0, /* tp_clear */
3534 0, /* tp_richcompare */
3535 0, /* tp_weaklistoffset */
3536 PyObject_SelfIter, /* tp_iter */
3537 (iternextfunc)izip_longest_next, /* tp_iternext */
3538 0, /* tp_methods */
3539 0, /* tp_members */
3540 0, /* tp_getset */
3541 0, /* tp_base */
3542 0, /* tp_dict */
3543 0, /* tp_descr_get */
3544 0, /* tp_descr_set */
3545 0, /* tp_dictoffset */
3546 0, /* tp_init */
3547 0, /* tp_alloc */
3548 izip_longest_new, /* tp_new */
3549 PyObject_GC_Del, /* tp_free */
3550};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003551
3552/* module level code ********************************************************/
3553
3554PyDoc_STRVAR(module_doc,
3555"Functional tools for creating and using iterators.\n\
3556\n\
3557Infinite iterators:\n\
3558count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003559cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003560repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003561\n\
3562Iterators terminating on the shortest input sequence:\n\
3563izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003564izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003565ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
3566ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003567islice(seq, [start,] stop [, step]) --> elements from\n\
3568 seq[start:stop:step]\n\
3569imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
3570starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003571tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003572chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003573takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
3574dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003575groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003576");
3577
3578
Raymond Hettingerad983e72003-11-12 14:32:26 +00003579static PyMethodDef module_methods[] = {
3580 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3581 {NULL, NULL} /* sentinel */
3582};
3583
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003584PyMODINIT_FUNC
3585inititertools(void)
3586{
Raymond Hettinger60eca932003-02-09 06:40:58 +00003587 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003588 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003589 char *name;
3590 PyTypeObject *typelist[] = {
Raymond Hettinger93e804d2008-02-26 23:40:50 +00003591 &combinations_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003592 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003593 &dropwhile_type,
3594 &takewhile_type,
3595 &islice_type,
3596 &starmap_type,
3597 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003598 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003599 &ifilter_type,
3600 &ifilterfalse_type,
3601 &count_type,
3602 &izip_type,
Raymond Hettinger50986cc2008-02-22 03:16:42 +00003603 &iziplongest_type,
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00003604 &permutations_type,
Raymond Hettinger93e804d2008-02-26 23:40:50 +00003605 &product_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003606 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003607 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003608 NULL
3609 };
3610
Christian Heimese93237d2007-12-19 02:37:44 +00003611 Py_TYPE(&teedataobject_type) = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00003612 m = Py_InitModule3("itertools", module_methods, module_doc);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00003613 if (m == NULL)
3614 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003615
Raymond Hettinger60eca932003-02-09 06:40:58 +00003616 for (i=0 ; typelist[i] != NULL ; i++) {
3617 if (PyType_Ready(typelist[i]) < 0)
3618 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003619 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003620 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003621 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003622 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003623 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003624
3625 if (PyType_Ready(&teedataobject_type) < 0)
3626 return;
3627 if (PyType_Ready(&tee_type) < 0)
3628 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003629 if (PyType_Ready(&_grouper_type) < 0)
3630 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003631}