blob: fe06ef48e38f2b6e5a1f6ff084c68fccddd36804 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
2#include "Python.h"
Fred Drake08ebfec2004-10-17 19:36:57 +00003#include "structmember.h"
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004
5/* Itertools module written and maintained
6 by Raymond D. Hettinger <python@rcn.com>
7 Copyright (c) 2003 Python Software Foundation.
8 All rights reserved.
9*/
10
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000011
12/* groupby object ***********************************************************/
13
14typedef struct {
15 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
21} groupbyobject;
22
23static PyTypeObject groupby_type;
24static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26static PyObject *
27groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28{
Martin v. Löwis02cbf4a2006-02-27 17:20:04 +000029 static char *kwargs[] = {"iterable", "key", NULL};
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000030 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
32
33 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
34 &it, &keyfunc))
35 return NULL;
36
37 gbo = (groupbyobject *)type->tp_alloc(type, 0);
38 if (gbo == NULL)
39 return NULL;
40 gbo->tgtkey = NULL;
41 gbo->currkey = NULL;
42 gbo->currvalue = NULL;
43 gbo->keyfunc = keyfunc;
44 Py_INCREF(keyfunc);
45 gbo->it = PyObject_GetIter(it);
46 if (gbo->it == NULL) {
47 Py_DECREF(gbo);
48 return NULL;
49 }
50 return (PyObject *)gbo;
51}
52
53static void
54groupby_dealloc(groupbyobject *gbo)
55{
56 PyObject_GC_UnTrack(gbo);
57 Py_XDECREF(gbo->it);
58 Py_XDECREF(gbo->keyfunc);
59 Py_XDECREF(gbo->tgtkey);
60 Py_XDECREF(gbo->currkey);
61 Py_XDECREF(gbo->currvalue);
Christian Heimese93237d2007-12-19 02:37:44 +000062 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000063}
64
65static int
66groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
67{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +000068 Py_VISIT(gbo->it);
69 Py_VISIT(gbo->keyfunc);
70 Py_VISIT(gbo->tgtkey);
71 Py_VISIT(gbo->currkey);
72 Py_VISIT(gbo->currvalue);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000073 return 0;
74}
75
76static PyObject *
77groupby_next(groupbyobject *gbo)
78{
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
81 /* skip to next iteration group */
82 for (;;) {
83 if (gbo->currkey == NULL)
84 /* pass */;
85 else if (gbo->tgtkey == NULL)
86 break;
87 else {
88 int rcmp;
89
90 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
91 gbo->currkey, Py_EQ);
92 if (rcmp == -1)
93 return NULL;
94 else if (rcmp == 0)
95 break;
96 }
97
98 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
101
102 if (gbo->keyfunc == Py_None) {
103 newkey = newvalue;
104 Py_INCREF(newvalue);
105 } else {
106 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
107 newvalue, NULL);
108 if (newkey == NULL) {
109 Py_DECREF(newvalue);
110 return NULL;
111 }
112 }
113
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000114 tmp = gbo->currkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000115 gbo->currkey = newkey;
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000116 Py_XDECREF(tmp);
117
118 tmp = gbo->currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000119 gbo->currvalue = newvalue;
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000120 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000121 }
122
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000123 Py_INCREF(gbo->currkey);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000124 tmp = gbo->tgtkey;
125 gbo->tgtkey = gbo->currkey;
126 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000127
128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
131
132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
135}
136
137PyDoc_STRVAR(groupby_doc,
138"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
139(key, sub-iterator) grouped by each value of key(value).\n");
140
141static PyTypeObject groupby_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000142 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000143 "itertools.groupby", /* tp_name */
144 sizeof(groupbyobject), /* tp_basicsize */
145 0, /* tp_itemsize */
146 /* methods */
147 (destructor)groupby_dealloc, /* tp_dealloc */
148 0, /* tp_print */
149 0, /* tp_getattr */
150 0, /* tp_setattr */
151 0, /* tp_compare */
152 0, /* tp_repr */
153 0, /* tp_as_number */
154 0, /* tp_as_sequence */
155 0, /* tp_as_mapping */
156 0, /* tp_hash */
157 0, /* tp_call */
158 0, /* tp_str */
159 PyObject_GenericGetAttr, /* tp_getattro */
160 0, /* tp_setattro */
161 0, /* tp_as_buffer */
162 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
163 Py_TPFLAGS_BASETYPE, /* tp_flags */
164 groupby_doc, /* tp_doc */
165 (traverseproc)groupby_traverse, /* tp_traverse */
166 0, /* tp_clear */
167 0, /* tp_richcompare */
168 0, /* tp_weaklistoffset */
169 PyObject_SelfIter, /* tp_iter */
170 (iternextfunc)groupby_next, /* tp_iternext */
171 0, /* tp_methods */
172 0, /* tp_members */
173 0, /* tp_getset */
174 0, /* tp_base */
175 0, /* tp_dict */
176 0, /* tp_descr_get */
177 0, /* tp_descr_set */
178 0, /* tp_dictoffset */
179 0, /* tp_init */
180 0, /* tp_alloc */
181 groupby_new, /* tp_new */
182 PyObject_GC_Del, /* tp_free */
183};
184
185
186/* _grouper object (internal) ************************************************/
187
188typedef struct {
189 PyObject_HEAD
190 PyObject *parent;
191 PyObject *tgtkey;
192} _grouperobject;
193
194static PyTypeObject _grouper_type;
195
196static PyObject *
197_grouper_create(groupbyobject *parent, PyObject *tgtkey)
198{
199 _grouperobject *igo;
200
201 igo = PyObject_New(_grouperobject, &_grouper_type);
202 if (igo == NULL)
203 return NULL;
204 igo->parent = (PyObject *)parent;
205 Py_INCREF(parent);
206 igo->tgtkey = tgtkey;
207 Py_INCREF(tgtkey);
208
209 return (PyObject *)igo;
210}
211
212static void
213_grouper_dealloc(_grouperobject *igo)
214{
215 Py_DECREF(igo->parent);
216 Py_DECREF(igo->tgtkey);
217 PyObject_Del(igo);
218}
219
220static PyObject *
221_grouper_next(_grouperobject *igo)
222{
223 groupbyobject *gbo = (groupbyobject *)igo->parent;
224 PyObject *newvalue, *newkey, *r;
225 int rcmp;
226
227 if (gbo->currvalue == NULL) {
228 newvalue = PyIter_Next(gbo->it);
229 if (newvalue == NULL)
230 return NULL;
231
232 if (gbo->keyfunc == Py_None) {
233 newkey = newvalue;
234 Py_INCREF(newvalue);
235 } else {
236 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
237 newvalue, NULL);
238 if (newkey == NULL) {
239 Py_DECREF(newvalue);
240 return NULL;
241 }
242 }
243
244 assert(gbo->currkey == NULL);
245 gbo->currkey = newkey;
246 gbo->currvalue = newvalue;
247 }
248
249 assert(gbo->currkey != NULL);
250 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
251 if (rcmp <= 0)
252 /* got any error or current group is end */
253 return NULL;
254
255 r = gbo->currvalue;
256 gbo->currvalue = NULL;
Raymond Hettinger75ccea32004-09-01 07:02:44 +0000257 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
259 return r;
260}
261
262static PyTypeObject _grouper_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000263 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000264 "itertools._grouper", /* tp_name */
265 sizeof(_grouperobject), /* tp_basicsize */
266 0, /* tp_itemsize */
267 /* methods */
268 (destructor)_grouper_dealloc, /* tp_dealloc */
269 0, /* tp_print */
270 0, /* tp_getattr */
271 0, /* tp_setattr */
272 0, /* tp_compare */
273 0, /* tp_repr */
274 0, /* tp_as_number */
275 0, /* tp_as_sequence */
276 0, /* tp_as_mapping */
277 0, /* tp_hash */
278 0, /* tp_call */
279 0, /* tp_str */
280 PyObject_GenericGetAttr, /* tp_getattro */
281 0, /* tp_setattro */
282 0, /* tp_as_buffer */
283 Py_TPFLAGS_DEFAULT, /* tp_flags */
284 0, /* tp_doc */
285 0, /* tp_traverse */
286 0, /* tp_clear */
287 0, /* tp_richcompare */
288 0, /* tp_weaklistoffset */
289 PyObject_SelfIter, /* tp_iter */
290 (iternextfunc)_grouper_next, /* tp_iternext */
291 0, /* tp_methods */
292 0, /* tp_members */
293 0, /* tp_getset */
294 0, /* tp_base */
295 0, /* tp_dict */
296 0, /* tp_descr_get */
297 0, /* tp_descr_set */
298 0, /* tp_dictoffset */
299 0, /* tp_init */
300 0, /* tp_alloc */
301 0, /* tp_new */
302 PyObject_Del, /* tp_free */
303};
304
305
306
Raymond Hettingerad983e72003-11-12 14:32:26 +0000307/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000308
Raymond Hettingerad983e72003-11-12 14:32:26 +0000309/* The teedataobject pre-allocates space for LINKCELLS number of objects.
310 To help the object fit neatly inside cache lines (space for 16 to 32
311 pointers), the value should be a multiple of 16 minus space for
312 the other structure members including PyHEAD overhead. The larger the
313 value, the less memory overhead per object and the less time spent
314 allocating/deallocating new links. The smaller the number, the less
315 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000316*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000317#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000318
319typedef struct {
320 PyObject_HEAD
321 PyObject *it;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000322 int numread;
323 PyObject *nextlink;
324 PyObject *(values[LINKCELLS]);
325} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000326
327typedef struct {
328 PyObject_HEAD
Raymond Hettingerad983e72003-11-12 14:32:26 +0000329 teedataobject *dataobj;
330 int index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000331 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000332} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000333
Raymond Hettingerad983e72003-11-12 14:32:26 +0000334static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000335
336static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000337teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000338{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000339 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000340
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000341 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000342 if (tdo == NULL)
Raymond Hettinger45143692003-10-25 06:37:47 +0000343 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000344
345 tdo->numread = 0;
346 tdo->nextlink = NULL;
347 Py_INCREF(it);
348 tdo->it = it;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000349 PyObject_GC_Track(tdo);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000350 return (PyObject *)tdo;
351}
352
353static PyObject *
354teedataobject_jumplink(teedataobject *tdo)
355{
356 if (tdo->nextlink == NULL)
357 tdo->nextlink = teedataobject_new(tdo->it);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000358 Py_XINCREF(tdo->nextlink);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000359 return tdo->nextlink;
360}
361
362static PyObject *
363teedataobject_getitem(teedataobject *tdo, int i)
364{
365 PyObject *value;
366
367 assert(i < LINKCELLS);
368 if (i < tdo->numread)
369 value = tdo->values[i];
370 else {
371 /* this is the lead iterator, so fetch more data */
372 assert(i == tdo->numread);
373 value = PyIter_Next(tdo->it);
374 if (value == NULL)
375 return NULL;
376 tdo->numread++;
377 tdo->values[i] = value;
Raymond Hettinger45143692003-10-25 06:37:47 +0000378 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000379 Py_INCREF(value);
380 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000381}
382
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000383static int
384teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
385{
386 int i;
387 Py_VISIT(tdo->it);
388 for (i = 0; i < tdo->numread; i++)
389 Py_VISIT(tdo->values[i]);
390 Py_VISIT(tdo->nextlink);
391 return 0;
392}
393
394static int
395teedataobject_clear(teedataobject *tdo)
396{
397 int i;
398 Py_CLEAR(tdo->it);
399 for (i=0 ; i<tdo->numread ; i++)
400 Py_CLEAR(tdo->values[i]);
401 Py_CLEAR(tdo->nextlink);
402 return 0;
403}
404
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000405static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000406teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000407{
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000408 PyObject_GC_UnTrack(tdo);
409 teedataobject_clear(tdo);
410 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000411}
412
Raymond Hettingerad983e72003-11-12 14:32:26 +0000413PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000414
Raymond Hettingerad983e72003-11-12 14:32:26 +0000415static PyTypeObject teedataobject_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000416 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000417 "itertools.tee_dataobject", /* tp_name */
418 sizeof(teedataobject), /* tp_basicsize */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000419 0, /* tp_itemsize */
420 /* methods */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000421 (destructor)teedataobject_dealloc, /* tp_dealloc */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000422 0, /* tp_print */
423 0, /* tp_getattr */
424 0, /* tp_setattr */
425 0, /* tp_compare */
426 0, /* tp_repr */
427 0, /* tp_as_number */
428 0, /* tp_as_sequence */
429 0, /* tp_as_mapping */
430 0, /* tp_hash */
431 0, /* tp_call */
432 0, /* tp_str */
433 PyObject_GenericGetAttr, /* tp_getattro */
434 0, /* tp_setattro */
435 0, /* tp_as_buffer */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000436 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000437 teedataobject_doc, /* tp_doc */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000438 (traverseproc)teedataobject_traverse, /* tp_traverse */
439 (inquiry)teedataobject_clear, /* tp_clear */
440 0, /* tp_richcompare */
441 0, /* tp_weaklistoffset */
442 0, /* tp_iter */
443 0, /* tp_iternext */
444 0, /* tp_methods */
445 0, /* tp_members */
446 0, /* tp_getset */
447 0, /* tp_base */
448 0, /* tp_dict */
449 0, /* tp_descr_get */
450 0, /* tp_descr_set */
451 0, /* tp_dictoffset */
452 0, /* tp_init */
453 0, /* tp_alloc */
454 0, /* tp_new */
455 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000456};
457
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000458
459static PyTypeObject tee_type;
460
461static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000462tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000463{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000464 PyObject *value, *link;
465
466 if (to->index >= LINKCELLS) {
467 link = teedataobject_jumplink(to->dataobj);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000468 Py_DECREF(to->dataobj);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000469 to->dataobj = (teedataobject *)link;
470 to->index = 0;
471 }
472 value = teedataobject_getitem(to->dataobj, to->index);
473 if (value == NULL)
474 return NULL;
475 to->index++;
476 return value;
477}
478
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000479static int
480tee_traverse(teeobject *to, visitproc visit, void *arg)
481{
482 Py_VISIT((PyObject *)to->dataobj);
483 return 0;
484}
485
Raymond Hettingerad983e72003-11-12 14:32:26 +0000486static PyObject *
487tee_copy(teeobject *to)
488{
489 teeobject *newto;
490
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000491 newto = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000492 if (newto == NULL)
493 return NULL;
494 Py_INCREF(to->dataobj);
495 newto->dataobj = to->dataobj;
496 newto->index = to->index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000497 newto->weakreflist = NULL;
Thomas Woutersb3deb942006-04-15 22:33:13 +0000498 PyObject_GC_Track(newto);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000499 return (PyObject *)newto;
500}
501
502PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
503
504static PyObject *
505tee_fromiterable(PyObject *iterable)
506{
507 teeobject *to;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000508 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000509
510 it = PyObject_GetIter(iterable);
511 if (it == NULL)
512 return NULL;
513 if (PyObject_TypeCheck(it, &tee_type)) {
514 to = (teeobject *)tee_copy((teeobject *)it);
515 goto done;
516 }
517
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000518 to = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000519 if (to == NULL)
520 goto done;
521 to->dataobj = (teedataobject *)teedataobject_new(it);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000522 if (!to->dataobj) {
523 PyObject_GC_Del(to);
524 to = NULL;
525 goto done;
526 }
527
Raymond Hettingerad983e72003-11-12 14:32:26 +0000528 to->index = 0;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000529 to->weakreflist = NULL;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000530 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000531done:
532 Py_XDECREF(it);
533 return (PyObject *)to;
534}
535
536static PyObject *
537tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
538{
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000539 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000540
541 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
542 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000543 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000544}
545
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000546static int
547tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000548{
Raymond Hettingera9f60922004-10-17 16:40:14 +0000549 if (to->weakreflist != NULL)
550 PyObject_ClearWeakRefs((PyObject *) to);
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000551 Py_CLEAR(to->dataobj);
552 return 0;
553}
554
555static void
556tee_dealloc(teeobject *to)
557{
558 PyObject_GC_UnTrack(to);
559 tee_clear(to);
560 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000561}
562
Raymond Hettingerad983e72003-11-12 14:32:26 +0000563PyDoc_STRVAR(teeobject_doc,
564"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000565
Raymond Hettingerad983e72003-11-12 14:32:26 +0000566static PyMethodDef tee_methods[] = {
567 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
568 {NULL, NULL} /* sentinel */
569};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000570
571static PyTypeObject tee_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000572 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000573 "itertools.tee", /* tp_name */
574 sizeof(teeobject), /* tp_basicsize */
575 0, /* tp_itemsize */
576 /* methods */
577 (destructor)tee_dealloc, /* tp_dealloc */
578 0, /* tp_print */
579 0, /* tp_getattr */
580 0, /* tp_setattr */
581 0, /* tp_compare */
582 0, /* tp_repr */
583 0, /* tp_as_number */
584 0, /* tp_as_sequence */
585 0, /* tp_as_mapping */
586 0, /* tp_hash */
587 0, /* tp_call */
588 0, /* tp_str */
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000589 0, /* tp_getattro */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000590 0, /* tp_setattro */
591 0, /* tp_as_buffer */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000592 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000593 teeobject_doc, /* tp_doc */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000594 (traverseproc)tee_traverse, /* tp_traverse */
595 (inquiry)tee_clear, /* tp_clear */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000596 0, /* tp_richcompare */
Raymond Hettingera9f60922004-10-17 16:40:14 +0000597 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000598 PyObject_SelfIter, /* tp_iter */
599 (iternextfunc)tee_next, /* tp_iternext */
600 tee_methods, /* tp_methods */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000601 0, /* tp_members */
602 0, /* tp_getset */
603 0, /* tp_base */
604 0, /* tp_dict */
605 0, /* tp_descr_get */
606 0, /* tp_descr_set */
607 0, /* tp_dictoffset */
608 0, /* tp_init */
609 0, /* tp_alloc */
610 tee_new, /* tp_new */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000611 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000612};
613
Raymond Hettingerad983e72003-11-12 14:32:26 +0000614static PyObject *
615tee(PyObject *self, PyObject *args)
616{
Neal Norwitz69e88972006-09-02 02:58:13 +0000617 Py_ssize_t i, n=2;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000618 PyObject *it, *iterable, *copyable, *result;
619
Neal Norwitz69e88972006-09-02 02:58:13 +0000620 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000621 return NULL;
Neal Norwitz69e88972006-09-02 02:58:13 +0000622 if (n < 0) {
623 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
624 return NULL;
625 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000626 result = PyTuple_New(n);
627 if (result == NULL)
628 return NULL;
629 if (n == 0)
630 return result;
631 it = PyObject_GetIter(iterable);
632 if (it == NULL) {
633 Py_DECREF(result);
634 return NULL;
635 }
636 if (!PyObject_HasAttrString(it, "__copy__")) {
637 copyable = tee_fromiterable(it);
638 Py_DECREF(it);
639 if (copyable == NULL) {
640 Py_DECREF(result);
641 return NULL;
642 }
643 } else
644 copyable = it;
645 PyTuple_SET_ITEM(result, 0, copyable);
646 for (i=1 ; i<n ; i++) {
647 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
648 if (copyable == NULL) {
649 Py_DECREF(result);
650 return NULL;
651 }
652 PyTuple_SET_ITEM(result, i, copyable);
653 }
654 return result;
655}
656
657PyDoc_STRVAR(tee_doc,
658"tee(iterable, n=2) --> tuple of n independent iterators.");
659
660
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000661/* cycle object **********************************************************/
662
663typedef struct {
664 PyObject_HEAD
665 PyObject *it;
666 PyObject *saved;
667 int firstpass;
668} cycleobject;
669
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000670static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000671
672static PyObject *
673cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
674{
675 PyObject *it;
676 PyObject *iterable;
677 PyObject *saved;
678 cycleobject *lz;
679
Georg Brandlb84c1372007-01-21 10:28:43 +0000680 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000681 return NULL;
682
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000683 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
684 return NULL;
685
686 /* Get iterator. */
687 it = PyObject_GetIter(iterable);
688 if (it == NULL)
689 return NULL;
690
691 saved = PyList_New(0);
692 if (saved == NULL) {
693 Py_DECREF(it);
694 return NULL;
695 }
696
697 /* create cycleobject structure */
698 lz = (cycleobject *)type->tp_alloc(type, 0);
699 if (lz == NULL) {
700 Py_DECREF(it);
701 Py_DECREF(saved);
702 return NULL;
703 }
704 lz->it = it;
705 lz->saved = saved;
706 lz->firstpass = 0;
707
708 return (PyObject *)lz;
709}
710
711static void
712cycle_dealloc(cycleobject *lz)
713{
714 PyObject_GC_UnTrack(lz);
715 Py_XDECREF(lz->saved);
716 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +0000717 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000718}
719
720static int
721cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
722{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000723 Py_VISIT(lz->it);
724 Py_VISIT(lz->saved);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000725 return 0;
726}
727
728static PyObject *
729cycle_next(cycleobject *lz)
730{
731 PyObject *item;
732 PyObject *it;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000733 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000734
735 while (1) {
736 item = PyIter_Next(lz->it);
737 if (item != NULL) {
738 if (!lz->firstpass)
739 PyList_Append(lz->saved, item);
740 return item;
741 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000742 if (PyErr_Occurred()) {
743 if (PyErr_ExceptionMatches(PyExc_StopIteration))
744 PyErr_Clear();
745 else
746 return NULL;
747 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000748 if (PyList_Size(lz->saved) == 0)
749 return NULL;
750 it = PyObject_GetIter(lz->saved);
751 if (it == NULL)
752 return NULL;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000753 tmp = lz->it;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000754 lz->it = it;
755 lz->firstpass = 1;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000756 Py_DECREF(tmp);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000757 }
758}
759
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000760PyDoc_STRVAR(cycle_doc,
761"cycle(iterable) --> cycle object\n\
762\n\
763Return elements from the iterable until it is exhausted.\n\
764Then repeat the sequence indefinitely.");
765
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000766static PyTypeObject cycle_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000767 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000768 "itertools.cycle", /* tp_name */
769 sizeof(cycleobject), /* tp_basicsize */
770 0, /* tp_itemsize */
771 /* methods */
772 (destructor)cycle_dealloc, /* tp_dealloc */
773 0, /* tp_print */
774 0, /* tp_getattr */
775 0, /* tp_setattr */
776 0, /* tp_compare */
777 0, /* tp_repr */
778 0, /* tp_as_number */
779 0, /* tp_as_sequence */
780 0, /* tp_as_mapping */
781 0, /* tp_hash */
782 0, /* tp_call */
783 0, /* tp_str */
784 PyObject_GenericGetAttr, /* tp_getattro */
785 0, /* tp_setattro */
786 0, /* tp_as_buffer */
787 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
788 Py_TPFLAGS_BASETYPE, /* tp_flags */
789 cycle_doc, /* tp_doc */
790 (traverseproc)cycle_traverse, /* tp_traverse */
791 0, /* tp_clear */
792 0, /* tp_richcompare */
793 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000794 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000795 (iternextfunc)cycle_next, /* tp_iternext */
796 0, /* tp_methods */
797 0, /* tp_members */
798 0, /* tp_getset */
799 0, /* tp_base */
800 0, /* tp_dict */
801 0, /* tp_descr_get */
802 0, /* tp_descr_set */
803 0, /* tp_dictoffset */
804 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000805 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000806 cycle_new, /* tp_new */
807 PyObject_GC_Del, /* tp_free */
808};
809
810
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000811/* dropwhile object **********************************************************/
812
813typedef struct {
814 PyObject_HEAD
815 PyObject *func;
816 PyObject *it;
817 long start;
818} dropwhileobject;
819
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000820static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000821
822static PyObject *
823dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
824{
825 PyObject *func, *seq;
826 PyObject *it;
827 dropwhileobject *lz;
828
Georg Brandlb84c1372007-01-21 10:28:43 +0000829 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000830 return NULL;
831
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000832 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
833 return NULL;
834
835 /* Get iterator. */
836 it = PyObject_GetIter(seq);
837 if (it == NULL)
838 return NULL;
839
840 /* create dropwhileobject structure */
841 lz = (dropwhileobject *)type->tp_alloc(type, 0);
842 if (lz == NULL) {
843 Py_DECREF(it);
844 return NULL;
845 }
846 Py_INCREF(func);
847 lz->func = func;
848 lz->it = it;
849 lz->start = 0;
850
851 return (PyObject *)lz;
852}
853
854static void
855dropwhile_dealloc(dropwhileobject *lz)
856{
857 PyObject_GC_UnTrack(lz);
858 Py_XDECREF(lz->func);
859 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +0000860 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000861}
862
863static int
864dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
865{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000866 Py_VISIT(lz->it);
867 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000868 return 0;
869}
870
871static PyObject *
872dropwhile_next(dropwhileobject *lz)
873{
874 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000875 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000876 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000877 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000878
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000879 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +0000880 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000881 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000882 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000883 if (item == NULL)
884 return NULL;
885 if (lz->start == 1)
886 return item;
887
888 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
889 if (good == NULL) {
890 Py_DECREF(item);
891 return NULL;
892 }
893 ok = PyObject_IsTrue(good);
894 Py_DECREF(good);
895 if (!ok) {
896 lz->start = 1;
897 return item;
898 }
899 Py_DECREF(item);
900 }
901}
902
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000903PyDoc_STRVAR(dropwhile_doc,
904"dropwhile(predicate, iterable) --> dropwhile object\n\
905\n\
906Drop items from the iterable while predicate(item) is true.\n\
907Afterwards, return every element until the iterable is exhausted.");
908
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000909static PyTypeObject dropwhile_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000910 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000911 "itertools.dropwhile", /* tp_name */
912 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000913 0, /* tp_itemsize */
914 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000915 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000916 0, /* tp_print */
917 0, /* tp_getattr */
918 0, /* tp_setattr */
919 0, /* tp_compare */
920 0, /* tp_repr */
921 0, /* tp_as_number */
922 0, /* tp_as_sequence */
923 0, /* tp_as_mapping */
924 0, /* tp_hash */
925 0, /* tp_call */
926 0, /* tp_str */
927 PyObject_GenericGetAttr, /* tp_getattro */
928 0, /* tp_setattro */
929 0, /* tp_as_buffer */
930 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
931 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000932 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000933 (traverseproc)dropwhile_traverse, /* tp_traverse */
934 0, /* tp_clear */
935 0, /* tp_richcompare */
936 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000937 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000938 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000939 0, /* tp_methods */
940 0, /* tp_members */
941 0, /* tp_getset */
942 0, /* tp_base */
943 0, /* tp_dict */
944 0, /* tp_descr_get */
945 0, /* tp_descr_set */
946 0, /* tp_dictoffset */
947 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000948 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000949 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000950 PyObject_GC_Del, /* tp_free */
951};
952
953
954/* takewhile object **********************************************************/
955
956typedef struct {
957 PyObject_HEAD
958 PyObject *func;
959 PyObject *it;
960 long stop;
961} takewhileobject;
962
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000963static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000964
965static PyObject *
966takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
967{
968 PyObject *func, *seq;
969 PyObject *it;
970 takewhileobject *lz;
971
Georg Brandlb84c1372007-01-21 10:28:43 +0000972 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000973 return NULL;
974
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000975 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
976 return NULL;
977
978 /* Get iterator. */
979 it = PyObject_GetIter(seq);
980 if (it == NULL)
981 return NULL;
982
983 /* create takewhileobject structure */
984 lz = (takewhileobject *)type->tp_alloc(type, 0);
985 if (lz == NULL) {
986 Py_DECREF(it);
987 return NULL;
988 }
989 Py_INCREF(func);
990 lz->func = func;
991 lz->it = it;
992 lz->stop = 0;
993
994 return (PyObject *)lz;
995}
996
997static void
998takewhile_dealloc(takewhileobject *lz)
999{
1000 PyObject_GC_UnTrack(lz);
1001 Py_XDECREF(lz->func);
1002 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00001003 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001004}
1005
1006static int
1007takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1008{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001009 Py_VISIT(lz->it);
1010 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001011 return 0;
1012}
1013
1014static PyObject *
1015takewhile_next(takewhileobject *lz)
1016{
1017 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001018 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001019 long ok;
1020
1021 if (lz->stop == 1)
1022 return NULL;
1023
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001024 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00001025 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001026 if (item == NULL)
1027 return NULL;
1028
1029 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1030 if (good == NULL) {
1031 Py_DECREF(item);
1032 return NULL;
1033 }
1034 ok = PyObject_IsTrue(good);
1035 Py_DECREF(good);
1036 if (ok)
1037 return item;
1038 Py_DECREF(item);
1039 lz->stop = 1;
1040 return NULL;
1041}
1042
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001043PyDoc_STRVAR(takewhile_doc,
1044"takewhile(predicate, iterable) --> takewhile object\n\
1045\n\
1046Return successive entries from an iterable as long as the \n\
1047predicate evaluates to true for each entry.");
1048
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001049static PyTypeObject takewhile_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001050 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001051 "itertools.takewhile", /* tp_name */
1052 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053 0, /* tp_itemsize */
1054 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001055 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001056 0, /* tp_print */
1057 0, /* tp_getattr */
1058 0, /* tp_setattr */
1059 0, /* tp_compare */
1060 0, /* tp_repr */
1061 0, /* tp_as_number */
1062 0, /* tp_as_sequence */
1063 0, /* tp_as_mapping */
1064 0, /* tp_hash */
1065 0, /* tp_call */
1066 0, /* tp_str */
1067 PyObject_GenericGetAttr, /* tp_getattro */
1068 0, /* tp_setattro */
1069 0, /* tp_as_buffer */
1070 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1071 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001072 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001073 (traverseproc)takewhile_traverse, /* tp_traverse */
1074 0, /* tp_clear */
1075 0, /* tp_richcompare */
1076 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001077 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001078 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001079 0, /* tp_methods */
1080 0, /* tp_members */
1081 0, /* tp_getset */
1082 0, /* tp_base */
1083 0, /* tp_dict */
1084 0, /* tp_descr_get */
1085 0, /* tp_descr_set */
1086 0, /* tp_dictoffset */
1087 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001088 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001089 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001090 PyObject_GC_Del, /* tp_free */
1091};
1092
1093
1094/* islice object ************************************************************/
1095
1096typedef struct {
1097 PyObject_HEAD
1098 PyObject *it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001099 Py_ssize_t next;
1100 Py_ssize_t stop;
1101 Py_ssize_t step;
1102 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001103} isliceobject;
1104
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001105static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001106
1107static PyObject *
1108islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1109{
1110 PyObject *seq;
Jack Diederich6c433a92006-05-26 11:15:17 +00001111 Py_ssize_t start=0, stop=-1, step=1;
Raymond Hettingerb2594052004-12-05 09:25:51 +00001112 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001113 Py_ssize_t numargs;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001114 isliceobject *lz;
1115
Georg Brandlb84c1372007-01-21 10:28:43 +00001116 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001117 return NULL;
1118
Raymond Hettingerb2594052004-12-05 09:25:51 +00001119 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001120 return NULL;
1121
Raymond Hettingerb2594052004-12-05 09:25:51 +00001122 numargs = PyTuple_Size(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001123 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001124 if (a1 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001125 stop = PyInt_AsSsize_t(a1);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001126 if (stop == -1) {
1127 if (PyErr_Occurred())
1128 PyErr_Clear();
1129 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001130 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001131 return NULL;
1132 }
1133 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001134 } else {
Raymond Hettingerb2594052004-12-05 09:25:51 +00001135 if (a1 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001136 start = PyInt_AsSsize_t(a1);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001137 if (start == -1 && PyErr_Occurred())
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001138 PyErr_Clear();
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001139 if (a2 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001140 stop = PyInt_AsSsize_t(a2);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001141 if (stop == -1) {
1142 if (PyErr_Occurred())
1143 PyErr_Clear();
1144 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001145 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001146 return NULL;
1147 }
1148 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001149 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001150 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001151 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001152 "Indices for islice() must be non-negative integers or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001153 return NULL;
1154 }
1155
Raymond Hettingerb2594052004-12-05 09:25:51 +00001156 if (a3 != NULL) {
1157 if (a3 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001158 step = PyInt_AsSsize_t(a3);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001159 if (step == -1 && PyErr_Occurred())
1160 PyErr_Clear();
1161 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001162 if (step<1) {
1163 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001164 "Step for islice() must be a positive integer or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001165 return NULL;
1166 }
1167
1168 /* Get iterator. */
1169 it = PyObject_GetIter(seq);
1170 if (it == NULL)
1171 return NULL;
1172
1173 /* create isliceobject structure */
1174 lz = (isliceobject *)type->tp_alloc(type, 0);
1175 if (lz == NULL) {
1176 Py_DECREF(it);
1177 return NULL;
1178 }
1179 lz->it = it;
1180 lz->next = start;
1181 lz->stop = stop;
1182 lz->step = step;
1183 lz->cnt = 0L;
1184
1185 return (PyObject *)lz;
1186}
1187
1188static void
1189islice_dealloc(isliceobject *lz)
1190{
1191 PyObject_GC_UnTrack(lz);
1192 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00001193 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001194}
1195
1196static int
1197islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1198{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001199 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001200 return 0;
1201}
1202
1203static PyObject *
1204islice_next(isliceobject *lz)
1205{
1206 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001207 PyObject *it = lz->it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001208 Py_ssize_t oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001209 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001210
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001211 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00001212 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001213 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001214 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001215 if (item == NULL)
1216 return NULL;
1217 Py_DECREF(item);
1218 lz->cnt++;
1219 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001220 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001221 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001222 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001223 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001224 if (item == NULL)
1225 return NULL;
1226 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001227 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001228 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001229 if (lz->next < oldnext) /* Check for overflow */
1230 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001231 return item;
1232}
1233
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001234PyDoc_STRVAR(islice_doc,
1235"islice(iterable, [start,] stop [, step]) --> islice object\n\
1236\n\
1237Return an iterator whose next() method returns selected values from an\n\
1238iterable. If start is specified, will skip all preceding elements;\n\
1239otherwise, start defaults to zero. Step defaults to one. If\n\
1240specified as another value, step determines how many values are \n\
1241skipped between successive calls. Works like a slice() on a list\n\
1242but returns an iterator.");
1243
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001244static PyTypeObject islice_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001245 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001246 "itertools.islice", /* tp_name */
1247 sizeof(isliceobject), /* tp_basicsize */
1248 0, /* tp_itemsize */
1249 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001250 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001251 0, /* tp_print */
1252 0, /* tp_getattr */
1253 0, /* tp_setattr */
1254 0, /* tp_compare */
1255 0, /* tp_repr */
1256 0, /* tp_as_number */
1257 0, /* tp_as_sequence */
1258 0, /* tp_as_mapping */
1259 0, /* tp_hash */
1260 0, /* tp_call */
1261 0, /* tp_str */
1262 PyObject_GenericGetAttr, /* tp_getattro */
1263 0, /* tp_setattro */
1264 0, /* tp_as_buffer */
1265 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1266 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001267 islice_doc, /* tp_doc */
1268 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001269 0, /* tp_clear */
1270 0, /* tp_richcompare */
1271 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001272 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001273 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001274 0, /* tp_methods */
1275 0, /* tp_members */
1276 0, /* tp_getset */
1277 0, /* tp_base */
1278 0, /* tp_dict */
1279 0, /* tp_descr_get */
1280 0, /* tp_descr_set */
1281 0, /* tp_dictoffset */
1282 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001283 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001284 islice_new, /* tp_new */
1285 PyObject_GC_Del, /* tp_free */
1286};
1287
1288
1289/* starmap object ************************************************************/
1290
1291typedef struct {
1292 PyObject_HEAD
1293 PyObject *func;
1294 PyObject *it;
1295} starmapobject;
1296
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001297static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001298
1299static PyObject *
1300starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1301{
1302 PyObject *func, *seq;
1303 PyObject *it;
1304 starmapobject *lz;
1305
Georg Brandlb84c1372007-01-21 10:28:43 +00001306 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001307 return NULL;
1308
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001309 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1310 return NULL;
1311
1312 /* Get iterator. */
1313 it = PyObject_GetIter(seq);
1314 if (it == NULL)
1315 return NULL;
1316
1317 /* create starmapobject structure */
1318 lz = (starmapobject *)type->tp_alloc(type, 0);
1319 if (lz == NULL) {
1320 Py_DECREF(it);
1321 return NULL;
1322 }
1323 Py_INCREF(func);
1324 lz->func = func;
1325 lz->it = it;
1326
1327 return (PyObject *)lz;
1328}
1329
1330static void
1331starmap_dealloc(starmapobject *lz)
1332{
1333 PyObject_GC_UnTrack(lz);
1334 Py_XDECREF(lz->func);
1335 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00001336 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001337}
1338
1339static int
1340starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1341{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001342 Py_VISIT(lz->it);
1343 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001344 return 0;
1345}
1346
1347static PyObject *
1348starmap_next(starmapobject *lz)
1349{
1350 PyObject *args;
1351 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001352 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001353
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001354 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00001355 args = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001356 if (args == NULL)
1357 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001358 if (!PyTuple_CheckExact(args)) {
Raymond Hettinger47317092008-01-17 03:02:14 +00001359 PyObject *newargs = PySequence_Tuple(args);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001360 Py_DECREF(args);
Raymond Hettinger47317092008-01-17 03:02:14 +00001361 if (newargs == NULL)
1362 return NULL;
1363 args = newargs;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001364 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001365 result = PyObject_Call(lz->func, args, NULL);
1366 Py_DECREF(args);
1367 return result;
1368}
1369
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001370PyDoc_STRVAR(starmap_doc,
1371"starmap(function, sequence) --> starmap object\n\
1372\n\
1373Return an iterator whose values are returned from the function evaluated\n\
1374with a argument tuple taken from the given sequence.");
1375
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001376static PyTypeObject starmap_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001377 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001378 "itertools.starmap", /* tp_name */
1379 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001380 0, /* tp_itemsize */
1381 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001382 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001383 0, /* tp_print */
1384 0, /* tp_getattr */
1385 0, /* tp_setattr */
1386 0, /* tp_compare */
1387 0, /* tp_repr */
1388 0, /* tp_as_number */
1389 0, /* tp_as_sequence */
1390 0, /* tp_as_mapping */
1391 0, /* tp_hash */
1392 0, /* tp_call */
1393 0, /* tp_str */
1394 PyObject_GenericGetAttr, /* tp_getattro */
1395 0, /* tp_setattro */
1396 0, /* tp_as_buffer */
1397 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1398 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001399 starmap_doc, /* tp_doc */
1400 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001401 0, /* tp_clear */
1402 0, /* tp_richcompare */
1403 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001404 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001405 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001406 0, /* tp_methods */
1407 0, /* tp_members */
1408 0, /* tp_getset */
1409 0, /* tp_base */
1410 0, /* tp_dict */
1411 0, /* tp_descr_get */
1412 0, /* tp_descr_set */
1413 0, /* tp_dictoffset */
1414 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001415 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001416 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001417 PyObject_GC_Del, /* tp_free */
1418};
1419
1420
1421/* imap object ************************************************************/
1422
1423typedef struct {
1424 PyObject_HEAD
1425 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001426 PyObject *func;
1427} imapobject;
1428
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001429static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001430
1431static PyObject *
1432imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1433{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001434 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001435 imapobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001436 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001437
Georg Brandlb84c1372007-01-21 10:28:43 +00001438 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001439 return NULL;
1440
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001441 numargs = PyTuple_Size(args);
1442 if (numargs < 2) {
1443 PyErr_SetString(PyExc_TypeError,
1444 "imap() must have at least two arguments.");
1445 return NULL;
1446 }
1447
1448 iters = PyTuple_New(numargs-1);
1449 if (iters == NULL)
1450 return NULL;
1451
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001452 for (i=1 ; i<numargs ; i++) {
1453 /* Get iterator. */
1454 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1455 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001456 Py_DECREF(iters);
1457 return NULL;
1458 }
1459 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001460 }
1461
1462 /* create imapobject structure */
1463 lz = (imapobject *)type->tp_alloc(type, 0);
1464 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001465 Py_DECREF(iters);
1466 return NULL;
1467 }
1468 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469 func = PyTuple_GET_ITEM(args, 0);
1470 Py_INCREF(func);
1471 lz->func = func;
1472
1473 return (PyObject *)lz;
1474}
1475
1476static void
1477imap_dealloc(imapobject *lz)
1478{
1479 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001480 Py_XDECREF(lz->iters);
1481 Py_XDECREF(lz->func);
Christian Heimese93237d2007-12-19 02:37:44 +00001482 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001483}
1484
1485static int
1486imap_traverse(imapobject *lz, visitproc visit, void *arg)
1487{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001488 Py_VISIT(lz->iters);
1489 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001490 return 0;
1491}
1492
Raymond Hettinger2012f172003-02-07 05:32:58 +00001493/*
1494imap() is an iterator version of __builtins__.map() except that it does
1495not have the None fill-in feature. That was intentionally left out for
1496the following reasons:
1497
1498 1) Itertools are designed to be easily combined and chained together.
1499 Having all tools stop with the shortest input is a unifying principle
1500 that makes it easier to combine finite iterators (supplying data) with
1501 infinite iterators like count() and repeat() (for supplying sequential
1502 or constant arguments to a function).
1503
1504 2) In typical use cases for combining itertools, having one finite data
1505 supplier run out before another is likely to be an error condition which
1506 should not pass silently by automatically supplying None.
1507
1508 3) The use cases for automatic None fill-in are rare -- not many functions
1509 do something useful when a parameter suddenly switches type and becomes
1510 None.
1511
1512 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001513 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001514
1515 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1516*/
1517
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001518static PyObject *
1519imap_next(imapobject *lz)
1520{
1521 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001522 PyObject *argtuple;
1523 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001524 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001525
1526 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001527 argtuple = PyTuple_New(numargs);
1528 if (argtuple == NULL)
1529 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001530
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001531 for (i=0 ; i<numargs ; i++) {
1532 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1533 if (val == NULL) {
1534 Py_DECREF(argtuple);
1535 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001536 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001537 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001538 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001539 if (lz->func == Py_None)
1540 return argtuple;
1541 result = PyObject_Call(lz->func, argtuple, NULL);
1542 Py_DECREF(argtuple);
1543 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001544}
1545
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001546PyDoc_STRVAR(imap_doc,
1547"imap(func, *iterables) --> imap object\n\
1548\n\
1549Make an iterator that computes the function using arguments from\n\
1550each of the iterables. Like map() except that it returns\n\
1551an iterator instead of a list and that it stops when the shortest\n\
1552iterable is exhausted instead of filling in None for shorter\n\
1553iterables.");
1554
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001555static PyTypeObject imap_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001556 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001557 "itertools.imap", /* tp_name */
1558 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001559 0, /* tp_itemsize */
1560 /* methods */
1561 (destructor)imap_dealloc, /* tp_dealloc */
1562 0, /* tp_print */
1563 0, /* tp_getattr */
1564 0, /* tp_setattr */
1565 0, /* tp_compare */
1566 0, /* tp_repr */
1567 0, /* tp_as_number */
1568 0, /* tp_as_sequence */
1569 0, /* tp_as_mapping */
1570 0, /* tp_hash */
1571 0, /* tp_call */
1572 0, /* tp_str */
1573 PyObject_GenericGetAttr, /* tp_getattro */
1574 0, /* tp_setattro */
1575 0, /* tp_as_buffer */
1576 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1577 Py_TPFLAGS_BASETYPE, /* tp_flags */
1578 imap_doc, /* tp_doc */
1579 (traverseproc)imap_traverse, /* tp_traverse */
1580 0, /* tp_clear */
1581 0, /* tp_richcompare */
1582 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001583 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001584 (iternextfunc)imap_next, /* tp_iternext */
1585 0, /* tp_methods */
1586 0, /* tp_members */
1587 0, /* tp_getset */
1588 0, /* tp_base */
1589 0, /* tp_dict */
1590 0, /* tp_descr_get */
1591 0, /* tp_descr_set */
1592 0, /* tp_dictoffset */
1593 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001594 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001595 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001596 PyObject_GC_Del, /* tp_free */
1597};
1598
1599
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001600/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001601
1602typedef struct {
1603 PyObject_HEAD
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001604 PyObject *source; /* Iterator over input iterables */
1605 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001606} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001607
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001608static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001609
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001610static PyObject *
1611chain_new_internal(PyTypeObject *type, PyObject *source)
1612{
1613 chainobject *lz;
1614
1615 lz = (chainobject *)type->tp_alloc(type, 0);
1616 if (lz == NULL) {
1617 Py_DECREF(source);
1618 return NULL;
1619 }
1620
1621 lz->source = source;
1622 lz->active = NULL;
1623 return (PyObject *)lz;
1624}
1625
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001626static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001627chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001628{
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001629 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001630
Georg Brandlb84c1372007-01-21 10:28:43 +00001631 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001632 return NULL;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001633
1634 source = PyObject_GetIter(args);
1635 if (source == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001636 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001637
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001638 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001639}
1640
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001641static PyObject *
1642chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1643{
1644 PyObject *source;
1645
1646 source = PyObject_GetIter(arg);
1647 if (source == NULL)
1648 return NULL;
1649
1650 return chain_new_internal(type, source);
1651}
1652
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001653static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001654chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001655{
1656 PyObject_GC_UnTrack(lz);
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001657 Py_XDECREF(lz->active);
1658 Py_XDECREF(lz->source);
Christian Heimese93237d2007-12-19 02:37:44 +00001659 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001660}
1661
Raymond Hettinger2012f172003-02-07 05:32:58 +00001662static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001663chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001664{
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001665 Py_VISIT(lz->source);
1666 Py_VISIT(lz->active);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001667 return 0;
1668}
1669
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001670static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001671chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001672{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001673 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001674
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001675 if (lz->source == NULL)
1676 return NULL; /* already stopped */
1677
1678 if (lz->active == NULL) {
1679 PyObject *iterable = PyIter_Next(lz->source);
1680 if (iterable == NULL) {
1681 Py_CLEAR(lz->source);
1682 return NULL; /* no more input sources */
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001683 }
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001684 lz->active = PyObject_GetIter(iterable);
Raymond Hettingerf1cca2b2008-03-04 22:29:44 +00001685 Py_DECREF(iterable);
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001686 if (lz->active == NULL) {
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001687 Py_CLEAR(lz->source);
1688 return NULL; /* input not iterable */
1689 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001690 }
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001691 item = PyIter_Next(lz->active);
1692 if (item != NULL)
1693 return item;
1694 if (PyErr_Occurred()) {
1695 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1696 PyErr_Clear();
1697 else
1698 return NULL; /* input raised an exception */
1699 }
1700 Py_CLEAR(lz->active);
1701 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001702}
1703
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001704PyDoc_STRVAR(chain_doc,
1705"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001706\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001707Return a chain object whose .next() method returns elements from the\n\
1708first iterable until it is exhausted, then elements from the next\n\
1709iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001710
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001711PyDoc_STRVAR(chain_from_iterable_doc,
1712"chain.from_iterable(iterable) --> chain object\n\
1713\n\
1714Alternate chain() contructor taking a single iterable argument\n\
1715that evaluates lazily.");
1716
1717static PyMethodDef chain_methods[] = {
1718 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1719 chain_from_iterable_doc},
1720 {NULL, NULL} /* sentinel */
1721};
1722
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001723static PyTypeObject chain_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001724 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001725 "itertools.chain", /* tp_name */
1726 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001727 0, /* tp_itemsize */
1728 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001729 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001730 0, /* tp_print */
1731 0, /* tp_getattr */
1732 0, /* tp_setattr */
1733 0, /* tp_compare */
1734 0, /* tp_repr */
1735 0, /* tp_as_number */
1736 0, /* tp_as_sequence */
1737 0, /* tp_as_mapping */
1738 0, /* tp_hash */
1739 0, /* tp_call */
1740 0, /* tp_str */
1741 PyObject_GenericGetAttr, /* tp_getattro */
1742 0, /* tp_setattro */
1743 0, /* tp_as_buffer */
1744 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1745 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001746 chain_doc, /* tp_doc */
1747 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001748 0, /* tp_clear */
1749 0, /* tp_richcompare */
1750 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001751 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001752 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001753 chain_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001754 0, /* tp_members */
1755 0, /* tp_getset */
1756 0, /* tp_base */
1757 0, /* tp_dict */
1758 0, /* tp_descr_get */
1759 0, /* tp_descr_set */
1760 0, /* tp_dictoffset */
1761 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001762 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001763 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001764 PyObject_GC_Del, /* tp_free */
1765};
1766
1767
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001768/* product object ************************************************************/
1769
1770typedef struct {
1771 PyObject_HEAD
1772 PyObject *pools; /* tuple of pool tuples */
Raymond Hettinger532316d2008-02-23 04:03:50 +00001773 Py_ssize_t *indices; /* one index per pool */
1774 PyObject *result; /* most recently returned result tuple */
1775 int stopped; /* set to 1 when the product iterator is exhausted */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001776} productobject;
1777
1778static PyTypeObject product_type;
1779
1780static PyObject *
1781product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1782{
1783 productobject *lz;
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001784 Py_ssize_t nargs, npools, repeat=1;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001785 PyObject *pools = NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001786 Py_ssize_t *indices = NULL;
1787 Py_ssize_t i;
1788
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001789 if (kwds != NULL) {
1790 char *kwlist[] = {"repeat", 0};
1791 PyObject *tmpargs = PyTuple_New(0);
1792 if (tmpargs == NULL)
1793 return NULL;
1794 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1795 Py_DECREF(tmpargs);
1796 return NULL;
1797 }
1798 Py_DECREF(tmpargs);
1799 if (repeat < 0) {
1800 PyErr_SetString(PyExc_ValueError,
1801 "repeat argument cannot be negative");
1802 return NULL;
1803 }
1804 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001805
1806 assert(PyTuple_Check(args));
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001807 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1808 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001809
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001810 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
Raymond Hettinger61024b92008-03-02 11:57:16 +00001811 if (indices == NULL) {
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001812 PyErr_NoMemory();
1813 goto error;
1814 }
1815
1816 pools = PyTuple_New(npools);
1817 if (pools == NULL)
1818 goto error;
1819
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001820 for (i=0; i < nargs ; ++i) {
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001821 PyObject *item = PyTuple_GET_ITEM(args, i);
1822 PyObject *pool = PySequence_Tuple(item);
1823 if (pool == NULL)
1824 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001825 PyTuple_SET_ITEM(pools, i, pool);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001826 indices[i] = 0;
1827 }
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001828 for ( ; i < npools; ++i) {
1829 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1830 Py_INCREF(pool);
1831 PyTuple_SET_ITEM(pools, i, pool);
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001832 indices[i] = 0;
1833 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001834
1835 /* create productobject structure */
1836 lz = (productobject *)type->tp_alloc(type, 0);
Raymond Hettinger3bd77122008-02-27 01:08:04 +00001837 if (lz == NULL)
Raymond Hettinger73d79632008-02-23 02:20:41 +00001838 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001839
1840 lz->pools = pools;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001841 lz->indices = indices;
1842 lz->result = NULL;
1843 lz->stopped = 0;
1844
1845 return (PyObject *)lz;
1846
1847error:
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001848 if (indices != NULL)
1849 PyMem_Free(indices);
1850 Py_XDECREF(pools);
1851 return NULL;
1852}
1853
1854static void
1855product_dealloc(productobject *lz)
1856{
1857 PyObject_GC_UnTrack(lz);
1858 Py_XDECREF(lz->pools);
1859 Py_XDECREF(lz->result);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001860 PyMem_Free(lz->indices);
1861 Py_TYPE(lz)->tp_free(lz);
1862}
1863
1864static int
1865product_traverse(productobject *lz, visitproc visit, void *arg)
1866{
1867 Py_VISIT(lz->pools);
1868 Py_VISIT(lz->result);
1869 return 0;
1870}
1871
1872static PyObject *
1873product_next(productobject *lz)
1874{
1875 PyObject *pool;
1876 PyObject *elem;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001877 PyObject *oldelem;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001878 PyObject *pools = lz->pools;
1879 PyObject *result = lz->result;
1880 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1881 Py_ssize_t i;
1882
1883 if (lz->stopped)
1884 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001885
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001886 if (result == NULL) {
Raymond Hettinger73d79632008-02-23 02:20:41 +00001887 /* On the first pass, return an initial tuple filled with the
Raymond Hettingerd553d852008-03-04 04:17:08 +00001888 first element from each pool. */
Raymond Hettinger73d79632008-02-23 02:20:41 +00001889 result = PyTuple_New(npools);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001890 if (result == NULL)
1891 goto empty;
1892 lz->result = result;
1893 for (i=0; i < npools; i++) {
1894 pool = PyTuple_GET_ITEM(pools, i);
1895 if (PyTuple_GET_SIZE(pool) == 0)
1896 goto empty;
1897 elem = PyTuple_GET_ITEM(pool, 0);
1898 Py_INCREF(elem);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001899 PyTuple_SET_ITEM(result, i, elem);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001900 }
1901 } else {
1902 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001903
1904 /* Copy the previous result tuple or re-use it if available */
1905 if (Py_REFCNT(result) > 1) {
1906 PyObject *old_result = result;
1907 result = PyTuple_New(npools);
1908 if (result == NULL)
1909 goto empty;
1910 lz->result = result;
1911 for (i=0; i < npools; i++) {
1912 elem = PyTuple_GET_ITEM(old_result, i);
1913 Py_INCREF(elem);
1914 PyTuple_SET_ITEM(result, i, elem);
1915 }
1916 Py_DECREF(old_result);
1917 }
1918 /* Now, we've got the only copy so we can update it in-place */
Raymond Hettingere3fabd12008-03-02 12:02:19 +00001919 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001920
1921 /* Update the pool indices right-to-left. Only advance to the
1922 next pool when the previous one rolls-over */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001923 for (i=npools-1 ; i >= 0 ; i--) {
1924 pool = PyTuple_GET_ITEM(pools, i);
1925 indices[i]++;
Raymond Hettinger61024b92008-03-02 11:57:16 +00001926 if (indices[i] == PyTuple_GET_SIZE(pool)) {
Raymond Hettinger73d79632008-02-23 02:20:41 +00001927 /* Roll-over and advance to next pool */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001928 indices[i] = 0;
1929 elem = PyTuple_GET_ITEM(pool, 0);
1930 Py_INCREF(elem);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001931 oldelem = PyTuple_GET_ITEM(result, i);
1932 PyTuple_SET_ITEM(result, i, elem);
1933 Py_DECREF(oldelem);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001934 } else {
Raymond Hettinger73d79632008-02-23 02:20:41 +00001935 /* No rollover. Just increment and stop here. */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001936 elem = PyTuple_GET_ITEM(pool, indices[i]);
1937 Py_INCREF(elem);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001938 oldelem = PyTuple_GET_ITEM(result, i);
1939 PyTuple_SET_ITEM(result, i, elem);
1940 Py_DECREF(oldelem);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001941 break;
1942 }
1943 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00001944
1945 /* If i is negative, then the indices have all rolled-over
1946 and we're done. */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001947 if (i < 0)
Raymond Hettinger73d79632008-02-23 02:20:41 +00001948 goto empty;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001949 }
1950
Raymond Hettinger73d79632008-02-23 02:20:41 +00001951 Py_INCREF(result);
1952 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001953
1954empty:
1955 lz->stopped = 1;
1956 return NULL;
1957}
1958
1959PyDoc_STRVAR(product_doc,
1960"product(*iterables) --> product object\n\
1961\n\
Raymond Hettinger73d79632008-02-23 02:20:41 +00001962Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001963For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1964The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1965cycle in a manner similar to an odometer (with the rightmost element changing\n\
1966on every iteration).\n\n\
1967product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1968product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1969
1970static PyTypeObject product_type = {
1971 PyVarObject_HEAD_INIT(NULL, 0)
1972 "itertools.product", /* tp_name */
1973 sizeof(productobject), /* tp_basicsize */
1974 0, /* tp_itemsize */
1975 /* methods */
1976 (destructor)product_dealloc, /* tp_dealloc */
1977 0, /* tp_print */
1978 0, /* tp_getattr */
1979 0, /* tp_setattr */
1980 0, /* tp_compare */
1981 0, /* tp_repr */
1982 0, /* tp_as_number */
1983 0, /* tp_as_sequence */
1984 0, /* tp_as_mapping */
1985 0, /* tp_hash */
1986 0, /* tp_call */
1987 0, /* tp_str */
1988 PyObject_GenericGetAttr, /* tp_getattro */
1989 0, /* tp_setattro */
1990 0, /* tp_as_buffer */
1991 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1992 Py_TPFLAGS_BASETYPE, /* tp_flags */
1993 product_doc, /* tp_doc */
1994 (traverseproc)product_traverse, /* tp_traverse */
1995 0, /* tp_clear */
1996 0, /* tp_richcompare */
1997 0, /* tp_weaklistoffset */
1998 PyObject_SelfIter, /* tp_iter */
1999 (iternextfunc)product_next, /* tp_iternext */
2000 0, /* tp_methods */
2001 0, /* tp_members */
2002 0, /* tp_getset */
2003 0, /* tp_base */
2004 0, /* tp_dict */
2005 0, /* tp_descr_get */
2006 0, /* tp_descr_set */
2007 0, /* tp_dictoffset */
2008 0, /* tp_init */
2009 0, /* tp_alloc */
2010 product_new, /* tp_new */
2011 PyObject_GC_Del, /* tp_free */
2012};
2013
2014
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002015/* combinations object ************************************************************/
2016
2017typedef struct {
2018 PyObject_HEAD
2019 PyObject *pool; /* input converted to a tuple */
2020 Py_ssize_t *indices; /* one index per result element */
2021 PyObject *result; /* most recently returned result tuple */
2022 Py_ssize_t r; /* size of result tuple */
2023 int stopped; /* set to 1 when the combinations iterator is exhausted */
2024} combinationsobject;
2025
2026static PyTypeObject combinations_type;
2027
2028static PyObject *
2029combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2030{
2031 combinationsobject *co;
2032 Py_ssize_t n;
2033 Py_ssize_t r;
2034 PyObject *pool = NULL;
2035 PyObject *iterable = NULL;
2036 Py_ssize_t *indices = NULL;
2037 Py_ssize_t i;
2038 static char *kwargs[] = {"iterable", "r", NULL};
2039
2040 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2041 &iterable, &r))
2042 return NULL;
2043
2044 pool = PySequence_Tuple(iterable);
2045 if (pool == NULL)
2046 goto error;
2047 n = PyTuple_GET_SIZE(pool);
2048 if (r < 0) {
2049 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2050 goto error;
2051 }
2052 if (r > n) {
2053 PyErr_SetString(PyExc_ValueError, "r cannot be bigger than the iterable");
2054 goto error;
2055 }
2056
2057 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2058 if (indices == NULL) {
2059 PyErr_NoMemory();
2060 goto error;
2061 }
2062
2063 for (i=0 ; i<r ; i++)
2064 indices[i] = i;
2065
2066 /* create combinationsobject structure */
2067 co = (combinationsobject *)type->tp_alloc(type, 0);
2068 if (co == NULL)
2069 goto error;
2070
2071 co->pool = pool;
2072 co->indices = indices;
2073 co->result = NULL;
2074 co->r = r;
2075 co->stopped = 0;
2076
2077 return (PyObject *)co;
2078
2079error:
2080 if (indices != NULL)
2081 PyMem_Free(indices);
2082 Py_XDECREF(pool);
2083 return NULL;
2084}
2085
2086static void
2087combinations_dealloc(combinationsobject *co)
2088{
2089 PyObject_GC_UnTrack(co);
2090 Py_XDECREF(co->pool);
2091 Py_XDECREF(co->result);
2092 PyMem_Free(co->indices);
2093 Py_TYPE(co)->tp_free(co);
2094}
2095
2096static int
2097combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2098{
2099 Py_VISIT(co->pool);
2100 Py_VISIT(co->result);
2101 return 0;
2102}
2103
2104static PyObject *
2105combinations_next(combinationsobject *co)
2106{
2107 PyObject *elem;
2108 PyObject *oldelem;
2109 PyObject *pool = co->pool;
2110 Py_ssize_t *indices = co->indices;
2111 PyObject *result = co->result;
2112 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2113 Py_ssize_t r = co->r;
2114 Py_ssize_t i, j, index;
2115
2116 if (co->stopped)
2117 return NULL;
2118
2119 if (result == NULL) {
2120 /* On the first pass, initialize result tuple using the indices */
2121 result = PyTuple_New(r);
2122 if (result == NULL)
2123 goto empty;
2124 co->result = result;
2125 for (i=0; i<r ; i++) {
2126 index = indices[i];
2127 elem = PyTuple_GET_ITEM(pool, index);
2128 Py_INCREF(elem);
2129 PyTuple_SET_ITEM(result, i, elem);
2130 }
2131 } else {
2132 /* Copy the previous result tuple or re-use it if available */
2133 if (Py_REFCNT(result) > 1) {
2134 PyObject *old_result = result;
2135 result = PyTuple_New(r);
2136 if (result == NULL)
2137 goto empty;
2138 co->result = result;
2139 for (i=0; i<r ; i++) {
2140 elem = PyTuple_GET_ITEM(old_result, i);
2141 Py_INCREF(elem);
2142 PyTuple_SET_ITEM(result, i, elem);
2143 }
2144 Py_DECREF(old_result);
2145 }
Christian Heimescdddf182008-02-28 11:18:49 +00002146 /* Now, we've got the only copy so we can update it in-place
2147 * CPython's empty tuple is a singleton and cached in
2148 * PyTuple's freelist.
2149 */
2150 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002151
2152 /* Scan indices right-to-left until finding one that is not
2153 at its maximum (i + n - r). */
2154 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2155 ;
2156
2157 /* If i is negative, then the indices are all at
2158 their maximum value and we're done. */
2159 if (i < 0)
2160 goto empty;
2161
2162 /* Increment the current index which we know is not at its
2163 maximum. Then move back to the right setting each index
2164 to its lowest possible value (one higher than the index
2165 to its left -- this maintains the sort order invariant). */
2166 indices[i]++;
2167 for (j=i+1 ; j<r ; j++)
2168 indices[j] = indices[j-1] + 1;
2169
2170 /* Update the result tuple for the new indices
2171 starting with i, the leftmost index that changed */
2172 for ( ; i<r ; i++) {
2173 index = indices[i];
2174 elem = PyTuple_GET_ITEM(pool, index);
2175 Py_INCREF(elem);
2176 oldelem = PyTuple_GET_ITEM(result, i);
2177 PyTuple_SET_ITEM(result, i, elem);
2178 Py_DECREF(oldelem);
2179 }
2180 }
2181
2182 Py_INCREF(result);
2183 return result;
2184
2185empty:
2186 co->stopped = 1;
2187 return NULL;
2188}
2189
2190PyDoc_STRVAR(combinations_doc,
2191"combinations(iterables) --> combinations object\n\
2192\n\
2193Return successive r-length combinations of elements in the iterable.\n\n\
2194combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2195
2196static PyTypeObject combinations_type = {
2197 PyVarObject_HEAD_INIT(NULL, 0)
2198 "itertools.combinations", /* tp_name */
2199 sizeof(combinationsobject), /* tp_basicsize */
2200 0, /* tp_itemsize */
2201 /* methods */
2202 (destructor)combinations_dealloc, /* tp_dealloc */
2203 0, /* tp_print */
2204 0, /* tp_getattr */
2205 0, /* tp_setattr */
2206 0, /* tp_compare */
2207 0, /* tp_repr */
2208 0, /* tp_as_number */
2209 0, /* tp_as_sequence */
2210 0, /* tp_as_mapping */
2211 0, /* tp_hash */
2212 0, /* tp_call */
2213 0, /* tp_str */
2214 PyObject_GenericGetAttr, /* tp_getattro */
2215 0, /* tp_setattro */
2216 0, /* tp_as_buffer */
2217 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2218 Py_TPFLAGS_BASETYPE, /* tp_flags */
2219 combinations_doc, /* tp_doc */
2220 (traverseproc)combinations_traverse, /* tp_traverse */
2221 0, /* tp_clear */
2222 0, /* tp_richcompare */
2223 0, /* tp_weaklistoffset */
2224 PyObject_SelfIter, /* tp_iter */
2225 (iternextfunc)combinations_next, /* tp_iternext */
2226 0, /* tp_methods */
2227 0, /* tp_members */
2228 0, /* tp_getset */
2229 0, /* tp_base */
2230 0, /* tp_dict */
2231 0, /* tp_descr_get */
2232 0, /* tp_descr_set */
2233 0, /* tp_dictoffset */
2234 0, /* tp_init */
2235 0, /* tp_alloc */
2236 combinations_new, /* tp_new */
2237 PyObject_GC_Del, /* tp_free */
2238};
2239
2240
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002241/* permutations object ************************************************************
2242
2243def permutations(iterable, r=None):
2244 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2245 pool = tuple(iterable)
2246 n = len(pool)
2247 r = n if r is None else r
2248 indices = range(n)
2249 cycles = range(n-r+1, n+1)[::-1]
2250 yield tuple(pool[i] for i in indices[:r])
2251 while n:
2252 for i in reversed(range(r)):
2253 cycles[i] -= 1
2254 if cycles[i] == 0:
2255 indices[i:] = indices[i+1:] + indices[i:i+1]
2256 cycles[i] = n - i
2257 else:
2258 j = cycles[i]
2259 indices[i], indices[-j] = indices[-j], indices[i]
2260 yield tuple(pool[i] for i in indices[:r])
2261 break
2262 else:
2263 return
2264*/
2265
2266typedef struct {
2267 PyObject_HEAD
2268 PyObject *pool; /* input converted to a tuple */
2269 Py_ssize_t *indices; /* one index per element in the pool */
2270 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2271 PyObject *result; /* most recently returned result tuple */
2272 Py_ssize_t r; /* size of result tuple */
2273 int stopped; /* set to 1 when the permutations iterator is exhausted */
2274} permutationsobject;
2275
2276static PyTypeObject permutations_type;
2277
2278static PyObject *
2279permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2280{
2281 permutationsobject *po;
2282 Py_ssize_t n;
2283 Py_ssize_t r;
2284 PyObject *robj = Py_None;
2285 PyObject *pool = NULL;
2286 PyObject *iterable = NULL;
2287 Py_ssize_t *indices = NULL;
2288 Py_ssize_t *cycles = NULL;
2289 Py_ssize_t i;
2290 static char *kwargs[] = {"iterable", "r", NULL};
2291
2292 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2293 &iterable, &robj))
2294 return NULL;
2295
2296 pool = PySequence_Tuple(iterable);
2297 if (pool == NULL)
2298 goto error;
2299 n = PyTuple_GET_SIZE(pool);
2300
2301 r = n;
2302 if (robj != Py_None) {
2303 r = PyInt_AsSsize_t(robj);
2304 if (r == -1 && PyErr_Occurred())
2305 goto error;
2306 }
2307 if (r < 0) {
2308 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2309 goto error;
2310 }
2311 if (r > n) {
2312 PyErr_SetString(PyExc_ValueError, "r cannot be bigger than the iterable");
2313 goto error;
2314 }
2315
2316 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2317 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2318 if (indices == NULL || cycles == NULL) {
2319 PyErr_NoMemory();
2320 goto error;
2321 }
2322
2323 for (i=0 ; i<n ; i++)
2324 indices[i] = i;
2325 for (i=0 ; i<r ; i++)
2326 cycles[i] = n - i;
2327
2328 /* create permutationsobject structure */
2329 po = (permutationsobject *)type->tp_alloc(type, 0);
2330 if (po == NULL)
2331 goto error;
2332
2333 po->pool = pool;
2334 po->indices = indices;
2335 po->cycles = cycles;
2336 po->result = NULL;
2337 po->r = r;
2338 po->stopped = 0;
2339
2340 return (PyObject *)po;
2341
2342error:
2343 if (indices != NULL)
2344 PyMem_Free(indices);
2345 if (cycles != NULL)
2346 PyMem_Free(cycles);
2347 Py_XDECREF(pool);
2348 return NULL;
2349}
2350
2351static void
2352permutations_dealloc(permutationsobject *po)
2353{
2354 PyObject_GC_UnTrack(po);
2355 Py_XDECREF(po->pool);
2356 Py_XDECREF(po->result);
2357 PyMem_Free(po->indices);
2358 PyMem_Free(po->cycles);
2359 Py_TYPE(po)->tp_free(po);
2360}
2361
2362static int
2363permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2364{
2365 if (po->pool != NULL)
2366 Py_VISIT(po->pool);
2367 if (po->result != NULL)
2368 Py_VISIT(po->result);
2369 return 0;
2370}
2371
2372static PyObject *
2373permutations_next(permutationsobject *po)
2374{
2375 PyObject *elem;
2376 PyObject *oldelem;
2377 PyObject *pool = po->pool;
2378 Py_ssize_t *indices = po->indices;
2379 Py_ssize_t *cycles = po->cycles;
2380 PyObject *result = po->result;
2381 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2382 Py_ssize_t r = po->r;
2383 Py_ssize_t i, j, k, index;
2384
2385 if (po->stopped)
2386 return NULL;
2387
2388 if (result == NULL) {
2389 /* On the first pass, initialize result tuple using the indices */
2390 result = PyTuple_New(r);
2391 if (result == NULL)
2392 goto empty;
2393 po->result = result;
2394 for (i=0; i<r ; i++) {
2395 index = indices[i];
2396 elem = PyTuple_GET_ITEM(pool, index);
2397 Py_INCREF(elem);
2398 PyTuple_SET_ITEM(result, i, elem);
2399 }
2400 } else {
2401 if (n == 0)
2402 goto empty;
2403
2404 /* Copy the previous result tuple or re-use it if available */
2405 if (Py_REFCNT(result) > 1) {
2406 PyObject *old_result = result;
2407 result = PyTuple_New(r);
2408 if (result == NULL)
2409 goto empty;
2410 po->result = result;
2411 for (i=0; i<r ; i++) {
2412 elem = PyTuple_GET_ITEM(old_result, i);
2413 Py_INCREF(elem);
2414 PyTuple_SET_ITEM(result, i, elem);
2415 }
2416 Py_DECREF(old_result);
2417 }
2418 /* Now, we've got the only copy so we can update it in-place */
2419 assert(r == 0 || Py_REFCNT(result) == 1);
2420
2421 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2422 for (i=r-1 ; i>=0 ; i--) {
2423 cycles[i] -= 1;
2424 if (cycles[i] == 0) {
2425 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2426 index = indices[i];
2427 for (j=i ; j<n-1 ; j++)
2428 indices[j] = indices[j+1];
2429 indices[n-1] = index;
2430 cycles[i] = n - i;
2431 } else {
2432 j = cycles[i];
2433 index = indices[i];
2434 indices[i] = indices[n-j];
2435 indices[n-j] = index;
2436
2437 for (k=i; k<r ; k++) {
2438 /* start with i, the leftmost element that changed */
2439 /* yield tuple(pool[k] for k in indices[:r]) */
2440 index = indices[k];
2441 elem = PyTuple_GET_ITEM(pool, index);
2442 Py_INCREF(elem);
2443 oldelem = PyTuple_GET_ITEM(result, k);
2444 PyTuple_SET_ITEM(result, k, elem);
2445 Py_DECREF(oldelem);
2446 }
2447 break;
2448 }
2449 }
2450 /* If i is negative, then the cycles have all
2451 rolled-over and we're done. */
2452 if (i < 0)
2453 goto empty;
2454 }
2455 Py_INCREF(result);
2456 return result;
2457
2458empty:
2459 po->stopped = 1;
2460 return NULL;
2461}
2462
2463PyDoc_STRVAR(permutations_doc,
2464"permutations(iterables[, r]) --> permutations object\n\
2465\n\
2466Return successive r-length permutations of elements in the iterable.\n\n\
2467permutations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2468
2469static PyTypeObject permutations_type = {
2470 PyVarObject_HEAD_INIT(NULL, 0)
2471 "itertools.permutations", /* tp_name */
2472 sizeof(permutationsobject), /* tp_basicsize */
2473 0, /* tp_itemsize */
2474 /* methods */
2475 (destructor)permutations_dealloc, /* tp_dealloc */
2476 0, /* tp_print */
2477 0, /* tp_getattr */
2478 0, /* tp_setattr */
2479 0, /* tp_compare */
2480 0, /* tp_repr */
2481 0, /* tp_as_number */
2482 0, /* tp_as_sequence */
2483 0, /* tp_as_mapping */
2484 0, /* tp_hash */
2485 0, /* tp_call */
2486 0, /* tp_str */
2487 PyObject_GenericGetAttr, /* tp_getattro */
2488 0, /* tp_setattro */
2489 0, /* tp_as_buffer */
2490 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2491 Py_TPFLAGS_BASETYPE, /* tp_flags */
2492 permutations_doc, /* tp_doc */
2493 (traverseproc)permutations_traverse, /* tp_traverse */
2494 0, /* tp_clear */
2495 0, /* tp_richcompare */
2496 0, /* tp_weaklistoffset */
2497 PyObject_SelfIter, /* tp_iter */
2498 (iternextfunc)permutations_next, /* tp_iternext */
2499 0, /* tp_methods */
2500 0, /* tp_members */
2501 0, /* tp_getset */
2502 0, /* tp_base */
2503 0, /* tp_dict */
2504 0, /* tp_descr_get */
2505 0, /* tp_descr_set */
2506 0, /* tp_dictoffset */
2507 0, /* tp_init */
2508 0, /* tp_alloc */
2509 permutations_new, /* tp_new */
2510 PyObject_GC_Del, /* tp_free */
2511};
2512
2513
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002514/* ifilter object ************************************************************/
2515
2516typedef struct {
2517 PyObject_HEAD
2518 PyObject *func;
2519 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002520} ifilterobject;
2521
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002522static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002523
2524static PyObject *
2525ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2526{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002527 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002528 PyObject *it;
2529 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002530
Georg Brandlb84c1372007-01-21 10:28:43 +00002531 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002532 return NULL;
2533
Raymond Hettinger60eca932003-02-09 06:40:58 +00002534 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002535 return NULL;
2536
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002537 /* Get iterator. */
2538 it = PyObject_GetIter(seq);
2539 if (it == NULL)
2540 return NULL;
2541
2542 /* create ifilterobject structure */
2543 lz = (ifilterobject *)type->tp_alloc(type, 0);
2544 if (lz == NULL) {
2545 Py_DECREF(it);
2546 return NULL;
2547 }
2548 Py_INCREF(func);
2549 lz->func = func;
2550 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002551
2552 return (PyObject *)lz;
2553}
2554
2555static void
2556ifilter_dealloc(ifilterobject *lz)
2557{
2558 PyObject_GC_UnTrack(lz);
2559 Py_XDECREF(lz->func);
2560 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00002561 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002562}
2563
2564static int
2565ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
2566{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002567 Py_VISIT(lz->it);
2568 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002569 return 0;
2570}
2571
2572static PyObject *
2573ifilter_next(ifilterobject *lz)
2574{
2575 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002576 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002577 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002578 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002579
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002580 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00002581 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002582 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002583 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002584 if (item == NULL)
2585 return NULL;
2586
Facundo Batistab1d70e22008-02-25 23:46:02 +00002587 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002588 ok = PyObject_IsTrue(item);
2589 } else {
2590 PyObject *good;
2591 good = PyObject_CallFunctionObjArgs(lz->func,
2592 item, NULL);
2593 if (good == NULL) {
2594 Py_DECREF(item);
2595 return NULL;
2596 }
2597 ok = PyObject_IsTrue(good);
2598 Py_DECREF(good);
2599 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00002600 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002601 return item;
2602 Py_DECREF(item);
2603 }
2604}
2605
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002606PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002607"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002608\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002609Return those items of sequence for which function(item) is true.\n\
2610If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002611
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002612static PyTypeObject ifilter_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002613 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002614 "itertools.ifilter", /* tp_name */
2615 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002616 0, /* tp_itemsize */
2617 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002618 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002619 0, /* tp_print */
2620 0, /* tp_getattr */
2621 0, /* tp_setattr */
2622 0, /* tp_compare */
2623 0, /* tp_repr */
2624 0, /* tp_as_number */
2625 0, /* tp_as_sequence */
2626 0, /* tp_as_mapping */
2627 0, /* tp_hash */
2628 0, /* tp_call */
2629 0, /* tp_str */
2630 PyObject_GenericGetAttr, /* tp_getattro */
2631 0, /* tp_setattro */
2632 0, /* tp_as_buffer */
2633 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2634 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002635 ifilter_doc, /* tp_doc */
2636 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002637 0, /* tp_clear */
2638 0, /* tp_richcompare */
2639 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002640 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002641 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002642 0, /* tp_methods */
2643 0, /* tp_members */
2644 0, /* tp_getset */
2645 0, /* tp_base */
2646 0, /* tp_dict */
2647 0, /* tp_descr_get */
2648 0, /* tp_descr_set */
2649 0, /* tp_dictoffset */
2650 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002651 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002652 ifilter_new, /* tp_new */
2653 PyObject_GC_Del, /* tp_free */
2654};
2655
2656
2657/* ifilterfalse object ************************************************************/
2658
2659typedef struct {
2660 PyObject_HEAD
2661 PyObject *func;
2662 PyObject *it;
2663} ifilterfalseobject;
2664
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002665static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002666
2667static PyObject *
2668ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2669{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00002670 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002671 PyObject *it;
2672 ifilterfalseobject *lz;
2673
Georg Brandlb84c1372007-01-21 10:28:43 +00002674 if (type == &ifilterfalse_type &&
2675 !_PyArg_NoKeywords("ifilterfalse()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002676 return NULL;
2677
Raymond Hettinger60eca932003-02-09 06:40:58 +00002678 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
2679 return NULL;
2680
2681 /* Get iterator. */
2682 it = PyObject_GetIter(seq);
2683 if (it == NULL)
2684 return NULL;
2685
2686 /* create ifilterfalseobject structure */
2687 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
2688 if (lz == NULL) {
2689 Py_DECREF(it);
2690 return NULL;
2691 }
2692 Py_INCREF(func);
2693 lz->func = func;
2694 lz->it = it;
2695
2696 return (PyObject *)lz;
2697}
2698
2699static void
2700ifilterfalse_dealloc(ifilterfalseobject *lz)
2701{
2702 PyObject_GC_UnTrack(lz);
2703 Py_XDECREF(lz->func);
2704 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00002705 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002706}
2707
2708static int
2709ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
2710{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002711 Py_VISIT(lz->it);
2712 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002713 return 0;
2714}
2715
2716static PyObject *
2717ifilterfalse_next(ifilterfalseobject *lz)
2718{
2719 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002720 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002721 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002722 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002723
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002724 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00002725 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002726 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002727 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002728 if (item == NULL)
2729 return NULL;
2730
Facundo Batistab1d70e22008-02-25 23:46:02 +00002731 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
Raymond Hettinger60eca932003-02-09 06:40:58 +00002732 ok = PyObject_IsTrue(item);
2733 } else {
2734 PyObject *good;
2735 good = PyObject_CallFunctionObjArgs(lz->func,
2736 item, NULL);
2737 if (good == NULL) {
2738 Py_DECREF(item);
2739 return NULL;
2740 }
2741 ok = PyObject_IsTrue(good);
2742 Py_DECREF(good);
2743 }
2744 if (!ok)
2745 return item;
2746 Py_DECREF(item);
2747 }
2748}
2749
Raymond Hettinger60eca932003-02-09 06:40:58 +00002750PyDoc_STRVAR(ifilterfalse_doc,
2751"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
2752\n\
2753Return those items of sequence for which function(item) is false.\n\
2754If function is None, return the items that are false.");
2755
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002756static PyTypeObject ifilterfalse_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002757 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002758 "itertools.ifilterfalse", /* tp_name */
2759 sizeof(ifilterfalseobject), /* tp_basicsize */
2760 0, /* tp_itemsize */
2761 /* methods */
2762 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
2763 0, /* tp_print */
2764 0, /* tp_getattr */
2765 0, /* tp_setattr */
2766 0, /* tp_compare */
2767 0, /* tp_repr */
2768 0, /* tp_as_number */
2769 0, /* tp_as_sequence */
2770 0, /* tp_as_mapping */
2771 0, /* tp_hash */
2772 0, /* tp_call */
2773 0, /* tp_str */
2774 PyObject_GenericGetAttr, /* tp_getattro */
2775 0, /* tp_setattro */
2776 0, /* tp_as_buffer */
2777 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2778 Py_TPFLAGS_BASETYPE, /* tp_flags */
2779 ifilterfalse_doc, /* tp_doc */
2780 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
2781 0, /* tp_clear */
2782 0, /* tp_richcompare */
2783 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002784 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002785 (iternextfunc)ifilterfalse_next, /* tp_iternext */
2786 0, /* tp_methods */
2787 0, /* tp_members */
2788 0, /* tp_getset */
2789 0, /* tp_base */
2790 0, /* tp_dict */
2791 0, /* tp_descr_get */
2792 0, /* tp_descr_set */
2793 0, /* tp_dictoffset */
2794 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002795 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002796 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002797 PyObject_GC_Del, /* tp_free */
2798};
2799
2800
2801/* count object ************************************************************/
2802
2803typedef struct {
2804 PyObject_HEAD
Jack Diederich6c433a92006-05-26 11:15:17 +00002805 Py_ssize_t cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002806 PyObject *long_cnt; /* Arbitrarily large count when cnt >= PY_SSIZE_T_MAX */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002807} countobject;
2808
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002809static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002810
2811static PyObject *
2812count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2813{
2814 countobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00002815 Py_ssize_t cnt = 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002816 PyObject *cnt_arg = NULL;
2817 PyObject *long_cnt = NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002818
Georg Brandlb84c1372007-01-21 10:28:43 +00002819 if (type == &count_type && !_PyArg_NoKeywords("count()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002820 return NULL;
2821
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002822 if (!PyArg_UnpackTuple(args, "count", 0, 1, &cnt_arg))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002823 return NULL;
2824
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002825 if (cnt_arg != NULL) {
2826 cnt = PyInt_AsSsize_t(cnt_arg);
2827 if (cnt == -1 && PyErr_Occurred()) {
2828 PyErr_Clear();
2829 if (!PyLong_Check(cnt_arg)) {
2830 PyErr_SetString(PyExc_TypeError, "an integer is required");
2831 return NULL;
2832 }
2833 long_cnt = cnt_arg;
2834 Py_INCREF(long_cnt);
2835 cnt = PY_SSIZE_T_MAX;
2836 }
2837 }
2838
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002839 /* create countobject structure */
2840 lz = (countobject *)PyObject_New(countobject, &count_type);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002841 if (lz == NULL) {
2842 Py_XDECREF(long_cnt);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002843 return NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002844 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002845 lz->cnt = cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002846 lz->long_cnt = long_cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002847
2848 return (PyObject *)lz;
2849}
2850
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002851static void
2852count_dealloc(countobject *lz)
2853{
2854 Py_XDECREF(lz->long_cnt);
2855 PyObject_Del(lz);
2856}
2857
2858static PyObject *
2859count_nextlong(countobject *lz)
2860{
2861 static PyObject *one = NULL;
2862 PyObject *cnt;
2863 PyObject *stepped_up;
2864
2865 if (lz->long_cnt == NULL) {
2866 lz->long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
2867 if (lz->long_cnt == NULL)
2868 return NULL;
2869 }
2870 if (one == NULL) {
2871 one = PyInt_FromLong(1);
2872 if (one == NULL)
2873 return NULL;
2874 }
2875 cnt = lz->long_cnt;
2876 assert(cnt != NULL);
2877 stepped_up = PyNumber_Add(cnt, one);
2878 if (stepped_up == NULL)
2879 return NULL;
2880 lz->long_cnt = stepped_up;
2881 return cnt;
2882}
2883
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002884static PyObject *
2885count_next(countobject *lz)
2886{
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002887 if (lz->cnt == PY_SSIZE_T_MAX)
2888 return count_nextlong(lz);
Jack Diederich36234e82006-09-21 17:50:26 +00002889 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002890}
2891
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002892static PyObject *
2893count_repr(countobject *lz)
2894{
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002895 PyObject *cnt_repr;
2896 PyObject *result;
2897
2898 if (lz->cnt != PY_SSIZE_T_MAX)
2899 return PyString_FromFormat("count(%zd)", lz->cnt);
2900
2901 cnt_repr = PyObject_Repr(lz->long_cnt);
2902 if (cnt_repr == NULL)
2903 return NULL;
2904 result = PyString_FromFormat("count(%s)", PyString_AS_STRING(cnt_repr));
2905 Py_DECREF(cnt_repr);
2906 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002907}
2908
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002909PyDoc_STRVAR(count_doc,
2910"count([firstval]) --> count object\n\
2911\n\
2912Return a count object whose .next() method returns consecutive\n\
2913integers starting from zero or, if specified, from firstval.");
2914
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002915static PyTypeObject count_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002916 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002917 "itertools.count", /* tp_name */
2918 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002919 0, /* tp_itemsize */
2920 /* methods */
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002921 (destructor)count_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002922 0, /* tp_print */
2923 0, /* tp_getattr */
2924 0, /* tp_setattr */
2925 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002926 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002927 0, /* tp_as_number */
2928 0, /* tp_as_sequence */
2929 0, /* tp_as_mapping */
2930 0, /* tp_hash */
2931 0, /* tp_call */
2932 0, /* tp_str */
2933 PyObject_GenericGetAttr, /* tp_getattro */
2934 0, /* tp_setattro */
2935 0, /* tp_as_buffer */
2936 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002937 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002938 0, /* tp_traverse */
2939 0, /* tp_clear */
2940 0, /* tp_richcompare */
2941 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002942 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002943 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002944 0, /* tp_methods */
2945 0, /* tp_members */
2946 0, /* tp_getset */
2947 0, /* tp_base */
2948 0, /* tp_dict */
2949 0, /* tp_descr_get */
2950 0, /* tp_descr_set */
2951 0, /* tp_dictoffset */
2952 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002953 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002954 count_new, /* tp_new */
2955};
2956
2957
2958/* izip object ************************************************************/
2959
2960#include "Python.h"
2961
2962typedef struct {
2963 PyObject_HEAD
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002964 Py_ssize_t tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002965 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002966 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002967} izipobject;
2968
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002969static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002970
2971static PyObject *
2972izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2973{
2974 izipobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00002975 Py_ssize_t i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002976 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002977 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002978 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002979
Georg Brandlb84c1372007-01-21 10:28:43 +00002980 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002981 return NULL;
2982
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002983 /* args must be a tuple */
2984 assert(PyTuple_Check(args));
2985
2986 /* obtain iterators */
2987 ittuple = PyTuple_New(tuplesize);
Neal Norwitz2f3136b2006-05-27 05:18:57 +00002988 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002989 return NULL;
2990 for (i=0; i < tuplesize; ++i) {
2991 PyObject *item = PyTuple_GET_ITEM(args, i);
2992 PyObject *it = PyObject_GetIter(item);
2993 if (it == NULL) {
2994 if (PyErr_ExceptionMatches(PyExc_TypeError))
2995 PyErr_Format(PyExc_TypeError,
Neal Norwitz2f3136b2006-05-27 05:18:57 +00002996 "izip argument #%zd must support iteration",
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002997 i+1);
2998 Py_DECREF(ittuple);
2999 return NULL;
3000 }
3001 PyTuple_SET_ITEM(ittuple, i, it);
3002 }
3003
Raymond Hettinger2012f172003-02-07 05:32:58 +00003004 /* create a result holder */
3005 result = PyTuple_New(tuplesize);
3006 if (result == NULL) {
3007 Py_DECREF(ittuple);
3008 return NULL;
3009 }
3010 for (i=0 ; i < tuplesize ; i++) {
3011 Py_INCREF(Py_None);
3012 PyTuple_SET_ITEM(result, i, Py_None);
3013 }
3014
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003015 /* create izipobject structure */
3016 lz = (izipobject *)type->tp_alloc(type, 0);
3017 if (lz == NULL) {
3018 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003019 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003020 return NULL;
3021 }
3022 lz->ittuple = ittuple;
3023 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00003024 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003025
3026 return (PyObject *)lz;
3027}
3028
3029static void
3030izip_dealloc(izipobject *lz)
3031{
3032 PyObject_GC_UnTrack(lz);
3033 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003034 Py_XDECREF(lz->result);
Christian Heimese93237d2007-12-19 02:37:44 +00003035 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003036}
3037
3038static int
3039izip_traverse(izipobject *lz, visitproc visit, void *arg)
3040{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00003041 Py_VISIT(lz->ittuple);
3042 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003043 return 0;
3044}
3045
3046static PyObject *
3047izip_next(izipobject *lz)
3048{
Jack Diederich6c433a92006-05-26 11:15:17 +00003049 Py_ssize_t i;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00003050 Py_ssize_t tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00003051 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003052 PyObject *it;
3053 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00003054 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003055
Raymond Hettingerb5a42082003-08-08 05:10:41 +00003056 if (tuplesize == 0)
3057 return NULL;
Christian Heimese93237d2007-12-19 02:37:44 +00003058 if (Py_REFCNT(result) == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00003059 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003060 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003061 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00003062 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00003063 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00003064 if (item == NULL) {
3065 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003066 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00003067 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00003068 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003069 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00003070 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003071 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003072 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00003073 result = PyTuple_New(tuplesize);
3074 if (result == NULL)
3075 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003076 for (i=0 ; i < tuplesize ; i++) {
3077 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00003078 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00003079 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003080 if (item == NULL) {
3081 Py_DECREF(result);
3082 return NULL;
3083 }
3084 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003085 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003086 }
3087 return result;
3088}
3089
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003090PyDoc_STRVAR(izip_doc,
3091"izip(iter1 [,iter2 [...]]) --> izip object\n\
3092\n\
3093Return a izip object whose .next() method returns a tuple where\n\
3094the i-th element comes from the i-th iterable argument. The .next()\n\
3095method continues until the shortest iterable in the argument sequence\n\
3096is exhausted and then it raises StopIteration. Works like the zip()\n\
3097function but consumes less memory by returning an iterator instead of\n\
3098a list.");
3099
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003100static PyTypeObject izip_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003101 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003102 "itertools.izip", /* tp_name */
3103 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003104 0, /* tp_itemsize */
3105 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003106 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003107 0, /* tp_print */
3108 0, /* tp_getattr */
3109 0, /* tp_setattr */
3110 0, /* tp_compare */
3111 0, /* tp_repr */
3112 0, /* tp_as_number */
3113 0, /* tp_as_sequence */
3114 0, /* tp_as_mapping */
3115 0, /* tp_hash */
3116 0, /* tp_call */
3117 0, /* tp_str */
3118 PyObject_GenericGetAttr, /* tp_getattro */
3119 0, /* tp_setattro */
3120 0, /* tp_as_buffer */
3121 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3122 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003123 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003124 (traverseproc)izip_traverse, /* tp_traverse */
3125 0, /* tp_clear */
3126 0, /* tp_richcompare */
3127 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003128 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003129 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003130 0, /* tp_methods */
3131 0, /* tp_members */
3132 0, /* tp_getset */
3133 0, /* tp_base */
3134 0, /* tp_dict */
3135 0, /* tp_descr_get */
3136 0, /* tp_descr_set */
3137 0, /* tp_dictoffset */
3138 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003139 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003140 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003141 PyObject_GC_Del, /* tp_free */
3142};
3143
3144
3145/* repeat object ************************************************************/
3146
3147typedef struct {
3148 PyObject_HEAD
3149 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00003150 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003151} repeatobject;
3152
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003153static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003154
3155static PyObject *
3156repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3157{
3158 repeatobject *ro;
3159 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00003160 Py_ssize_t cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003161
Georg Brandlb84c1372007-01-21 10:28:43 +00003162 if (type == &repeat_type && !_PyArg_NoKeywords("repeat()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00003163 return NULL;
3164
Jack Diederich6c433a92006-05-26 11:15:17 +00003165 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003166 return NULL;
3167
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003168 if (PyTuple_Size(args) == 2 && cnt < 0)
3169 cnt = 0;
3170
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003171 ro = (repeatobject *)type->tp_alloc(type, 0);
3172 if (ro == NULL)
3173 return NULL;
3174 Py_INCREF(element);
3175 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003176 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003177 return (PyObject *)ro;
3178}
3179
3180static void
3181repeat_dealloc(repeatobject *ro)
3182{
3183 PyObject_GC_UnTrack(ro);
3184 Py_XDECREF(ro->element);
Christian Heimese93237d2007-12-19 02:37:44 +00003185 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003186}
3187
3188static int
3189repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3190{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00003191 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003192 return 0;
3193}
3194
3195static PyObject *
3196repeat_next(repeatobject *ro)
3197{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003198 if (ro->cnt == 0)
3199 return NULL;
3200 if (ro->cnt > 0)
3201 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003202 Py_INCREF(ro->element);
3203 return ro->element;
3204}
3205
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003206static PyObject *
3207repeat_repr(repeatobject *ro)
3208{
3209 PyObject *result, *objrepr;
3210
3211 objrepr = PyObject_Repr(ro->element);
3212 if (objrepr == NULL)
3213 return NULL;
3214
3215 if (ro->cnt == -1)
3216 result = PyString_FromFormat("repeat(%s)",
3217 PyString_AS_STRING(objrepr));
3218 else
Jack Diederich6c433a92006-05-26 11:15:17 +00003219 result = PyString_FromFormat("repeat(%s, %zd)",
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003220 PyString_AS_STRING(objrepr), ro->cnt);
3221 Py_DECREF(objrepr);
3222 return result;
3223}
3224
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003225static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003226repeat_len(repeatobject *ro)
3227{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003228 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003229 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003230 return NULL;
3231 }
Jack Diederich6c433a92006-05-26 11:15:17 +00003232 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003233}
3234
Armin Rigof5b3e362006-02-11 21:32:43 +00003235PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003236
3237static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00003238 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003239 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003240};
3241
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003242PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003243"repeat(element [,times]) -> create an iterator which returns the element\n\
3244for the specified number of times. If not specified, returns the element\n\
3245endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003246
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003247static PyTypeObject repeat_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003248 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003249 "itertools.repeat", /* tp_name */
3250 sizeof(repeatobject), /* tp_basicsize */
3251 0, /* tp_itemsize */
3252 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003253 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003254 0, /* tp_print */
3255 0, /* tp_getattr */
3256 0, /* tp_setattr */
3257 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003258 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003259 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003260 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003261 0, /* tp_as_mapping */
3262 0, /* tp_hash */
3263 0, /* tp_call */
3264 0, /* tp_str */
3265 PyObject_GenericGetAttr, /* tp_getattro */
3266 0, /* tp_setattro */
3267 0, /* tp_as_buffer */
3268 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3269 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003270 repeat_doc, /* tp_doc */
3271 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003272 0, /* tp_clear */
3273 0, /* tp_richcompare */
3274 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003275 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003276 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003277 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003278 0, /* tp_members */
3279 0, /* tp_getset */
3280 0, /* tp_base */
3281 0, /* tp_dict */
3282 0, /* tp_descr_get */
3283 0, /* tp_descr_set */
3284 0, /* tp_dictoffset */
3285 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003286 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003287 repeat_new, /* tp_new */
3288 PyObject_GC_Del, /* tp_free */
3289};
3290
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003291/* iziplongest object ************************************************************/
3292
3293#include "Python.h"
3294
3295typedef struct {
3296 PyObject_HEAD
3297 Py_ssize_t tuplesize;
3298 Py_ssize_t numactive;
3299 PyObject *ittuple; /* tuple of iterators */
3300 PyObject *result;
3301 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003302} iziplongestobject;
3303
3304static PyTypeObject iziplongest_type;
3305
3306static PyObject *
3307izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3308{
3309 iziplongestobject *lz;
3310 Py_ssize_t i;
3311 PyObject *ittuple; /* tuple of iterators */
3312 PyObject *result;
3313 PyObject *fillvalue = Py_None;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003314 Py_ssize_t tuplesize = PySequence_Length(args);
3315
3316 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3317 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3318 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3319 PyErr_SetString(PyExc_TypeError,
3320 "izip_longest() got an unexpected keyword argument");
3321 return NULL;
3322 }
3323 }
3324
3325 /* args must be a tuple */
3326 assert(PyTuple_Check(args));
3327
3328 /* obtain iterators */
3329 ittuple = PyTuple_New(tuplesize);
3330 if (ittuple == NULL)
3331 return NULL;
3332 for (i=0; i < tuplesize; ++i) {
3333 PyObject *item = PyTuple_GET_ITEM(args, i);
3334 PyObject *it = PyObject_GetIter(item);
3335 if (it == NULL) {
3336 if (PyErr_ExceptionMatches(PyExc_TypeError))
3337 PyErr_Format(PyExc_TypeError,
3338 "izip_longest argument #%zd must support iteration",
3339 i+1);
3340 Py_DECREF(ittuple);
3341 return NULL;
3342 }
3343 PyTuple_SET_ITEM(ittuple, i, it);
3344 }
3345
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003346 /* create a result holder */
3347 result = PyTuple_New(tuplesize);
3348 if (result == NULL) {
3349 Py_DECREF(ittuple);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003350 return NULL;
3351 }
3352 for (i=0 ; i < tuplesize ; i++) {
3353 Py_INCREF(Py_None);
3354 PyTuple_SET_ITEM(result, i, Py_None);
3355 }
3356
3357 /* create iziplongestobject structure */
3358 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3359 if (lz == NULL) {
3360 Py_DECREF(ittuple);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003361 Py_DECREF(result);
3362 return NULL;
3363 }
3364 lz->ittuple = ittuple;
3365 lz->tuplesize = tuplesize;
3366 lz->numactive = tuplesize;
3367 lz->result = result;
3368 Py_INCREF(fillvalue);
3369 lz->fillvalue = fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003370 return (PyObject *)lz;
3371}
3372
3373static void
3374izip_longest_dealloc(iziplongestobject *lz)
3375{
3376 PyObject_GC_UnTrack(lz);
3377 Py_XDECREF(lz->ittuple);
3378 Py_XDECREF(lz->result);
3379 Py_XDECREF(lz->fillvalue);
Christian Heimese93237d2007-12-19 02:37:44 +00003380 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003381}
3382
3383static int
3384izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3385{
3386 Py_VISIT(lz->ittuple);
3387 Py_VISIT(lz->result);
3388 Py_VISIT(lz->fillvalue);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003389 return 0;
3390}
3391
3392static PyObject *
3393izip_longest_next(iziplongestobject *lz)
3394{
3395 Py_ssize_t i;
3396 Py_ssize_t tuplesize = lz->tuplesize;
3397 PyObject *result = lz->result;
3398 PyObject *it;
3399 PyObject *item;
3400 PyObject *olditem;
3401
3402 if (tuplesize == 0)
3403 return NULL;
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003404 if (lz->numactive == 0)
3405 return NULL;
Christian Heimese93237d2007-12-19 02:37:44 +00003406 if (Py_REFCNT(result) == 1) {
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003407 Py_INCREF(result);
3408 for (i=0 ; i < tuplesize ; i++) {
3409 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003410 if (it == NULL) {
3411 Py_INCREF(lz->fillvalue);
3412 item = lz->fillvalue;
3413 } else {
3414 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00003415 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003416 if (item == NULL) {
3417 lz->numactive -= 1;
3418 if (lz->numactive == 0) {
3419 Py_DECREF(result);
3420 return NULL;
3421 } else {
3422 Py_INCREF(lz->fillvalue);
3423 item = lz->fillvalue;
3424 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3425 Py_DECREF(it);
3426 }
3427 }
3428 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003429 olditem = PyTuple_GET_ITEM(result, i);
3430 PyTuple_SET_ITEM(result, i, item);
3431 Py_DECREF(olditem);
3432 }
3433 } else {
3434 result = PyTuple_New(tuplesize);
3435 if (result == NULL)
3436 return NULL;
3437 for (i=0 ; i < tuplesize ; i++) {
3438 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003439 if (it == NULL) {
3440 Py_INCREF(lz->fillvalue);
3441 item = lz->fillvalue;
3442 } else {
3443 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00003444 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003445 if (item == NULL) {
3446 lz->numactive -= 1;
3447 if (lz->numactive == 0) {
3448 Py_DECREF(result);
3449 return NULL;
3450 } else {
3451 Py_INCREF(lz->fillvalue);
3452 item = lz->fillvalue;
3453 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3454 Py_DECREF(it);
3455 }
3456 }
3457 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003458 PyTuple_SET_ITEM(result, i, item);
3459 }
3460 }
3461 return result;
3462}
3463
3464PyDoc_STRVAR(izip_longest_doc,
3465"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3466\n\
3467Return an izip_longest object whose .next() method returns a tuple where\n\
3468the i-th element comes from the i-th iterable argument. The .next()\n\
3469method continues until the longest iterable in the argument sequence\n\
3470is exhausted and then it raises StopIteration. When the shorter iterables\n\
3471are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3472defaults to None or can be specified by a keyword argument.\n\
3473");
3474
3475static PyTypeObject iziplongest_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003476 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003477 "itertools.izip_longest", /* tp_name */
3478 sizeof(iziplongestobject), /* tp_basicsize */
3479 0, /* tp_itemsize */
3480 /* methods */
3481 (destructor)izip_longest_dealloc, /* tp_dealloc */
3482 0, /* tp_print */
3483 0, /* tp_getattr */
3484 0, /* tp_setattr */
3485 0, /* tp_compare */
3486 0, /* tp_repr */
3487 0, /* tp_as_number */
3488 0, /* tp_as_sequence */
3489 0, /* tp_as_mapping */
3490 0, /* tp_hash */
3491 0, /* tp_call */
3492 0, /* tp_str */
3493 PyObject_GenericGetAttr, /* tp_getattro */
3494 0, /* tp_setattro */
3495 0, /* tp_as_buffer */
3496 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3497 Py_TPFLAGS_BASETYPE, /* tp_flags */
3498 izip_longest_doc, /* tp_doc */
3499 (traverseproc)izip_longest_traverse, /* tp_traverse */
3500 0, /* tp_clear */
3501 0, /* tp_richcompare */
3502 0, /* tp_weaklistoffset */
3503 PyObject_SelfIter, /* tp_iter */
3504 (iternextfunc)izip_longest_next, /* tp_iternext */
3505 0, /* tp_methods */
3506 0, /* tp_members */
3507 0, /* tp_getset */
3508 0, /* tp_base */
3509 0, /* tp_dict */
3510 0, /* tp_descr_get */
3511 0, /* tp_descr_set */
3512 0, /* tp_dictoffset */
3513 0, /* tp_init */
3514 0, /* tp_alloc */
3515 izip_longest_new, /* tp_new */
3516 PyObject_GC_Del, /* tp_free */
3517};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003518
3519/* module level code ********************************************************/
3520
3521PyDoc_STRVAR(module_doc,
3522"Functional tools for creating and using iterators.\n\
3523\n\
3524Infinite iterators:\n\
3525count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003526cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003527repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003528\n\
3529Iterators terminating on the shortest input sequence:\n\
3530izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003531izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003532ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
3533ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003534islice(seq, [start,] stop [, step]) --> elements from\n\
3535 seq[start:stop:step]\n\
3536imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
3537starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003538tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003539chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003540takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
3541dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003542groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003543");
3544
3545
Raymond Hettingerad983e72003-11-12 14:32:26 +00003546static PyMethodDef module_methods[] = {
3547 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3548 {NULL, NULL} /* sentinel */
3549};
3550
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003551PyMODINIT_FUNC
3552inititertools(void)
3553{
Raymond Hettinger60eca932003-02-09 06:40:58 +00003554 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003555 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003556 char *name;
3557 PyTypeObject *typelist[] = {
Raymond Hettinger93e804d2008-02-26 23:40:50 +00003558 &combinations_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003559 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003560 &dropwhile_type,
3561 &takewhile_type,
3562 &islice_type,
3563 &starmap_type,
3564 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003565 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003566 &ifilter_type,
3567 &ifilterfalse_type,
3568 &count_type,
3569 &izip_type,
Raymond Hettinger50986cc2008-02-22 03:16:42 +00003570 &iziplongest_type,
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00003571 &permutations_type,
Raymond Hettinger93e804d2008-02-26 23:40:50 +00003572 &product_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003573 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003574 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003575 NULL
3576 };
3577
Christian Heimese93237d2007-12-19 02:37:44 +00003578 Py_TYPE(&teedataobject_type) = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00003579 m = Py_InitModule3("itertools", module_methods, module_doc);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00003580 if (m == NULL)
3581 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003582
Raymond Hettinger60eca932003-02-09 06:40:58 +00003583 for (i=0 ; typelist[i] != NULL ; i++) {
3584 if (PyType_Ready(typelist[i]) < 0)
3585 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003586 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003587 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003588 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003589 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003590 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003591
3592 if (PyType_Ready(&teedataobject_type) < 0)
3593 return;
3594 if (PyType_Ready(&tee_type) < 0)
3595 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003596 if (PyType_Ready(&_grouper_type) < 0)
3597 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003598}