blob: fa7360516665b9265c3f0ffb520784516dae334c [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 Heimes90aa7642007-12-19 02:45:37 +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öwis9f2e3462007-07-21 17:22:18 +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öwis9f2e3462007-07-21 17:22:18 +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 Wouters49fd7fa2006-04-21 10:40:58 +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 Wouters49fd7fa2006-04-21 10:40:58 +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);
Thomas Wouters0e3f5912006-08-11 14:57:12 +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 Wouters49fd7fa2006-04-21 10:40:58 +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 Wouters49fd7fa2006-04-21 10:40:58 +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öwis9f2e3462007-07-21 17:22:18 +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 Wouters49fd7fa2006-04-21 10:40:58 +0000436 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000437 teedataobject_doc, /* tp_doc */
Thomas Wouters49fd7fa2006-04-21 10:40:58 +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);
Thomas Wouters0e3f5912006-08-11 14:57:12 +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 Wouters49fd7fa2006-04-21 10:40:58 +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 Wouters49fd7fa2006-04-21 10:40:58 +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 Wouters49fd7fa2006-04-21 10:40:58 +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 Wouters49fd7fa2006-04-21 10:40:58 +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);
Thomas Wouters0e3f5912006-08-11 14:57:12 +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 Wouters49fd7fa2006-04-21 10:40:58 +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 Wouters49fd7fa2006-04-21 10:40:58 +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 Wouters49fd7fa2006-04-21 10:40:58 +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öwis9f2e3462007-07-21 17:22:18 +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 Wouters49fd7fa2006-04-21 10:40:58 +0000592 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000593 teeobject_doc, /* tp_doc */
Thomas Wouters49fd7fa2006-04-21 10:40:58 +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 Wouters49fd7fa2006-04-21 10:40:58 +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{
Thomas Wouters89f507f2006-12-13 04:49:30 +0000617 Py_ssize_t i, n=2;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000618 PyObject *it, *iterable, *copyable, *result;
619
Thomas Wouters89f507f2006-12-13 04:49:30 +0000620 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000621 return NULL;
Thomas Wouters89f507f2006-12-13 04:49:30 +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
Thomas Woutersb2137042007-02-01 18:02:27 +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 Heimes90aa7642007-12-19 02:45:37 +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öwis9f2e3462007-07-21 17:22:18 +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
Thomas Woutersb2137042007-02-01 18:02:27 +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 Heimes90aa7642007-12-19 02:45:37 +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 Heimes90aa7642007-12-19 02:45:37 +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öwis9f2e3462007-07-21 17:22:18 +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
Thomas Woutersb2137042007-02-01 18:02:27 +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 Heimes90aa7642007-12-19 02:45:37 +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 Heimes90aa7642007-12-19 02:45:37 +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öwis9f2e3462007-07-21 17:22:18 +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;
Thomas Wouters477c8d52006-05-27 19:21:47 +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;
Thomas Wouters477c8d52006-05-27 19:21:47 +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
Thomas Woutersb2137042007-02-01 18:02:27 +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) {
Christian Heimes217cfd12007-12-02 14:31:20 +00001125 stop = PyLong_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)
Christian Heimes217cfd12007-12-02 14:31:20 +00001136 start = PyLong_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) {
Christian Heimes217cfd12007-12-02 14:31:20 +00001140 stop = PyLong_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)
Christian Heimes217cfd12007-12-02 14:31:20 +00001158 step = PyLong_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 Heimes90aa7642007-12-19 02:45:37 +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;
Thomas Wouters477c8d52006-05-27 19:21:47 +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 Heimes90aa7642007-12-19 02:45:37 +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öwis9f2e3462007-07-21 17:22:18 +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
Thomas Woutersb2137042007-02-01 18:02:27 +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 Heimes90aa7642007-12-19 02:45:37 +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 Heimes90aa7642007-12-19 02:45:37 +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)) {
Christian Heimes679db4a2008-01-18 09:56:22 +00001359 PyObject *newargs = PySequence_Tuple(args);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001360 Py_DECREF(args);
Christian Heimes679db4a2008-01-18 09:56:22 +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öwis9f2e3462007-07-21 17:22:18 +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
Thomas Woutersb2137042007-02-01 18:02:27 +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 Heimes90aa7642007-12-19 02:45:37 +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
1493static PyObject *
1494imap_next(imapobject *lz)
1495{
1496 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001497 PyObject *argtuple;
1498 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001499 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001500
1501 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001502 argtuple = PyTuple_New(numargs);
1503 if (argtuple == NULL)
1504 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001505
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001506 for (i=0 ; i<numargs ; i++) {
1507 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1508 if (val == NULL) {
1509 Py_DECREF(argtuple);
1510 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001511 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001512 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001513 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001514 result = PyObject_Call(lz->func, argtuple, NULL);
1515 Py_DECREF(argtuple);
1516 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001517}
1518
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001519PyDoc_STRVAR(imap_doc,
1520"imap(func, *iterables) --> imap object\n\
1521\n\
1522Make an iterator that computes the function using arguments from\n\
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +00001523each of the iterables. Stops when the shortest iterable is exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001524
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001525static PyTypeObject imap_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001526 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001527 "itertools.imap", /* tp_name */
1528 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001529 0, /* tp_itemsize */
1530 /* methods */
1531 (destructor)imap_dealloc, /* tp_dealloc */
1532 0, /* tp_print */
1533 0, /* tp_getattr */
1534 0, /* tp_setattr */
1535 0, /* tp_compare */
1536 0, /* tp_repr */
1537 0, /* tp_as_number */
1538 0, /* tp_as_sequence */
1539 0, /* tp_as_mapping */
1540 0, /* tp_hash */
1541 0, /* tp_call */
1542 0, /* tp_str */
1543 PyObject_GenericGetAttr, /* tp_getattro */
1544 0, /* tp_setattro */
1545 0, /* tp_as_buffer */
1546 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1547 Py_TPFLAGS_BASETYPE, /* tp_flags */
1548 imap_doc, /* tp_doc */
1549 (traverseproc)imap_traverse, /* tp_traverse */
1550 0, /* tp_clear */
1551 0, /* tp_richcompare */
1552 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001553 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001554 (iternextfunc)imap_next, /* tp_iternext */
1555 0, /* tp_methods */
1556 0, /* tp_members */
1557 0, /* tp_getset */
1558 0, /* tp_base */
1559 0, /* tp_dict */
1560 0, /* tp_descr_get */
1561 0, /* tp_descr_set */
1562 0, /* tp_dictoffset */
1563 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001564 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001565 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001566 PyObject_GC_Del, /* tp_free */
1567};
1568
1569
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001570/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001571
1572typedef struct {
1573 PyObject_HEAD
Thomas Wouters477c8d52006-05-27 19:21:47 +00001574 Py_ssize_t tuplesize;
1575 Py_ssize_t iternum; /* which iterator is active */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001576 PyObject *ittuple; /* tuple of iterators */
1577} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001578
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001579static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001580
1581static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001582chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001583{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001584 chainobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001585 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouters477c8d52006-05-27 19:21:47 +00001586 Py_ssize_t i;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001587 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001588
Thomas Woutersb2137042007-02-01 18:02:27 +00001589 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001590 return NULL;
1591
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001592 /* obtain iterators */
1593 assert(PyTuple_Check(args));
1594 ittuple = PyTuple_New(tuplesize);
Thomas Wouters477c8d52006-05-27 19:21:47 +00001595 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001596 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001597 for (i=0; i < tuplesize; ++i) {
1598 PyObject *item = PyTuple_GET_ITEM(args, i);
1599 PyObject *it = PyObject_GetIter(item);
1600 if (it == NULL) {
1601 if (PyErr_ExceptionMatches(PyExc_TypeError))
1602 PyErr_Format(PyExc_TypeError,
Thomas Wouters477c8d52006-05-27 19:21:47 +00001603 "chain argument #%zd must support iteration",
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001604 i+1);
1605 Py_DECREF(ittuple);
1606 return NULL;
1607 }
1608 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001609 }
1610
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001611 /* create chainobject structure */
1612 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001613 if (lz == NULL) {
1614 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001615 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001616 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001617
1618 lz->ittuple = ittuple;
1619 lz->iternum = 0;
1620 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001621
1622 return (PyObject *)lz;
1623}
1624
1625static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001626chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001627{
1628 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001629 Py_XDECREF(lz->ittuple);
Christian Heimes90aa7642007-12-19 02:45:37 +00001630 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001631}
1632
Raymond Hettinger2012f172003-02-07 05:32:58 +00001633static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001634chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001635{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001636 Py_VISIT(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001637 return 0;
1638}
1639
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001640static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001641chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001642{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001643 PyObject *it;
1644 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001645
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001646 while (lz->iternum < lz->tuplesize) {
1647 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1648 item = PyIter_Next(it);
1649 if (item != NULL)
1650 return item;
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001651 if (PyErr_Occurred()) {
1652 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1653 PyErr_Clear();
1654 else
1655 return NULL;
1656 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001657 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001658 }
1659 return NULL;
1660}
1661
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001662PyDoc_STRVAR(chain_doc,
1663"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001664\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001665Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001666first iterable until it is exhausted, then elements from the next\n\
1667iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001668
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001669static PyTypeObject chain_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001670 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001671 "itertools.chain", /* tp_name */
1672 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001673 0, /* tp_itemsize */
1674 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001675 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001676 0, /* tp_print */
1677 0, /* tp_getattr */
1678 0, /* tp_setattr */
1679 0, /* tp_compare */
1680 0, /* tp_repr */
1681 0, /* tp_as_number */
1682 0, /* tp_as_sequence */
1683 0, /* tp_as_mapping */
1684 0, /* tp_hash */
1685 0, /* tp_call */
1686 0, /* tp_str */
1687 PyObject_GenericGetAttr, /* tp_getattro */
1688 0, /* tp_setattro */
1689 0, /* tp_as_buffer */
1690 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1691 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001692 chain_doc, /* tp_doc */
1693 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001694 0, /* tp_clear */
1695 0, /* tp_richcompare */
1696 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001697 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001698 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001699 0, /* tp_methods */
1700 0, /* tp_members */
1701 0, /* tp_getset */
1702 0, /* tp_base */
1703 0, /* tp_dict */
1704 0, /* tp_descr_get */
1705 0, /* tp_descr_set */
1706 0, /* tp_dictoffset */
1707 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001708 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001709 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001710 PyObject_GC_Del, /* tp_free */
1711};
1712
1713
Christian Heimesc3f30c42008-02-22 16:37:40 +00001714/* product object ************************************************************/
1715
1716typedef struct {
1717 PyObject_HEAD
1718 PyObject *pools; /* tuple of pool tuples */
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001719 Py_ssize_t *maxvec; /* size of each pool */
1720 Py_ssize_t *indices; /* one index per pool */
1721 PyObject *result; /* most recently returned result tuple */
1722 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001723} productobject;
1724
1725static PyTypeObject product_type;
1726
1727static PyObject *
1728product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1729{
1730 productobject *lz;
1731 Py_ssize_t npools;
1732 PyObject *pools = NULL;
1733 Py_ssize_t *maxvec = NULL;
1734 Py_ssize_t *indices = NULL;
1735 Py_ssize_t i;
1736
1737 if (type == &product_type && !_PyArg_NoKeywords("product()", kwds))
1738 return NULL;
1739
1740 assert(PyTuple_Check(args));
1741 npools = PyTuple_GET_SIZE(args);
1742
1743 maxvec = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1744 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1745 if (maxvec == NULL || indices == NULL) {
1746 PyErr_NoMemory();
1747 goto error;
1748 }
1749
1750 pools = PyTuple_New(npools);
1751 if (pools == NULL)
1752 goto error;
1753
1754 for (i=0; i < npools; ++i) {
1755 PyObject *item = PyTuple_GET_ITEM(args, i);
1756 PyObject *pool = PySequence_Tuple(item);
1757 if (pool == NULL)
1758 goto error;
1759
1760 PyTuple_SET_ITEM(pools, i, pool);
1761 maxvec[i] = PyTuple_GET_SIZE(pool);
1762 indices[i] = 0;
1763 }
1764
1765 /* create productobject structure */
1766 lz = (productobject *)type->tp_alloc(type, 0);
Christian Heimes380f7f22008-02-28 11:19:05 +00001767 if (lz == NULL)
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001768 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001769
1770 lz->pools = pools;
1771 lz->maxvec = maxvec;
1772 lz->indices = indices;
1773 lz->result = NULL;
1774 lz->stopped = 0;
1775
1776 return (PyObject *)lz;
1777
1778error:
1779 if (maxvec != NULL)
1780 PyMem_Free(maxvec);
1781 if (indices != NULL)
1782 PyMem_Free(indices);
1783 Py_XDECREF(pools);
1784 return NULL;
1785}
1786
1787static void
1788product_dealloc(productobject *lz)
1789{
1790 PyObject_GC_UnTrack(lz);
1791 Py_XDECREF(lz->pools);
1792 Py_XDECREF(lz->result);
1793 PyMem_Free(lz->maxvec);
1794 PyMem_Free(lz->indices);
1795 Py_TYPE(lz)->tp_free(lz);
1796}
1797
1798static int
1799product_traverse(productobject *lz, visitproc visit, void *arg)
1800{
1801 Py_VISIT(lz->pools);
1802 Py_VISIT(lz->result);
1803 return 0;
1804}
1805
1806static PyObject *
1807product_next(productobject *lz)
1808{
1809 PyObject *pool;
1810 PyObject *elem;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001811 PyObject *oldelem;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001812 PyObject *pools = lz->pools;
1813 PyObject *result = lz->result;
1814 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1815 Py_ssize_t i;
1816
1817 if (lz->stopped)
1818 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001819
Christian Heimesc3f30c42008-02-22 16:37:40 +00001820 if (result == NULL) {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001821 /* On the first pass, return an initial tuple filled with the
1822 first element from each pool. If any pool is empty, then
1823 whole product is empty and we're already done */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001824 if (npools == 0)
1825 goto empty;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001826 result = PyTuple_New(npools);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001827 if (result == NULL)
1828 goto empty;
1829 lz->result = result;
1830 for (i=0; i < npools; i++) {
1831 pool = PyTuple_GET_ITEM(pools, i);
1832 if (PyTuple_GET_SIZE(pool) == 0)
1833 goto empty;
1834 elem = PyTuple_GET_ITEM(pool, 0);
1835 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001836 PyTuple_SET_ITEM(result, i, elem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001837 }
1838 } else {
1839 Py_ssize_t *indices = lz->indices;
1840 Py_ssize_t *maxvec = lz->maxvec;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001841
1842 /* Copy the previous result tuple or re-use it if available */
1843 if (Py_REFCNT(result) > 1) {
1844 PyObject *old_result = result;
1845 result = PyTuple_New(npools);
1846 if (result == NULL)
1847 goto empty;
1848 lz->result = result;
1849 for (i=0; i < npools; i++) {
1850 elem = PyTuple_GET_ITEM(old_result, i);
1851 Py_INCREF(elem);
1852 PyTuple_SET_ITEM(result, i, elem);
1853 }
1854 Py_DECREF(old_result);
1855 }
1856 /* Now, we've got the only copy so we can update it in-place */
1857 assert (Py_REFCNT(result) == 1);
1858
1859 /* Update the pool indices right-to-left. Only advance to the
1860 next pool when the previous one rolls-over */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001861 for (i=npools-1 ; i >= 0 ; i--) {
1862 pool = PyTuple_GET_ITEM(pools, i);
1863 indices[i]++;
1864 if (indices[i] == maxvec[i]) {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001865 /* Roll-over and advance to next pool */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001866 indices[i] = 0;
1867 elem = PyTuple_GET_ITEM(pool, 0);
1868 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001869 oldelem = PyTuple_GET_ITEM(result, i);
1870 PyTuple_SET_ITEM(result, i, elem);
1871 Py_DECREF(oldelem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001872 } else {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001873 /* No rollover. Just increment and stop here. */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001874 elem = PyTuple_GET_ITEM(pool, indices[i]);
1875 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001876 oldelem = PyTuple_GET_ITEM(result, i);
1877 PyTuple_SET_ITEM(result, i, elem);
1878 Py_DECREF(oldelem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001879 break;
1880 }
1881 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001882
1883 /* If i is negative, then the indices have all rolled-over
1884 and we're done. */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001885 if (i < 0)
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001886 goto empty;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001887 }
1888
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001889 Py_INCREF(result);
1890 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001891
1892empty:
1893 lz->stopped = 1;
1894 return NULL;
1895}
1896
1897PyDoc_STRVAR(product_doc,
1898"product(*iterables) --> product object\n\
1899\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001900Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001901For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1902The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1903cycle in a manner similar to an odometer (with the rightmost element changing\n\
1904on every iteration).\n\n\
1905product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1906product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1907
1908static PyTypeObject product_type = {
1909 PyVarObject_HEAD_INIT(NULL, 0)
1910 "itertools.product", /* tp_name */
1911 sizeof(productobject), /* tp_basicsize */
1912 0, /* tp_itemsize */
1913 /* methods */
1914 (destructor)product_dealloc, /* tp_dealloc */
1915 0, /* tp_print */
1916 0, /* tp_getattr */
1917 0, /* tp_setattr */
1918 0, /* tp_compare */
1919 0, /* tp_repr */
1920 0, /* tp_as_number */
1921 0, /* tp_as_sequence */
1922 0, /* tp_as_mapping */
1923 0, /* tp_hash */
1924 0, /* tp_call */
1925 0, /* tp_str */
1926 PyObject_GenericGetAttr, /* tp_getattro */
1927 0, /* tp_setattro */
1928 0, /* tp_as_buffer */
1929 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1930 Py_TPFLAGS_BASETYPE, /* tp_flags */
1931 product_doc, /* tp_doc */
1932 (traverseproc)product_traverse, /* tp_traverse */
1933 0, /* tp_clear */
1934 0, /* tp_richcompare */
1935 0, /* tp_weaklistoffset */
1936 PyObject_SelfIter, /* tp_iter */
1937 (iternextfunc)product_next, /* tp_iternext */
1938 0, /* tp_methods */
1939 0, /* tp_members */
1940 0, /* tp_getset */
1941 0, /* tp_base */
1942 0, /* tp_dict */
1943 0, /* tp_descr_get */
1944 0, /* tp_descr_set */
1945 0, /* tp_dictoffset */
1946 0, /* tp_init */
1947 0, /* tp_alloc */
1948 product_new, /* tp_new */
1949 PyObject_GC_Del, /* tp_free */
1950};
1951
1952
Christian Heimes380f7f22008-02-28 11:19:05 +00001953/* combinations object ************************************************************/
1954
1955typedef struct {
1956 PyObject_HEAD
1957 PyObject *pool; /* input converted to a tuple */
1958 Py_ssize_t *indices; /* one index per result element */
1959 PyObject *result; /* most recently returned result tuple */
1960 Py_ssize_t r; /* size of result tuple */
1961 int stopped; /* set to 1 when the combinations iterator is exhausted */
1962} combinationsobject;
1963
1964static PyTypeObject combinations_type;
1965
1966static PyObject *
1967combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1968{
1969 combinationsobject *co;
1970 Py_ssize_t n;
1971 Py_ssize_t r;
1972 PyObject *pool = NULL;
1973 PyObject *iterable = NULL;
1974 Py_ssize_t *indices = NULL;
1975 Py_ssize_t i;
1976 static char *kwargs[] = {"iterable", "r", NULL};
1977
1978 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
1979 &iterable, &r))
1980 return NULL;
1981
1982 pool = PySequence_Tuple(iterable);
1983 if (pool == NULL)
1984 goto error;
1985 n = PyTuple_GET_SIZE(pool);
1986 if (r < 0) {
1987 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
1988 goto error;
1989 }
1990 if (r > n) {
1991 PyErr_SetString(PyExc_ValueError, "r cannot be bigger than the iterable");
1992 goto error;
1993 }
1994
1995 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
1996 if (indices == NULL) {
1997 PyErr_NoMemory();
1998 goto error;
1999 }
2000
2001 for (i=0 ; i<r ; i++)
2002 indices[i] = i;
2003
2004 /* create combinationsobject structure */
2005 co = (combinationsobject *)type->tp_alloc(type, 0);
2006 if (co == NULL)
2007 goto error;
2008
2009 co->pool = pool;
2010 co->indices = indices;
2011 co->result = NULL;
2012 co->r = r;
2013 co->stopped = 0;
2014
2015 return (PyObject *)co;
2016
2017error:
2018 if (indices != NULL)
2019 PyMem_Free(indices);
2020 Py_XDECREF(pool);
2021 return NULL;
2022}
2023
2024static void
2025combinations_dealloc(combinationsobject *co)
2026{
2027 PyObject_GC_UnTrack(co);
2028 Py_XDECREF(co->pool);
2029 Py_XDECREF(co->result);
2030 PyMem_Free(co->indices);
2031 Py_TYPE(co)->tp_free(co);
2032}
2033
2034static int
2035combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2036{
2037 Py_VISIT(co->pool);
2038 Py_VISIT(co->result);
2039 return 0;
2040}
2041
2042static PyObject *
2043combinations_next(combinationsobject *co)
2044{
2045 PyObject *elem;
2046 PyObject *oldelem;
2047 PyObject *pool = co->pool;
2048 Py_ssize_t *indices = co->indices;
2049 PyObject *result = co->result;
2050 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2051 Py_ssize_t r = co->r;
2052 Py_ssize_t i, j, index;
2053
2054 if (co->stopped)
2055 return NULL;
2056
2057 if (result == NULL) {
2058 /* On the first pass, initialize result tuple using the indices */
2059 result = PyTuple_New(r);
2060 if (result == NULL)
2061 goto empty;
2062 co->result = result;
2063 for (i=0; i<r ; i++) {
2064 index = indices[i];
2065 elem = PyTuple_GET_ITEM(pool, index);
2066 Py_INCREF(elem);
2067 PyTuple_SET_ITEM(result, i, elem);
2068 }
2069 } else {
2070 /* Copy the previous result tuple or re-use it if available */
2071 if (Py_REFCNT(result) > 1) {
2072 PyObject *old_result = result;
2073 result = PyTuple_New(r);
2074 if (result == NULL)
2075 goto empty;
2076 co->result = result;
2077 for (i=0; i<r ; i++) {
2078 elem = PyTuple_GET_ITEM(old_result, i);
2079 Py_INCREF(elem);
2080 PyTuple_SET_ITEM(result, i, elem);
2081 }
2082 Py_DECREF(old_result);
2083 }
2084 /* Now, we've got the only copy so we can update it in-place
2085 * CPython's empty tuple is a singleton and cached in
2086 * PyTuple's freelist.
2087 */
2088 assert(r == 0 || Py_REFCNT(result) == 1);
2089
2090 /* Scan indices right-to-left until finding one that is not
2091 at its maximum (i + n - r). */
2092 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2093 ;
2094
2095 /* If i is negative, then the indices are all at
2096 their maximum value and we're done. */
2097 if (i < 0)
2098 goto empty;
2099
2100 /* Increment the current index which we know is not at its
2101 maximum. Then move back to the right setting each index
2102 to its lowest possible value (one higher than the index
2103 to its left -- this maintains the sort order invariant). */
2104 indices[i]++;
2105 for (j=i+1 ; j<r ; j++)
2106 indices[j] = indices[j-1] + 1;
2107
2108 /* Update the result tuple for the new indices
2109 starting with i, the leftmost index that changed */
2110 for ( ; i<r ; i++) {
2111 index = indices[i];
2112 elem = PyTuple_GET_ITEM(pool, index);
2113 Py_INCREF(elem);
2114 oldelem = PyTuple_GET_ITEM(result, i);
2115 PyTuple_SET_ITEM(result, i, elem);
2116 Py_DECREF(oldelem);
2117 }
2118 }
2119
2120 Py_INCREF(result);
2121 return result;
2122
2123empty:
2124 co->stopped = 1;
2125 return NULL;
2126}
2127
2128PyDoc_STRVAR(combinations_doc,
2129"combinations(iterables) --> combinations object\n\
2130\n\
2131Return successive r-length combinations of elements in the iterable.\n\n\
2132combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2133
2134static PyTypeObject combinations_type = {
2135 PyVarObject_HEAD_INIT(NULL, 0)
2136 "itertools.combinations", /* tp_name */
2137 sizeof(combinationsobject), /* tp_basicsize */
2138 0, /* tp_itemsize */
2139 /* methods */
2140 (destructor)combinations_dealloc, /* tp_dealloc */
2141 0, /* tp_print */
2142 0, /* tp_getattr */
2143 0, /* tp_setattr */
2144 0, /* tp_compare */
2145 0, /* tp_repr */
2146 0, /* tp_as_number */
2147 0, /* tp_as_sequence */
2148 0, /* tp_as_mapping */
2149 0, /* tp_hash */
2150 0, /* tp_call */
2151 0, /* tp_str */
2152 PyObject_GenericGetAttr, /* tp_getattro */
2153 0, /* tp_setattro */
2154 0, /* tp_as_buffer */
2155 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2156 Py_TPFLAGS_BASETYPE, /* tp_flags */
2157 combinations_doc, /* tp_doc */
2158 (traverseproc)combinations_traverse, /* tp_traverse */
2159 0, /* tp_clear */
2160 0, /* tp_richcompare */
2161 0, /* tp_weaklistoffset */
2162 PyObject_SelfIter, /* tp_iter */
2163 (iternextfunc)combinations_next, /* tp_iternext */
2164 0, /* tp_methods */
2165 0, /* tp_members */
2166 0, /* tp_getset */
2167 0, /* tp_base */
2168 0, /* tp_dict */
2169 0, /* tp_descr_get */
2170 0, /* tp_descr_set */
2171 0, /* tp_dictoffset */
2172 0, /* tp_init */
2173 0, /* tp_alloc */
2174 combinations_new, /* tp_new */
2175 PyObject_GC_Del, /* tp_free */
2176};
2177
2178
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002179/* ifilter object ************************************************************/
2180
2181typedef struct {
2182 PyObject_HEAD
2183 PyObject *func;
2184 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002185} ifilterobject;
2186
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002187static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002188
2189static PyObject *
2190ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2191{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002192 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002193 PyObject *it;
2194 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002195
Thomas Woutersb2137042007-02-01 18:02:27 +00002196 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002197 return NULL;
2198
Raymond Hettinger60eca932003-02-09 06:40:58 +00002199 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002200 return NULL;
2201
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002202 /* Get iterator. */
2203 it = PyObject_GetIter(seq);
2204 if (it == NULL)
2205 return NULL;
2206
2207 /* create ifilterobject structure */
2208 lz = (ifilterobject *)type->tp_alloc(type, 0);
2209 if (lz == NULL) {
2210 Py_DECREF(it);
2211 return NULL;
2212 }
2213 Py_INCREF(func);
2214 lz->func = func;
2215 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002216
2217 return (PyObject *)lz;
2218}
2219
2220static void
2221ifilter_dealloc(ifilterobject *lz)
2222{
2223 PyObject_GC_UnTrack(lz);
2224 Py_XDECREF(lz->func);
2225 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00002226 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002227}
2228
2229static int
2230ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
2231{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002232 Py_VISIT(lz->it);
2233 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002234 return 0;
2235}
2236
2237static PyObject *
2238ifilter_next(ifilterobject *lz)
2239{
2240 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002241 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002242 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002243 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002244
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002245 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00002246 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002247 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002248 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002249 if (item == NULL)
2250 return NULL;
2251
Christian Heimes836baa52008-02-26 08:18:30 +00002252 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002253 ok = PyObject_IsTrue(item);
2254 } else {
2255 PyObject *good;
2256 good = PyObject_CallFunctionObjArgs(lz->func,
2257 item, NULL);
2258 if (good == NULL) {
2259 Py_DECREF(item);
2260 return NULL;
2261 }
2262 ok = PyObject_IsTrue(good);
2263 Py_DECREF(good);
2264 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00002265 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002266 return item;
2267 Py_DECREF(item);
2268 }
2269}
2270
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002271PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002272"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002273\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002274Return those items of sequence for which function(item) is true.\n\
2275If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002276
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002277static PyTypeObject ifilter_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002278 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002279 "itertools.ifilter", /* tp_name */
2280 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002281 0, /* tp_itemsize */
2282 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002283 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002284 0, /* tp_print */
2285 0, /* tp_getattr */
2286 0, /* tp_setattr */
2287 0, /* tp_compare */
2288 0, /* tp_repr */
2289 0, /* tp_as_number */
2290 0, /* tp_as_sequence */
2291 0, /* tp_as_mapping */
2292 0, /* tp_hash */
2293 0, /* tp_call */
2294 0, /* tp_str */
2295 PyObject_GenericGetAttr, /* tp_getattro */
2296 0, /* tp_setattro */
2297 0, /* tp_as_buffer */
2298 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2299 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002300 ifilter_doc, /* tp_doc */
2301 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002302 0, /* tp_clear */
2303 0, /* tp_richcompare */
2304 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002305 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002306 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002307 0, /* tp_methods */
2308 0, /* tp_members */
2309 0, /* tp_getset */
2310 0, /* tp_base */
2311 0, /* tp_dict */
2312 0, /* tp_descr_get */
2313 0, /* tp_descr_set */
2314 0, /* tp_dictoffset */
2315 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002316 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002317 ifilter_new, /* tp_new */
2318 PyObject_GC_Del, /* tp_free */
2319};
2320
2321
2322/* ifilterfalse object ************************************************************/
2323
2324typedef struct {
2325 PyObject_HEAD
2326 PyObject *func;
2327 PyObject *it;
2328} ifilterfalseobject;
2329
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002330static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002331
2332static PyObject *
2333ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2334{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00002335 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002336 PyObject *it;
2337 ifilterfalseobject *lz;
2338
Thomas Woutersb2137042007-02-01 18:02:27 +00002339 if (type == &ifilterfalse_type &&
2340 !_PyArg_NoKeywords("ifilterfalse()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002341 return NULL;
2342
Raymond Hettinger60eca932003-02-09 06:40:58 +00002343 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
2344 return NULL;
2345
2346 /* Get iterator. */
2347 it = PyObject_GetIter(seq);
2348 if (it == NULL)
2349 return NULL;
2350
2351 /* create ifilterfalseobject structure */
2352 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
2353 if (lz == NULL) {
2354 Py_DECREF(it);
2355 return NULL;
2356 }
2357 Py_INCREF(func);
2358 lz->func = func;
2359 lz->it = it;
2360
2361 return (PyObject *)lz;
2362}
2363
2364static void
2365ifilterfalse_dealloc(ifilterfalseobject *lz)
2366{
2367 PyObject_GC_UnTrack(lz);
2368 Py_XDECREF(lz->func);
2369 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00002370 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002371}
2372
2373static int
2374ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
2375{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002376 Py_VISIT(lz->it);
2377 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002378 return 0;
2379}
2380
2381static PyObject *
2382ifilterfalse_next(ifilterfalseobject *lz)
2383{
2384 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002385 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002386 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002387 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002388
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002389 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00002390 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002391 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002392 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002393 if (item == NULL)
2394 return NULL;
2395
Christian Heimes836baa52008-02-26 08:18:30 +00002396 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
Raymond Hettinger60eca932003-02-09 06:40:58 +00002397 ok = PyObject_IsTrue(item);
2398 } else {
2399 PyObject *good;
2400 good = PyObject_CallFunctionObjArgs(lz->func,
2401 item, NULL);
2402 if (good == NULL) {
2403 Py_DECREF(item);
2404 return NULL;
2405 }
2406 ok = PyObject_IsTrue(good);
2407 Py_DECREF(good);
2408 }
2409 if (!ok)
2410 return item;
2411 Py_DECREF(item);
2412 }
2413}
2414
Raymond Hettinger60eca932003-02-09 06:40:58 +00002415PyDoc_STRVAR(ifilterfalse_doc,
2416"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
2417\n\
2418Return those items of sequence for which function(item) is false.\n\
2419If function is None, return the items that are false.");
2420
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002421static PyTypeObject ifilterfalse_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002422 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002423 "itertools.ifilterfalse", /* tp_name */
2424 sizeof(ifilterfalseobject), /* tp_basicsize */
2425 0, /* tp_itemsize */
2426 /* methods */
2427 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
2428 0, /* tp_print */
2429 0, /* tp_getattr */
2430 0, /* tp_setattr */
2431 0, /* tp_compare */
2432 0, /* tp_repr */
2433 0, /* tp_as_number */
2434 0, /* tp_as_sequence */
2435 0, /* tp_as_mapping */
2436 0, /* tp_hash */
2437 0, /* tp_call */
2438 0, /* tp_str */
2439 PyObject_GenericGetAttr, /* tp_getattro */
2440 0, /* tp_setattro */
2441 0, /* tp_as_buffer */
2442 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2443 Py_TPFLAGS_BASETYPE, /* tp_flags */
2444 ifilterfalse_doc, /* tp_doc */
2445 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
2446 0, /* tp_clear */
2447 0, /* tp_richcompare */
2448 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002449 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002450 (iternextfunc)ifilterfalse_next, /* tp_iternext */
2451 0, /* tp_methods */
2452 0, /* tp_members */
2453 0, /* tp_getset */
2454 0, /* tp_base */
2455 0, /* tp_dict */
2456 0, /* tp_descr_get */
2457 0, /* tp_descr_set */
2458 0, /* tp_dictoffset */
2459 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002460 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002461 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002462 PyObject_GC_Del, /* tp_free */
2463};
2464
2465
2466/* count object ************************************************************/
2467
2468typedef struct {
2469 PyObject_HEAD
Thomas Wouters477c8d52006-05-27 19:21:47 +00002470 Py_ssize_t cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002471 PyObject *long_cnt; /* Arbitrarily large count when cnt >= PY_SSIZE_T_MAX */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002472} countobject;
2473
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002474static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002475
2476static PyObject *
2477count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2478{
2479 countobject *lz;
Thomas Wouters477c8d52006-05-27 19:21:47 +00002480 Py_ssize_t cnt = 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002481 PyObject *cnt_arg = NULL;
2482 PyObject *long_cnt = NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002483
Thomas Woutersb2137042007-02-01 18:02:27 +00002484 if (type == &count_type && !_PyArg_NoKeywords("count()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002485 return NULL;
2486
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002487 if (!PyArg_UnpackTuple(args, "count", 0, 1, &cnt_arg))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002488 return NULL;
2489
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002490 if (cnt_arg != NULL) {
Christian Heimes217cfd12007-12-02 14:31:20 +00002491 cnt = PyLong_AsSsize_t(cnt_arg);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002492 if (cnt == -1 && PyErr_Occurred()) {
2493 PyErr_Clear();
2494 if (!PyLong_Check(cnt_arg)) {
2495 PyErr_SetString(PyExc_TypeError, "an integer is required");
2496 return NULL;
2497 }
2498 long_cnt = cnt_arg;
2499 Py_INCREF(long_cnt);
2500 cnt = PY_SSIZE_T_MAX;
2501 }
2502 }
2503
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002504 /* create countobject structure */
2505 lz = (countobject *)PyObject_New(countobject, &count_type);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002506 if (lz == NULL) {
2507 Py_XDECREF(long_cnt);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002508 return NULL;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002509 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002510 lz->cnt = cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002511 lz->long_cnt = long_cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002512
2513 return (PyObject *)lz;
2514}
2515
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002516static void
2517count_dealloc(countobject *lz)
2518{
2519 Py_XDECREF(lz->long_cnt);
2520 PyObject_Del(lz);
2521}
2522
2523static PyObject *
2524count_nextlong(countobject *lz)
2525{
2526 static PyObject *one = NULL;
2527 PyObject *cnt;
2528 PyObject *stepped_up;
2529
2530 if (lz->long_cnt == NULL) {
Christian Heimes217cfd12007-12-02 14:31:20 +00002531 lz->long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002532 if (lz->long_cnt == NULL)
2533 return NULL;
2534 }
2535 if (one == NULL) {
Christian Heimes217cfd12007-12-02 14:31:20 +00002536 one = PyLong_FromLong(1);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002537 if (one == NULL)
2538 return NULL;
2539 }
2540 cnt = lz->long_cnt;
2541 assert(cnt != NULL);
2542 stepped_up = PyNumber_Add(cnt, one);
2543 if (stepped_up == NULL)
2544 return NULL;
2545 lz->long_cnt = stepped_up;
2546 return cnt;
2547}
2548
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002549static PyObject *
2550count_next(countobject *lz)
2551{
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002552 if (lz->cnt == PY_SSIZE_T_MAX)
2553 return count_nextlong(lz);
Christian Heimes217cfd12007-12-02 14:31:20 +00002554 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002555}
2556
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002557static PyObject *
2558count_repr(countobject *lz)
2559{
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002560 if (lz->cnt != PY_SSIZE_T_MAX)
2561 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
2562
2563 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002564}
2565
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002566PyDoc_STRVAR(count_doc,
2567"count([firstval]) --> count object\n\
2568\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002569Return a count object whose .__next__() method returns consecutive\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002570integers starting from zero or, if specified, from firstval.");
2571
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002572static PyTypeObject count_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002573 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002574 "itertools.count", /* tp_name */
2575 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002576 0, /* tp_itemsize */
2577 /* methods */
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002578 (destructor)count_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002579 0, /* tp_print */
2580 0, /* tp_getattr */
2581 0, /* tp_setattr */
2582 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002583 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002584 0, /* tp_as_number */
2585 0, /* tp_as_sequence */
2586 0, /* tp_as_mapping */
2587 0, /* tp_hash */
2588 0, /* tp_call */
2589 0, /* tp_str */
2590 PyObject_GenericGetAttr, /* tp_getattro */
2591 0, /* tp_setattro */
2592 0, /* tp_as_buffer */
2593 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002594 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002595 0, /* tp_traverse */
2596 0, /* tp_clear */
2597 0, /* tp_richcompare */
2598 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002599 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002600 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002601 0, /* tp_methods */
2602 0, /* tp_members */
2603 0, /* tp_getset */
2604 0, /* tp_base */
2605 0, /* tp_dict */
2606 0, /* tp_descr_get */
2607 0, /* tp_descr_set */
2608 0, /* tp_dictoffset */
2609 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002610 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002611 count_new, /* tp_new */
2612};
2613
2614
2615/* izip object ************************************************************/
2616
2617#include "Python.h"
2618
2619typedef struct {
2620 PyObject_HEAD
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002621 Py_ssize_t tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002622 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002623 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002624} izipobject;
2625
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002626static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002627
2628static PyObject *
2629izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2630{
2631 izipobject *lz;
Thomas Wouters477c8d52006-05-27 19:21:47 +00002632 Py_ssize_t i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002633 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002634 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002635 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002636
Thomas Woutersb2137042007-02-01 18:02:27 +00002637 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002638 return NULL;
2639
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002640 /* args must be a tuple */
2641 assert(PyTuple_Check(args));
2642
2643 /* obtain iterators */
2644 ittuple = PyTuple_New(tuplesize);
Thomas Wouters477c8d52006-05-27 19:21:47 +00002645 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002646 return NULL;
2647 for (i=0; i < tuplesize; ++i) {
2648 PyObject *item = PyTuple_GET_ITEM(args, i);
2649 PyObject *it = PyObject_GetIter(item);
2650 if (it == NULL) {
2651 if (PyErr_ExceptionMatches(PyExc_TypeError))
2652 PyErr_Format(PyExc_TypeError,
Thomas Wouters477c8d52006-05-27 19:21:47 +00002653 "izip argument #%zd must support iteration",
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002654 i+1);
2655 Py_DECREF(ittuple);
2656 return NULL;
2657 }
2658 PyTuple_SET_ITEM(ittuple, i, it);
2659 }
2660
Raymond Hettinger2012f172003-02-07 05:32:58 +00002661 /* create a result holder */
2662 result = PyTuple_New(tuplesize);
2663 if (result == NULL) {
2664 Py_DECREF(ittuple);
2665 return NULL;
2666 }
2667 for (i=0 ; i < tuplesize ; i++) {
2668 Py_INCREF(Py_None);
2669 PyTuple_SET_ITEM(result, i, Py_None);
2670 }
2671
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002672 /* create izipobject structure */
2673 lz = (izipobject *)type->tp_alloc(type, 0);
2674 if (lz == NULL) {
2675 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002676 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002677 return NULL;
2678 }
2679 lz->ittuple = ittuple;
2680 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002681 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002682
2683 return (PyObject *)lz;
2684}
2685
2686static void
2687izip_dealloc(izipobject *lz)
2688{
2689 PyObject_GC_UnTrack(lz);
2690 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002691 Py_XDECREF(lz->result);
Christian Heimes90aa7642007-12-19 02:45:37 +00002692 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002693}
2694
2695static int
2696izip_traverse(izipobject *lz, visitproc visit, void *arg)
2697{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002698 Py_VISIT(lz->ittuple);
2699 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002700 return 0;
2701}
2702
2703static PyObject *
2704izip_next(izipobject *lz)
2705{
Thomas Wouters477c8d52006-05-27 19:21:47 +00002706 Py_ssize_t i;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002707 Py_ssize_t tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002708 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002709 PyObject *it;
2710 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002711 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002712
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002713 if (tuplesize == 0)
2714 return NULL;
Christian Heimes90aa7642007-12-19 02:45:37 +00002715 if (Py_REFCNT(result) == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002716 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002717 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002718 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002719 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00002720 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002721 if (item == NULL) {
2722 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002723 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002724 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002725 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002726 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002727 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002728 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002729 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002730 result = PyTuple_New(tuplesize);
2731 if (result == NULL)
2732 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002733 for (i=0 ; i < tuplesize ; i++) {
2734 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002735 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00002736 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002737 if (item == NULL) {
2738 Py_DECREF(result);
2739 return NULL;
2740 }
2741 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002742 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002743 }
2744 return result;
2745}
2746
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002747PyDoc_STRVAR(izip_doc,
2748"izip(iter1 [,iter2 [...]]) --> izip object\n\
2749\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002750Return a izip object whose .__next__() method returns a tuple where\n\
2751the i-th element comes from the i-th iterable argument. The .__next__()\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002752method continues until the shortest iterable in the argument sequence\n\
2753is exhausted and then it raises StopIteration. Works like the zip()\n\
2754function but consumes less memory by returning an iterator instead of\n\
2755a list.");
2756
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002757static PyTypeObject izip_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002758 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002759 "itertools.izip", /* tp_name */
2760 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002761 0, /* tp_itemsize */
2762 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002763 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002764 0, /* tp_print */
2765 0, /* tp_getattr */
2766 0, /* tp_setattr */
2767 0, /* tp_compare */
2768 0, /* tp_repr */
2769 0, /* tp_as_number */
2770 0, /* tp_as_sequence */
2771 0, /* tp_as_mapping */
2772 0, /* tp_hash */
2773 0, /* tp_call */
2774 0, /* tp_str */
2775 PyObject_GenericGetAttr, /* tp_getattro */
2776 0, /* tp_setattro */
2777 0, /* tp_as_buffer */
2778 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2779 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002780 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002781 (traverseproc)izip_traverse, /* tp_traverse */
2782 0, /* tp_clear */
2783 0, /* tp_richcompare */
2784 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002785 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002786 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002787 0, /* tp_methods */
2788 0, /* tp_members */
2789 0, /* tp_getset */
2790 0, /* tp_base */
2791 0, /* tp_dict */
2792 0, /* tp_descr_get */
2793 0, /* tp_descr_set */
2794 0, /* tp_dictoffset */
2795 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002796 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002797 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002798 PyObject_GC_Del, /* tp_free */
2799};
2800
2801
2802/* repeat object ************************************************************/
2803
2804typedef struct {
2805 PyObject_HEAD
2806 PyObject *element;
Thomas Wouters477c8d52006-05-27 19:21:47 +00002807 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002808} repeatobject;
2809
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002810static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002811
2812static PyObject *
2813repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2814{
2815 repeatobject *ro;
2816 PyObject *element;
Thomas Wouters477c8d52006-05-27 19:21:47 +00002817 Py_ssize_t cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002818
Thomas Woutersb2137042007-02-01 18:02:27 +00002819 if (type == &repeat_type && !_PyArg_NoKeywords("repeat()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002820 return NULL;
2821
Thomas Wouters477c8d52006-05-27 19:21:47 +00002822 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002823 return NULL;
2824
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002825 if (PyTuple_Size(args) == 2 && cnt < 0)
2826 cnt = 0;
2827
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002828 ro = (repeatobject *)type->tp_alloc(type, 0);
2829 if (ro == NULL)
2830 return NULL;
2831 Py_INCREF(element);
2832 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002833 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002834 return (PyObject *)ro;
2835}
2836
2837static void
2838repeat_dealloc(repeatobject *ro)
2839{
2840 PyObject_GC_UnTrack(ro);
2841 Py_XDECREF(ro->element);
Christian Heimes90aa7642007-12-19 02:45:37 +00002842 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002843}
2844
2845static int
2846repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2847{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002848 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002849 return 0;
2850}
2851
2852static PyObject *
2853repeat_next(repeatobject *ro)
2854{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002855 if (ro->cnt == 0)
2856 return NULL;
2857 if (ro->cnt > 0)
2858 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002859 Py_INCREF(ro->element);
2860 return ro->element;
2861}
2862
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002863static PyObject *
2864repeat_repr(repeatobject *ro)
2865{
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002866 if (ro->cnt == -1)
Walter Dörwald7569dfe2007-05-19 21:49:49 +00002867 return PyUnicode_FromFormat("repeat(%R)", ro->element);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002868 else
Walter Dörwald7569dfe2007-05-19 21:49:49 +00002869 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002870}
2871
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002872static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002873repeat_len(repeatobject *ro)
2874{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002875 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002876 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002877 return NULL;
2878 }
Christian Heimes217cfd12007-12-02 14:31:20 +00002879 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002880}
2881
Armin Rigof5b3e362006-02-11 21:32:43 +00002882PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002883
2884static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00002885 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002886 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002887};
2888
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002889PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002890"repeat(element [,times]) -> create an iterator which returns the element\n\
2891for the specified number of times. If not specified, returns the element\n\
2892endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002893
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002894static PyTypeObject repeat_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002895 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002896 "itertools.repeat", /* tp_name */
2897 sizeof(repeatobject), /* tp_basicsize */
2898 0, /* tp_itemsize */
2899 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002900 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002901 0, /* tp_print */
2902 0, /* tp_getattr */
2903 0, /* tp_setattr */
2904 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002905 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002906 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002907 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002908 0, /* tp_as_mapping */
2909 0, /* tp_hash */
2910 0, /* tp_call */
2911 0, /* tp_str */
2912 PyObject_GenericGetAttr, /* tp_getattro */
2913 0, /* tp_setattro */
2914 0, /* tp_as_buffer */
2915 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2916 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002917 repeat_doc, /* tp_doc */
2918 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002919 0, /* tp_clear */
2920 0, /* tp_richcompare */
2921 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002922 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002923 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002924 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002925 0, /* tp_members */
2926 0, /* tp_getset */
2927 0, /* tp_base */
2928 0, /* tp_dict */
2929 0, /* tp_descr_get */
2930 0, /* tp_descr_set */
2931 0, /* tp_dictoffset */
2932 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002933 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002934 repeat_new, /* tp_new */
2935 PyObject_GC_Del, /* tp_free */
2936};
2937
Thomas Wouterscf297e42007-02-23 15:07:44 +00002938/* iziplongest object ************************************************************/
2939
2940#include "Python.h"
2941
2942typedef struct {
2943 PyObject_HEAD
2944 Py_ssize_t tuplesize;
2945 Py_ssize_t numactive;
2946 PyObject *ittuple; /* tuple of iterators */
2947 PyObject *result;
2948 PyObject *fillvalue;
2949} iziplongestobject;
2950
2951static PyTypeObject iziplongest_type;
2952
2953static PyObject *
2954izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2955{
2956 iziplongestobject *lz;
2957 Py_ssize_t i;
2958 PyObject *ittuple; /* tuple of iterators */
2959 PyObject *result;
2960 PyObject *fillvalue = Py_None;
2961 Py_ssize_t tuplesize = PySequence_Length(args);
2962
2963 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
2964 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
2965 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
2966 PyErr_SetString(PyExc_TypeError,
2967 "izip_longest() got an unexpected keyword argument");
2968 return NULL;
2969 }
2970 }
2971
2972 /* args must be a tuple */
2973 assert(PyTuple_Check(args));
2974
2975 /* obtain iterators */
2976 ittuple = PyTuple_New(tuplesize);
2977 if (ittuple == NULL)
2978 return NULL;
2979 for (i=0; i < tuplesize; ++i) {
2980 PyObject *item = PyTuple_GET_ITEM(args, i);
2981 PyObject *it = PyObject_GetIter(item);
2982 if (it == NULL) {
2983 if (PyErr_ExceptionMatches(PyExc_TypeError))
2984 PyErr_Format(PyExc_TypeError,
2985 "izip_longest argument #%zd must support iteration",
2986 i+1);
2987 Py_DECREF(ittuple);
2988 return NULL;
2989 }
2990 PyTuple_SET_ITEM(ittuple, i, it);
2991 }
2992
2993 /* create a result holder */
2994 result = PyTuple_New(tuplesize);
2995 if (result == NULL) {
2996 Py_DECREF(ittuple);
2997 return NULL;
2998 }
2999 for (i=0 ; i < tuplesize ; i++) {
3000 Py_INCREF(Py_None);
3001 PyTuple_SET_ITEM(result, i, Py_None);
3002 }
3003
3004 /* create iziplongestobject structure */
3005 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3006 if (lz == NULL) {
3007 Py_DECREF(ittuple);
3008 Py_DECREF(result);
3009 return NULL;
3010 }
3011 lz->ittuple = ittuple;
3012 lz->tuplesize = tuplesize;
3013 lz->numactive = tuplesize;
3014 lz->result = result;
3015 Py_INCREF(fillvalue);
3016 lz->fillvalue = fillvalue;
3017 return (PyObject *)lz;
3018}
3019
3020static void
3021izip_longest_dealloc(iziplongestobject *lz)
3022{
3023 PyObject_GC_UnTrack(lz);
3024 Py_XDECREF(lz->ittuple);
3025 Py_XDECREF(lz->result);
3026 Py_XDECREF(lz->fillvalue);
Christian Heimes90aa7642007-12-19 02:45:37 +00003027 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003028}
3029
3030static int
3031izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3032{
3033 Py_VISIT(lz->ittuple);
3034 Py_VISIT(lz->result);
3035 Py_VISIT(lz->fillvalue);
3036 return 0;
3037}
3038
3039static PyObject *
3040izip_longest_next(iziplongestobject *lz)
3041{
3042 Py_ssize_t i;
3043 Py_ssize_t tuplesize = lz->tuplesize;
3044 PyObject *result = lz->result;
3045 PyObject *it;
3046 PyObject *item;
3047 PyObject *olditem;
3048
3049 if (tuplesize == 0)
3050 return NULL;
3051 if (lz->numactive == 0)
3052 return NULL;
Christian Heimes90aa7642007-12-19 02:45:37 +00003053 if (Py_REFCNT(result) == 1) {
Thomas Wouterscf297e42007-02-23 15:07:44 +00003054 Py_INCREF(result);
3055 for (i=0 ; i < tuplesize ; i++) {
3056 it = PyTuple_GET_ITEM(lz->ittuple, i);
3057 if (it == NULL) {
3058 Py_INCREF(lz->fillvalue);
3059 item = lz->fillvalue;
3060 } else {
3061 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00003062 item = (*Py_TYPE(it)->tp_iternext)(it);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003063 if (item == NULL) {
3064 lz->numactive -= 1;
3065 if (lz->numactive == 0) {
3066 Py_DECREF(result);
3067 return NULL;
3068 } else {
3069 Py_INCREF(lz->fillvalue);
3070 item = lz->fillvalue;
3071 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3072 Py_DECREF(it);
3073 }
3074 }
3075 }
3076 olditem = PyTuple_GET_ITEM(result, i);
3077 PyTuple_SET_ITEM(result, i, item);
3078 Py_DECREF(olditem);
3079 }
3080 } else {
3081 result = PyTuple_New(tuplesize);
3082 if (result == NULL)
3083 return NULL;
3084 for (i=0 ; i < tuplesize ; i++) {
3085 it = PyTuple_GET_ITEM(lz->ittuple, i);
3086 if (it == NULL) {
3087 Py_INCREF(lz->fillvalue);
3088 item = lz->fillvalue;
3089 } else {
3090 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00003091 item = (*Py_TYPE(it)->tp_iternext)(it);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003092 if (item == NULL) {
3093 lz->numactive -= 1;
3094 if (lz->numactive == 0) {
3095 Py_DECREF(result);
3096 return NULL;
3097 } else {
3098 Py_INCREF(lz->fillvalue);
3099 item = lz->fillvalue;
3100 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3101 Py_DECREF(it);
3102 }
3103 }
3104 }
3105 PyTuple_SET_ITEM(result, i, item);
3106 }
3107 }
3108 return result;
3109}
3110
3111PyDoc_STRVAR(izip_longest_doc,
3112"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3113\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00003114Return an izip_longest object whose .__next__() method returns a tuple where\n\
3115the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003116method continues until the longest iterable in the argument sequence\n\
3117is exhausted and then it raises StopIteration. When the shorter iterables\n\
3118are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3119defaults to None or can be specified by a keyword argument.\n\
3120");
3121
3122static PyTypeObject iziplongest_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00003123 PyVarObject_HEAD_INIT(NULL, 0)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003124 "itertools.izip_longest", /* tp_name */
3125 sizeof(iziplongestobject), /* tp_basicsize */
3126 0, /* tp_itemsize */
3127 /* methods */
3128 (destructor)izip_longest_dealloc, /* tp_dealloc */
3129 0, /* tp_print */
3130 0, /* tp_getattr */
3131 0, /* tp_setattr */
3132 0, /* tp_compare */
3133 0, /* tp_repr */
3134 0, /* tp_as_number */
3135 0, /* tp_as_sequence */
3136 0, /* tp_as_mapping */
3137 0, /* tp_hash */
3138 0, /* tp_call */
3139 0, /* tp_str */
3140 PyObject_GenericGetAttr, /* tp_getattro */
3141 0, /* tp_setattro */
3142 0, /* tp_as_buffer */
3143 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3144 Py_TPFLAGS_BASETYPE, /* tp_flags */
3145 izip_longest_doc, /* tp_doc */
3146 (traverseproc)izip_longest_traverse, /* tp_traverse */
3147 0, /* tp_clear */
3148 0, /* tp_richcompare */
3149 0, /* tp_weaklistoffset */
3150 PyObject_SelfIter, /* tp_iter */
3151 (iternextfunc)izip_longest_next, /* tp_iternext */
3152 0, /* tp_methods */
3153 0, /* tp_members */
3154 0, /* tp_getset */
3155 0, /* tp_base */
3156 0, /* tp_dict */
3157 0, /* tp_descr_get */
3158 0, /* tp_descr_set */
3159 0, /* tp_dictoffset */
3160 0, /* tp_init */
3161 0, /* tp_alloc */
3162 izip_longest_new, /* tp_new */
3163 PyObject_GC_Del, /* tp_free */
3164};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003165
3166/* module level code ********************************************************/
3167
3168PyDoc_STRVAR(module_doc,
3169"Functional tools for creating and using iterators.\n\
3170\n\
3171Infinite iterators:\n\
3172count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003173cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003174repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003175\n\
3176Iterators terminating on the shortest input sequence:\n\
3177izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003178izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003179ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
3180ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003181islice(seq, [start,] stop [, step]) --> elements from\n\
3182 seq[start:stop:step]\n\
3183imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
3184starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003185tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003186chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003187takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
3188dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003189groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003190");
3191
3192
Raymond Hettingerad983e72003-11-12 14:32:26 +00003193static PyMethodDef module_methods[] = {
3194 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3195 {NULL, NULL} /* sentinel */
3196};
3197
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003198PyMODINIT_FUNC
3199inititertools(void)
3200{
Raymond Hettinger60eca932003-02-09 06:40:58 +00003201 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003202 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003203 char *name;
3204 PyTypeObject *typelist[] = {
Christian Heimes380f7f22008-02-28 11:19:05 +00003205 &combinations_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003206 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003207 &dropwhile_type,
3208 &takewhile_type,
3209 &islice_type,
3210 &starmap_type,
3211 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003212 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003213 &ifilter_type,
3214 &ifilterfalse_type,
3215 &count_type,
3216 &izip_type,
Christian Heimesc3f30c42008-02-22 16:37:40 +00003217 &iziplongest_type,
Christian Heimes380f7f22008-02-28 11:19:05 +00003218 &product_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003219 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003220 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003221 NULL
3222 };
3223
Christian Heimes90aa7642007-12-19 02:45:37 +00003224 Py_TYPE(&teedataobject_type) = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00003225 m = Py_InitModule3("itertools", module_methods, module_doc);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00003226 if (m == NULL)
3227 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003228
Raymond Hettinger60eca932003-02-09 06:40:58 +00003229 for (i=0 ; typelist[i] != NULL ; i++) {
3230 if (PyType_Ready(typelist[i]) < 0)
3231 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003232 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003233 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003234 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003235 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003236 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003237
3238 if (PyType_Ready(&teedataobject_type) < 0)
3239 return;
3240 if (PyType_Ready(&tee_type) < 0)
3241 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003242 if (PyType_Ready(&_grouper_type) < 0)
3243 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003244}