blob: 12e862d5f3c6527b870939a772a275157db0c8cc [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
Christian Heimesf16baeb2008-02-29 14:57:44 +00001574 PyObject *source; /* Iterator over input iterables */
1575 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001576} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001577
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001578static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001579
Christian Heimesf16baeb2008-02-29 14:57:44 +00001580static PyObject *
1581chain_new_internal(PyTypeObject *type, PyObject *source)
1582{
1583 chainobject *lz;
1584
1585 lz = (chainobject *)type->tp_alloc(type, 0);
1586 if (lz == NULL) {
1587 Py_DECREF(source);
1588 return NULL;
1589 }
1590
1591 lz->source = source;
1592 lz->active = NULL;
1593 return (PyObject *)lz;
1594}
1595
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001596static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001597chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001598{
Christian Heimesf16baeb2008-02-29 14:57:44 +00001599 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001600
Thomas Woutersb2137042007-02-01 18:02:27 +00001601 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001602 return NULL;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001603
1604 source = PyObject_GetIter(args);
1605 if (source == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001606 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001607
Christian Heimesf16baeb2008-02-29 14:57:44 +00001608 return chain_new_internal(type, source);
1609}
1610
1611static PyObject *
1612chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1613{
1614 PyObject *source;
1615
1616 source = PyObject_GetIter(arg);
1617 if (source == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001618 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001619
Christian Heimesf16baeb2008-02-29 14:57:44 +00001620 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001621}
1622
1623static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001624chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001625{
1626 PyObject_GC_UnTrack(lz);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001627 Py_XDECREF(lz->active);
1628 Py_XDECREF(lz->source);
Christian Heimes90aa7642007-12-19 02:45:37 +00001629 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001630}
1631
Raymond Hettinger2012f172003-02-07 05:32:58 +00001632static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001633chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001634{
Christian Heimesf16baeb2008-02-29 14:57:44 +00001635 Py_VISIT(lz->source);
1636 Py_VISIT(lz->active);
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 *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001644
Christian Heimesf16baeb2008-02-29 14:57:44 +00001645 if (lz->source == NULL)
1646 return NULL; /* already stopped */
1647
1648 if (lz->active == NULL) {
1649 PyObject *iterable = PyIter_Next(lz->source);
1650 if (iterable == NULL) {
1651 Py_CLEAR(lz->source);
1652 return NULL; /* no more input sources */
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001653 }
Christian Heimesf16baeb2008-02-29 14:57:44 +00001654 lz->active = PyObject_GetIter(iterable);
1655 if (lz->active == NULL) {
1656 Py_DECREF(iterable);
1657 Py_CLEAR(lz->source);
1658 return NULL; /* input not iterable */
1659 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001660 }
Christian Heimesf16baeb2008-02-29 14:57:44 +00001661 item = PyIter_Next(lz->active);
1662 if (item != NULL)
1663 return item;
1664 if (PyErr_Occurred()) {
1665 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1666 PyErr_Clear();
1667 else
1668 return NULL; /* input raised an exception */
1669 }
1670 Py_CLEAR(lz->active);
1671 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001672}
1673
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001674PyDoc_STRVAR(chain_doc,
1675"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001676\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001677Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001678first iterable until it is exhausted, then elements from the next\n\
1679iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001680
Christian Heimesf16baeb2008-02-29 14:57:44 +00001681PyDoc_STRVAR(chain_from_iterable_doc,
1682"chain.from_iterable(iterable) --> chain object\n\
1683\n\
1684Alternate chain() contructor taking a single iterable argument\n\
1685that evaluates lazily.");
1686
1687static PyMethodDef chain_methods[] = {
1688 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1689 chain_from_iterable_doc},
1690 {NULL, NULL} /* sentinel */
1691};
1692
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001693static PyTypeObject chain_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001694 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001695 "itertools.chain", /* tp_name */
1696 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001697 0, /* tp_itemsize */
1698 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001699 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001700 0, /* tp_print */
1701 0, /* tp_getattr */
1702 0, /* tp_setattr */
1703 0, /* tp_compare */
1704 0, /* tp_repr */
1705 0, /* tp_as_number */
1706 0, /* tp_as_sequence */
1707 0, /* tp_as_mapping */
1708 0, /* tp_hash */
1709 0, /* tp_call */
1710 0, /* tp_str */
1711 PyObject_GenericGetAttr, /* tp_getattro */
1712 0, /* tp_setattro */
1713 0, /* tp_as_buffer */
1714 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1715 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001716 chain_doc, /* tp_doc */
1717 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001718 0, /* tp_clear */
1719 0, /* tp_richcompare */
1720 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001721 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001722 (iternextfunc)chain_next, /* tp_iternext */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001723 chain_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001724 0, /* tp_members */
1725 0, /* tp_getset */
1726 0, /* tp_base */
1727 0, /* tp_dict */
1728 0, /* tp_descr_get */
1729 0, /* tp_descr_set */
1730 0, /* tp_dictoffset */
1731 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001732 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001733 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001734 PyObject_GC_Del, /* tp_free */
1735};
1736
1737
Christian Heimesc3f30c42008-02-22 16:37:40 +00001738/* product object ************************************************************/
1739
1740typedef struct {
1741 PyObject_HEAD
1742 PyObject *pools; /* tuple of pool tuples */
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001743 Py_ssize_t *maxvec; /* size of each pool */
1744 Py_ssize_t *indices; /* one index per pool */
1745 PyObject *result; /* most recently returned result tuple */
1746 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001747} productobject;
1748
1749static PyTypeObject product_type;
1750
1751static PyObject *
1752product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1753{
1754 productobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001755 Py_ssize_t nargs, npools, repeat=1;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001756 PyObject *pools = NULL;
1757 Py_ssize_t *maxvec = NULL;
1758 Py_ssize_t *indices = NULL;
1759 Py_ssize_t i;
1760
Christian Heimesf16baeb2008-02-29 14:57:44 +00001761 if (kwds != NULL) {
1762 char *kwlist[] = {"repeat", 0};
1763 PyObject *tmpargs = PyTuple_New(0);
1764 if (tmpargs == NULL)
1765 return NULL;
1766 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1767 Py_DECREF(tmpargs);
1768 return NULL;
1769 }
1770 Py_DECREF(tmpargs);
1771 if (repeat < 0) {
1772 PyErr_SetString(PyExc_ValueError,
1773 "repeat argument cannot be negative");
1774 return NULL;
1775 }
1776 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001777
1778 assert(PyTuple_Check(args));
Christian Heimesf16baeb2008-02-29 14:57:44 +00001779 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1780 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001781
1782 maxvec = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1783 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1784 if (maxvec == NULL || indices == NULL) {
1785 PyErr_NoMemory();
1786 goto error;
1787 }
1788
1789 pools = PyTuple_New(npools);
1790 if (pools == NULL)
1791 goto error;
1792
Christian Heimesf16baeb2008-02-29 14:57:44 +00001793 for (i=0; i < nargs ; ++i) {
Christian Heimesc3f30c42008-02-22 16:37:40 +00001794 PyObject *item = PyTuple_GET_ITEM(args, i);
1795 PyObject *pool = PySequence_Tuple(item);
1796 if (pool == NULL)
1797 goto error;
1798
1799 PyTuple_SET_ITEM(pools, i, pool);
1800 maxvec[i] = PyTuple_GET_SIZE(pool);
1801 indices[i] = 0;
1802 }
Christian Heimesf16baeb2008-02-29 14:57:44 +00001803 for ( ; i < npools; ++i) {
1804 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1805 Py_INCREF(pool);
1806 PyTuple_SET_ITEM(pools, i, pool);
1807 maxvec[i] = maxvec[i - nargs];
1808 indices[i] = 0;
1809 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001810
1811 /* create productobject structure */
1812 lz = (productobject *)type->tp_alloc(type, 0);
Christian Heimes380f7f22008-02-28 11:19:05 +00001813 if (lz == NULL)
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001814 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001815
1816 lz->pools = pools;
1817 lz->maxvec = maxvec;
1818 lz->indices = indices;
1819 lz->result = NULL;
1820 lz->stopped = 0;
1821
1822 return (PyObject *)lz;
1823
1824error:
1825 if (maxvec != NULL)
1826 PyMem_Free(maxvec);
1827 if (indices != NULL)
1828 PyMem_Free(indices);
1829 Py_XDECREF(pools);
1830 return NULL;
1831}
1832
1833static void
1834product_dealloc(productobject *lz)
1835{
1836 PyObject_GC_UnTrack(lz);
1837 Py_XDECREF(lz->pools);
1838 Py_XDECREF(lz->result);
1839 PyMem_Free(lz->maxvec);
1840 PyMem_Free(lz->indices);
1841 Py_TYPE(lz)->tp_free(lz);
1842}
1843
1844static int
1845product_traverse(productobject *lz, visitproc visit, void *arg)
1846{
1847 Py_VISIT(lz->pools);
1848 Py_VISIT(lz->result);
1849 return 0;
1850}
1851
1852static PyObject *
1853product_next(productobject *lz)
1854{
1855 PyObject *pool;
1856 PyObject *elem;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001857 PyObject *oldelem;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001858 PyObject *pools = lz->pools;
1859 PyObject *result = lz->result;
1860 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1861 Py_ssize_t i;
1862
1863 if (lz->stopped)
1864 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001865
Christian Heimesc3f30c42008-02-22 16:37:40 +00001866 if (result == NULL) {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001867 /* On the first pass, return an initial tuple filled with the
1868 first element from each pool. If any pool is empty, then
1869 whole product is empty and we're already done */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001870 if (npools == 0)
1871 goto empty;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001872 result = PyTuple_New(npools);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001873 if (result == NULL)
1874 goto empty;
1875 lz->result = result;
1876 for (i=0; i < npools; i++) {
1877 pool = PyTuple_GET_ITEM(pools, i);
1878 if (PyTuple_GET_SIZE(pool) == 0)
1879 goto empty;
1880 elem = PyTuple_GET_ITEM(pool, 0);
1881 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001882 PyTuple_SET_ITEM(result, i, elem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001883 }
1884 } else {
1885 Py_ssize_t *indices = lz->indices;
1886 Py_ssize_t *maxvec = lz->maxvec;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001887
1888 /* Copy the previous result tuple or re-use it if available */
1889 if (Py_REFCNT(result) > 1) {
1890 PyObject *old_result = result;
1891 result = PyTuple_New(npools);
1892 if (result == NULL)
1893 goto empty;
1894 lz->result = result;
1895 for (i=0; i < npools; i++) {
1896 elem = PyTuple_GET_ITEM(old_result, i);
1897 Py_INCREF(elem);
1898 PyTuple_SET_ITEM(result, i, elem);
1899 }
1900 Py_DECREF(old_result);
1901 }
1902 /* Now, we've got the only copy so we can update it in-place */
1903 assert (Py_REFCNT(result) == 1);
1904
1905 /* Update the pool indices right-to-left. Only advance to the
1906 next pool when the previous one rolls-over */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001907 for (i=npools-1 ; i >= 0 ; i--) {
1908 pool = PyTuple_GET_ITEM(pools, i);
1909 indices[i]++;
1910 if (indices[i] == maxvec[i]) {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001911 /* Roll-over and advance to next pool */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001912 indices[i] = 0;
1913 elem = PyTuple_GET_ITEM(pool, 0);
1914 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001915 oldelem = PyTuple_GET_ITEM(result, i);
1916 PyTuple_SET_ITEM(result, i, elem);
1917 Py_DECREF(oldelem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001918 } else {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001919 /* No rollover. Just increment and stop here. */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001920 elem = PyTuple_GET_ITEM(pool, indices[i]);
1921 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001922 oldelem = PyTuple_GET_ITEM(result, i);
1923 PyTuple_SET_ITEM(result, i, elem);
1924 Py_DECREF(oldelem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001925 break;
1926 }
1927 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001928
1929 /* If i is negative, then the indices have all rolled-over
1930 and we're done. */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001931 if (i < 0)
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001932 goto empty;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001933 }
1934
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001935 Py_INCREF(result);
1936 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001937
1938empty:
1939 lz->stopped = 1;
1940 return NULL;
1941}
1942
1943PyDoc_STRVAR(product_doc,
1944"product(*iterables) --> product object\n\
1945\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001946Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001947For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1948The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1949cycle in a manner similar to an odometer (with the rightmost element changing\n\
1950on every iteration).\n\n\
1951product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1952product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1953
1954static PyTypeObject product_type = {
1955 PyVarObject_HEAD_INIT(NULL, 0)
1956 "itertools.product", /* tp_name */
1957 sizeof(productobject), /* tp_basicsize */
1958 0, /* tp_itemsize */
1959 /* methods */
1960 (destructor)product_dealloc, /* tp_dealloc */
1961 0, /* tp_print */
1962 0, /* tp_getattr */
1963 0, /* tp_setattr */
1964 0, /* tp_compare */
1965 0, /* tp_repr */
1966 0, /* tp_as_number */
1967 0, /* tp_as_sequence */
1968 0, /* tp_as_mapping */
1969 0, /* tp_hash */
1970 0, /* tp_call */
1971 0, /* tp_str */
1972 PyObject_GenericGetAttr, /* tp_getattro */
1973 0, /* tp_setattro */
1974 0, /* tp_as_buffer */
1975 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1976 Py_TPFLAGS_BASETYPE, /* tp_flags */
1977 product_doc, /* tp_doc */
1978 (traverseproc)product_traverse, /* tp_traverse */
1979 0, /* tp_clear */
1980 0, /* tp_richcompare */
1981 0, /* tp_weaklistoffset */
1982 PyObject_SelfIter, /* tp_iter */
1983 (iternextfunc)product_next, /* tp_iternext */
1984 0, /* tp_methods */
1985 0, /* tp_members */
1986 0, /* tp_getset */
1987 0, /* tp_base */
1988 0, /* tp_dict */
1989 0, /* tp_descr_get */
1990 0, /* tp_descr_set */
1991 0, /* tp_dictoffset */
1992 0, /* tp_init */
1993 0, /* tp_alloc */
1994 product_new, /* tp_new */
1995 PyObject_GC_Del, /* tp_free */
1996};
1997
1998
Christian Heimes380f7f22008-02-28 11:19:05 +00001999/* combinations object ************************************************************/
2000
2001typedef struct {
2002 PyObject_HEAD
2003 PyObject *pool; /* input converted to a tuple */
2004 Py_ssize_t *indices; /* one index per result element */
2005 PyObject *result; /* most recently returned result tuple */
2006 Py_ssize_t r; /* size of result tuple */
2007 int stopped; /* set to 1 when the combinations iterator is exhausted */
2008} combinationsobject;
2009
2010static PyTypeObject combinations_type;
2011
2012static PyObject *
2013combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2014{
2015 combinationsobject *co;
2016 Py_ssize_t n;
2017 Py_ssize_t r;
2018 PyObject *pool = NULL;
2019 PyObject *iterable = NULL;
2020 Py_ssize_t *indices = NULL;
2021 Py_ssize_t i;
2022 static char *kwargs[] = {"iterable", "r", NULL};
2023
2024 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2025 &iterable, &r))
2026 return NULL;
2027
2028 pool = PySequence_Tuple(iterable);
2029 if (pool == NULL)
2030 goto error;
2031 n = PyTuple_GET_SIZE(pool);
2032 if (r < 0) {
2033 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2034 goto error;
2035 }
2036 if (r > n) {
2037 PyErr_SetString(PyExc_ValueError, "r cannot be bigger than the iterable");
2038 goto error;
2039 }
2040
2041 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2042 if (indices == NULL) {
2043 PyErr_NoMemory();
2044 goto error;
2045 }
2046
2047 for (i=0 ; i<r ; i++)
2048 indices[i] = i;
2049
2050 /* create combinationsobject structure */
2051 co = (combinationsobject *)type->tp_alloc(type, 0);
2052 if (co == NULL)
2053 goto error;
2054
2055 co->pool = pool;
2056 co->indices = indices;
2057 co->result = NULL;
2058 co->r = r;
2059 co->stopped = 0;
2060
2061 return (PyObject *)co;
2062
2063error:
2064 if (indices != NULL)
2065 PyMem_Free(indices);
2066 Py_XDECREF(pool);
2067 return NULL;
2068}
2069
2070static void
2071combinations_dealloc(combinationsobject *co)
2072{
2073 PyObject_GC_UnTrack(co);
2074 Py_XDECREF(co->pool);
2075 Py_XDECREF(co->result);
2076 PyMem_Free(co->indices);
2077 Py_TYPE(co)->tp_free(co);
2078}
2079
2080static int
2081combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2082{
2083 Py_VISIT(co->pool);
2084 Py_VISIT(co->result);
2085 return 0;
2086}
2087
2088static PyObject *
2089combinations_next(combinationsobject *co)
2090{
2091 PyObject *elem;
2092 PyObject *oldelem;
2093 PyObject *pool = co->pool;
2094 Py_ssize_t *indices = co->indices;
2095 PyObject *result = co->result;
2096 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2097 Py_ssize_t r = co->r;
2098 Py_ssize_t i, j, index;
2099
2100 if (co->stopped)
2101 return NULL;
2102
2103 if (result == NULL) {
2104 /* On the first pass, initialize result tuple using the indices */
2105 result = PyTuple_New(r);
2106 if (result == NULL)
2107 goto empty;
2108 co->result = result;
2109 for (i=0; i<r ; i++) {
2110 index = indices[i];
2111 elem = PyTuple_GET_ITEM(pool, index);
2112 Py_INCREF(elem);
2113 PyTuple_SET_ITEM(result, i, elem);
2114 }
2115 } else {
2116 /* Copy the previous result tuple or re-use it if available */
2117 if (Py_REFCNT(result) > 1) {
2118 PyObject *old_result = result;
2119 result = PyTuple_New(r);
2120 if (result == NULL)
2121 goto empty;
2122 co->result = result;
2123 for (i=0; i<r ; i++) {
2124 elem = PyTuple_GET_ITEM(old_result, i);
2125 Py_INCREF(elem);
2126 PyTuple_SET_ITEM(result, i, elem);
2127 }
2128 Py_DECREF(old_result);
2129 }
2130 /* Now, we've got the only copy so we can update it in-place
2131 * CPython's empty tuple is a singleton and cached in
2132 * PyTuple's freelist.
2133 */
2134 assert(r == 0 || Py_REFCNT(result) == 1);
2135
2136 /* Scan indices right-to-left until finding one that is not
2137 at its maximum (i + n - r). */
2138 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2139 ;
2140
2141 /* If i is negative, then the indices are all at
2142 their maximum value and we're done. */
2143 if (i < 0)
2144 goto empty;
2145
2146 /* Increment the current index which we know is not at its
2147 maximum. Then move back to the right setting each index
2148 to its lowest possible value (one higher than the index
2149 to its left -- this maintains the sort order invariant). */
2150 indices[i]++;
2151 for (j=i+1 ; j<r ; j++)
2152 indices[j] = indices[j-1] + 1;
2153
2154 /* Update the result tuple for the new indices
2155 starting with i, the leftmost index that changed */
2156 for ( ; i<r ; i++) {
2157 index = indices[i];
2158 elem = PyTuple_GET_ITEM(pool, index);
2159 Py_INCREF(elem);
2160 oldelem = PyTuple_GET_ITEM(result, i);
2161 PyTuple_SET_ITEM(result, i, elem);
2162 Py_DECREF(oldelem);
2163 }
2164 }
2165
2166 Py_INCREF(result);
2167 return result;
2168
2169empty:
2170 co->stopped = 1;
2171 return NULL;
2172}
2173
2174PyDoc_STRVAR(combinations_doc,
2175"combinations(iterables) --> combinations object\n\
2176\n\
2177Return successive r-length combinations of elements in the iterable.\n\n\
2178combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2179
2180static PyTypeObject combinations_type = {
2181 PyVarObject_HEAD_INIT(NULL, 0)
2182 "itertools.combinations", /* tp_name */
2183 sizeof(combinationsobject), /* tp_basicsize */
2184 0, /* tp_itemsize */
2185 /* methods */
2186 (destructor)combinations_dealloc, /* tp_dealloc */
2187 0, /* tp_print */
2188 0, /* tp_getattr */
2189 0, /* tp_setattr */
2190 0, /* tp_compare */
2191 0, /* tp_repr */
2192 0, /* tp_as_number */
2193 0, /* tp_as_sequence */
2194 0, /* tp_as_mapping */
2195 0, /* tp_hash */
2196 0, /* tp_call */
2197 0, /* tp_str */
2198 PyObject_GenericGetAttr, /* tp_getattro */
2199 0, /* tp_setattro */
2200 0, /* tp_as_buffer */
2201 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2202 Py_TPFLAGS_BASETYPE, /* tp_flags */
2203 combinations_doc, /* tp_doc */
2204 (traverseproc)combinations_traverse, /* tp_traverse */
2205 0, /* tp_clear */
2206 0, /* tp_richcompare */
2207 0, /* tp_weaklistoffset */
2208 PyObject_SelfIter, /* tp_iter */
2209 (iternextfunc)combinations_next, /* tp_iternext */
2210 0, /* tp_methods */
2211 0, /* tp_members */
2212 0, /* tp_getset */
2213 0, /* tp_base */
2214 0, /* tp_dict */
2215 0, /* tp_descr_get */
2216 0, /* tp_descr_set */
2217 0, /* tp_dictoffset */
2218 0, /* tp_init */
2219 0, /* tp_alloc */
2220 combinations_new, /* tp_new */
2221 PyObject_GC_Del, /* tp_free */
2222};
2223
2224
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002225/* ifilter object ************************************************************/
2226
2227typedef struct {
2228 PyObject_HEAD
2229 PyObject *func;
2230 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002231} ifilterobject;
2232
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002233static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002234
2235static PyObject *
2236ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2237{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002238 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002239 PyObject *it;
2240 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002241
Thomas Woutersb2137042007-02-01 18:02:27 +00002242 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002243 return NULL;
2244
Raymond Hettinger60eca932003-02-09 06:40:58 +00002245 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002246 return NULL;
2247
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002248 /* Get iterator. */
2249 it = PyObject_GetIter(seq);
2250 if (it == NULL)
2251 return NULL;
2252
2253 /* create ifilterobject structure */
2254 lz = (ifilterobject *)type->tp_alloc(type, 0);
2255 if (lz == NULL) {
2256 Py_DECREF(it);
2257 return NULL;
2258 }
2259 Py_INCREF(func);
2260 lz->func = func;
2261 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002262
2263 return (PyObject *)lz;
2264}
2265
2266static void
2267ifilter_dealloc(ifilterobject *lz)
2268{
2269 PyObject_GC_UnTrack(lz);
2270 Py_XDECREF(lz->func);
2271 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00002272 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002273}
2274
2275static int
2276ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
2277{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002278 Py_VISIT(lz->it);
2279 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002280 return 0;
2281}
2282
2283static PyObject *
2284ifilter_next(ifilterobject *lz)
2285{
2286 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002287 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002288 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002289 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002290
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002291 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00002292 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002293 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002294 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002295 if (item == NULL)
2296 return NULL;
2297
Christian Heimes836baa52008-02-26 08:18:30 +00002298 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002299 ok = PyObject_IsTrue(item);
2300 } else {
2301 PyObject *good;
2302 good = PyObject_CallFunctionObjArgs(lz->func,
2303 item, NULL);
2304 if (good == NULL) {
2305 Py_DECREF(item);
2306 return NULL;
2307 }
2308 ok = PyObject_IsTrue(good);
2309 Py_DECREF(good);
2310 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00002311 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002312 return item;
2313 Py_DECREF(item);
2314 }
2315}
2316
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002317PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002318"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002319\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002320Return those items of sequence for which function(item) is true.\n\
2321If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002322
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002323static PyTypeObject ifilter_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002324 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002325 "itertools.ifilter", /* tp_name */
2326 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002327 0, /* tp_itemsize */
2328 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002329 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002330 0, /* tp_print */
2331 0, /* tp_getattr */
2332 0, /* tp_setattr */
2333 0, /* tp_compare */
2334 0, /* tp_repr */
2335 0, /* tp_as_number */
2336 0, /* tp_as_sequence */
2337 0, /* tp_as_mapping */
2338 0, /* tp_hash */
2339 0, /* tp_call */
2340 0, /* tp_str */
2341 PyObject_GenericGetAttr, /* tp_getattro */
2342 0, /* tp_setattro */
2343 0, /* tp_as_buffer */
2344 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2345 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002346 ifilter_doc, /* tp_doc */
2347 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002348 0, /* tp_clear */
2349 0, /* tp_richcompare */
2350 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002351 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002352 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002353 0, /* tp_methods */
2354 0, /* tp_members */
2355 0, /* tp_getset */
2356 0, /* tp_base */
2357 0, /* tp_dict */
2358 0, /* tp_descr_get */
2359 0, /* tp_descr_set */
2360 0, /* tp_dictoffset */
2361 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002362 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002363 ifilter_new, /* tp_new */
2364 PyObject_GC_Del, /* tp_free */
2365};
2366
2367
2368/* ifilterfalse object ************************************************************/
2369
2370typedef struct {
2371 PyObject_HEAD
2372 PyObject *func;
2373 PyObject *it;
2374} ifilterfalseobject;
2375
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002376static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002377
2378static PyObject *
2379ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2380{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00002381 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002382 PyObject *it;
2383 ifilterfalseobject *lz;
2384
Thomas Woutersb2137042007-02-01 18:02:27 +00002385 if (type == &ifilterfalse_type &&
2386 !_PyArg_NoKeywords("ifilterfalse()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002387 return NULL;
2388
Raymond Hettinger60eca932003-02-09 06:40:58 +00002389 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
2390 return NULL;
2391
2392 /* Get iterator. */
2393 it = PyObject_GetIter(seq);
2394 if (it == NULL)
2395 return NULL;
2396
2397 /* create ifilterfalseobject structure */
2398 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
2399 if (lz == NULL) {
2400 Py_DECREF(it);
2401 return NULL;
2402 }
2403 Py_INCREF(func);
2404 lz->func = func;
2405 lz->it = it;
2406
2407 return (PyObject *)lz;
2408}
2409
2410static void
2411ifilterfalse_dealloc(ifilterfalseobject *lz)
2412{
2413 PyObject_GC_UnTrack(lz);
2414 Py_XDECREF(lz->func);
2415 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00002416 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002417}
2418
2419static int
2420ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
2421{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002422 Py_VISIT(lz->it);
2423 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002424 return 0;
2425}
2426
2427static PyObject *
2428ifilterfalse_next(ifilterfalseobject *lz)
2429{
2430 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002431 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002432 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002433 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002434
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002435 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00002436 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002437 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002438 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002439 if (item == NULL)
2440 return NULL;
2441
Christian Heimes836baa52008-02-26 08:18:30 +00002442 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
Raymond Hettinger60eca932003-02-09 06:40:58 +00002443 ok = PyObject_IsTrue(item);
2444 } else {
2445 PyObject *good;
2446 good = PyObject_CallFunctionObjArgs(lz->func,
2447 item, NULL);
2448 if (good == NULL) {
2449 Py_DECREF(item);
2450 return NULL;
2451 }
2452 ok = PyObject_IsTrue(good);
2453 Py_DECREF(good);
2454 }
2455 if (!ok)
2456 return item;
2457 Py_DECREF(item);
2458 }
2459}
2460
Raymond Hettinger60eca932003-02-09 06:40:58 +00002461PyDoc_STRVAR(ifilterfalse_doc,
2462"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
2463\n\
2464Return those items of sequence for which function(item) is false.\n\
2465If function is None, return the items that are false.");
2466
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002467static PyTypeObject ifilterfalse_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002468 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002469 "itertools.ifilterfalse", /* tp_name */
2470 sizeof(ifilterfalseobject), /* tp_basicsize */
2471 0, /* tp_itemsize */
2472 /* methods */
2473 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
2474 0, /* tp_print */
2475 0, /* tp_getattr */
2476 0, /* tp_setattr */
2477 0, /* tp_compare */
2478 0, /* tp_repr */
2479 0, /* tp_as_number */
2480 0, /* tp_as_sequence */
2481 0, /* tp_as_mapping */
2482 0, /* tp_hash */
2483 0, /* tp_call */
2484 0, /* tp_str */
2485 PyObject_GenericGetAttr, /* tp_getattro */
2486 0, /* tp_setattro */
2487 0, /* tp_as_buffer */
2488 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2489 Py_TPFLAGS_BASETYPE, /* tp_flags */
2490 ifilterfalse_doc, /* tp_doc */
2491 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
2492 0, /* tp_clear */
2493 0, /* tp_richcompare */
2494 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002495 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002496 (iternextfunc)ifilterfalse_next, /* tp_iternext */
2497 0, /* tp_methods */
2498 0, /* tp_members */
2499 0, /* tp_getset */
2500 0, /* tp_base */
2501 0, /* tp_dict */
2502 0, /* tp_descr_get */
2503 0, /* tp_descr_set */
2504 0, /* tp_dictoffset */
2505 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002506 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002507 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002508 PyObject_GC_Del, /* tp_free */
2509};
2510
2511
2512/* count object ************************************************************/
2513
2514typedef struct {
2515 PyObject_HEAD
Thomas Wouters477c8d52006-05-27 19:21:47 +00002516 Py_ssize_t cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002517 PyObject *long_cnt; /* Arbitrarily large count when cnt >= PY_SSIZE_T_MAX */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002518} countobject;
2519
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002520static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002521
2522static PyObject *
2523count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2524{
2525 countobject *lz;
Thomas Wouters477c8d52006-05-27 19:21:47 +00002526 Py_ssize_t cnt = 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002527 PyObject *cnt_arg = NULL;
2528 PyObject *long_cnt = NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002529
Thomas Woutersb2137042007-02-01 18:02:27 +00002530 if (type == &count_type && !_PyArg_NoKeywords("count()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002531 return NULL;
2532
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002533 if (!PyArg_UnpackTuple(args, "count", 0, 1, &cnt_arg))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002534 return NULL;
2535
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002536 if (cnt_arg != NULL) {
Christian Heimes217cfd12007-12-02 14:31:20 +00002537 cnt = PyLong_AsSsize_t(cnt_arg);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002538 if (cnt == -1 && PyErr_Occurred()) {
2539 PyErr_Clear();
2540 if (!PyLong_Check(cnt_arg)) {
2541 PyErr_SetString(PyExc_TypeError, "an integer is required");
2542 return NULL;
2543 }
2544 long_cnt = cnt_arg;
2545 Py_INCREF(long_cnt);
2546 cnt = PY_SSIZE_T_MAX;
2547 }
2548 }
2549
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002550 /* create countobject structure */
2551 lz = (countobject *)PyObject_New(countobject, &count_type);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002552 if (lz == NULL) {
2553 Py_XDECREF(long_cnt);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002554 return NULL;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002555 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002556 lz->cnt = cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002557 lz->long_cnt = long_cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002558
2559 return (PyObject *)lz;
2560}
2561
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002562static void
2563count_dealloc(countobject *lz)
2564{
2565 Py_XDECREF(lz->long_cnt);
2566 PyObject_Del(lz);
2567}
2568
2569static PyObject *
2570count_nextlong(countobject *lz)
2571{
2572 static PyObject *one = NULL;
2573 PyObject *cnt;
2574 PyObject *stepped_up;
2575
2576 if (lz->long_cnt == NULL) {
Christian Heimes217cfd12007-12-02 14:31:20 +00002577 lz->long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002578 if (lz->long_cnt == NULL)
2579 return NULL;
2580 }
2581 if (one == NULL) {
Christian Heimes217cfd12007-12-02 14:31:20 +00002582 one = PyLong_FromLong(1);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002583 if (one == NULL)
2584 return NULL;
2585 }
2586 cnt = lz->long_cnt;
2587 assert(cnt != NULL);
2588 stepped_up = PyNumber_Add(cnt, one);
2589 if (stepped_up == NULL)
2590 return NULL;
2591 lz->long_cnt = stepped_up;
2592 return cnt;
2593}
2594
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002595static PyObject *
2596count_next(countobject *lz)
2597{
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002598 if (lz->cnt == PY_SSIZE_T_MAX)
2599 return count_nextlong(lz);
Christian Heimes217cfd12007-12-02 14:31:20 +00002600 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002601}
2602
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002603static PyObject *
2604count_repr(countobject *lz)
2605{
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002606 if (lz->cnt != PY_SSIZE_T_MAX)
2607 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
2608
2609 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002610}
2611
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002612PyDoc_STRVAR(count_doc,
2613"count([firstval]) --> count object\n\
2614\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002615Return a count object whose .__next__() method returns consecutive\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002616integers starting from zero or, if specified, from firstval.");
2617
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002618static PyTypeObject count_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002619 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002620 "itertools.count", /* tp_name */
2621 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002622 0, /* tp_itemsize */
2623 /* methods */
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002624 (destructor)count_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002625 0, /* tp_print */
2626 0, /* tp_getattr */
2627 0, /* tp_setattr */
2628 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002629 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002630 0, /* tp_as_number */
2631 0, /* tp_as_sequence */
2632 0, /* tp_as_mapping */
2633 0, /* tp_hash */
2634 0, /* tp_call */
2635 0, /* tp_str */
2636 PyObject_GenericGetAttr, /* tp_getattro */
2637 0, /* tp_setattro */
2638 0, /* tp_as_buffer */
2639 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002640 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002641 0, /* tp_traverse */
2642 0, /* tp_clear */
2643 0, /* tp_richcompare */
2644 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002645 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002646 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002647 0, /* tp_methods */
2648 0, /* tp_members */
2649 0, /* tp_getset */
2650 0, /* tp_base */
2651 0, /* tp_dict */
2652 0, /* tp_descr_get */
2653 0, /* tp_descr_set */
2654 0, /* tp_dictoffset */
2655 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002656 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002657 count_new, /* tp_new */
2658};
2659
2660
2661/* izip object ************************************************************/
2662
2663#include "Python.h"
2664
2665typedef struct {
2666 PyObject_HEAD
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002667 Py_ssize_t tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002668 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002669 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002670} izipobject;
2671
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002672static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002673
2674static PyObject *
2675izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2676{
2677 izipobject *lz;
Thomas Wouters477c8d52006-05-27 19:21:47 +00002678 Py_ssize_t i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002679 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002680 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002681 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002682
Thomas Woutersb2137042007-02-01 18:02:27 +00002683 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002684 return NULL;
2685
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002686 /* args must be a tuple */
2687 assert(PyTuple_Check(args));
2688
2689 /* obtain iterators */
2690 ittuple = PyTuple_New(tuplesize);
Thomas Wouters477c8d52006-05-27 19:21:47 +00002691 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002692 return NULL;
2693 for (i=0; i < tuplesize; ++i) {
2694 PyObject *item = PyTuple_GET_ITEM(args, i);
2695 PyObject *it = PyObject_GetIter(item);
2696 if (it == NULL) {
2697 if (PyErr_ExceptionMatches(PyExc_TypeError))
2698 PyErr_Format(PyExc_TypeError,
Thomas Wouters477c8d52006-05-27 19:21:47 +00002699 "izip argument #%zd must support iteration",
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002700 i+1);
2701 Py_DECREF(ittuple);
2702 return NULL;
2703 }
2704 PyTuple_SET_ITEM(ittuple, i, it);
2705 }
2706
Raymond Hettinger2012f172003-02-07 05:32:58 +00002707 /* create a result holder */
2708 result = PyTuple_New(tuplesize);
2709 if (result == NULL) {
2710 Py_DECREF(ittuple);
2711 return NULL;
2712 }
2713 for (i=0 ; i < tuplesize ; i++) {
2714 Py_INCREF(Py_None);
2715 PyTuple_SET_ITEM(result, i, Py_None);
2716 }
2717
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002718 /* create izipobject structure */
2719 lz = (izipobject *)type->tp_alloc(type, 0);
2720 if (lz == NULL) {
2721 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002722 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002723 return NULL;
2724 }
2725 lz->ittuple = ittuple;
2726 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002727 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002728
2729 return (PyObject *)lz;
2730}
2731
2732static void
2733izip_dealloc(izipobject *lz)
2734{
2735 PyObject_GC_UnTrack(lz);
2736 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002737 Py_XDECREF(lz->result);
Christian Heimes90aa7642007-12-19 02:45:37 +00002738 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002739}
2740
2741static int
2742izip_traverse(izipobject *lz, visitproc visit, void *arg)
2743{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002744 Py_VISIT(lz->ittuple);
2745 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002746 return 0;
2747}
2748
2749static PyObject *
2750izip_next(izipobject *lz)
2751{
Thomas Wouters477c8d52006-05-27 19:21:47 +00002752 Py_ssize_t i;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002753 Py_ssize_t tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002754 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002755 PyObject *it;
2756 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002757 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002758
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002759 if (tuplesize == 0)
2760 return NULL;
Christian Heimes90aa7642007-12-19 02:45:37 +00002761 if (Py_REFCNT(result) == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002762 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002763 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002764 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002765 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00002766 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002767 if (item == NULL) {
2768 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002769 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002770 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002771 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002772 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002773 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002774 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002775 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002776 result = PyTuple_New(tuplesize);
2777 if (result == NULL)
2778 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002779 for (i=0 ; i < tuplesize ; i++) {
2780 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002781 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00002782 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002783 if (item == NULL) {
2784 Py_DECREF(result);
2785 return NULL;
2786 }
2787 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002788 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002789 }
2790 return result;
2791}
2792
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002793PyDoc_STRVAR(izip_doc,
2794"izip(iter1 [,iter2 [...]]) --> izip object\n\
2795\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002796Return a izip object whose .__next__() method returns a tuple where\n\
2797the i-th element comes from the i-th iterable argument. The .__next__()\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002798method continues until the shortest iterable in the argument sequence\n\
2799is exhausted and then it raises StopIteration. Works like the zip()\n\
2800function but consumes less memory by returning an iterator instead of\n\
2801a list.");
2802
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002803static PyTypeObject izip_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002804 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002805 "itertools.izip", /* tp_name */
2806 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002807 0, /* tp_itemsize */
2808 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002809 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002810 0, /* tp_print */
2811 0, /* tp_getattr */
2812 0, /* tp_setattr */
2813 0, /* tp_compare */
2814 0, /* tp_repr */
2815 0, /* tp_as_number */
2816 0, /* tp_as_sequence */
2817 0, /* tp_as_mapping */
2818 0, /* tp_hash */
2819 0, /* tp_call */
2820 0, /* tp_str */
2821 PyObject_GenericGetAttr, /* tp_getattro */
2822 0, /* tp_setattro */
2823 0, /* tp_as_buffer */
2824 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2825 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002826 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002827 (traverseproc)izip_traverse, /* tp_traverse */
2828 0, /* tp_clear */
2829 0, /* tp_richcompare */
2830 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002831 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002832 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002833 0, /* tp_methods */
2834 0, /* tp_members */
2835 0, /* tp_getset */
2836 0, /* tp_base */
2837 0, /* tp_dict */
2838 0, /* tp_descr_get */
2839 0, /* tp_descr_set */
2840 0, /* tp_dictoffset */
2841 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002842 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002843 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002844 PyObject_GC_Del, /* tp_free */
2845};
2846
2847
2848/* repeat object ************************************************************/
2849
2850typedef struct {
2851 PyObject_HEAD
2852 PyObject *element;
Thomas Wouters477c8d52006-05-27 19:21:47 +00002853 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002854} repeatobject;
2855
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002856static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002857
2858static PyObject *
2859repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2860{
2861 repeatobject *ro;
2862 PyObject *element;
Thomas Wouters477c8d52006-05-27 19:21:47 +00002863 Py_ssize_t cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002864
Thomas Woutersb2137042007-02-01 18:02:27 +00002865 if (type == &repeat_type && !_PyArg_NoKeywords("repeat()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002866 return NULL;
2867
Thomas Wouters477c8d52006-05-27 19:21:47 +00002868 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002869 return NULL;
2870
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002871 if (PyTuple_Size(args) == 2 && cnt < 0)
2872 cnt = 0;
2873
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002874 ro = (repeatobject *)type->tp_alloc(type, 0);
2875 if (ro == NULL)
2876 return NULL;
2877 Py_INCREF(element);
2878 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002879 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002880 return (PyObject *)ro;
2881}
2882
2883static void
2884repeat_dealloc(repeatobject *ro)
2885{
2886 PyObject_GC_UnTrack(ro);
2887 Py_XDECREF(ro->element);
Christian Heimes90aa7642007-12-19 02:45:37 +00002888 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002889}
2890
2891static int
2892repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2893{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002894 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002895 return 0;
2896}
2897
2898static PyObject *
2899repeat_next(repeatobject *ro)
2900{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002901 if (ro->cnt == 0)
2902 return NULL;
2903 if (ro->cnt > 0)
2904 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002905 Py_INCREF(ro->element);
2906 return ro->element;
2907}
2908
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002909static PyObject *
2910repeat_repr(repeatobject *ro)
2911{
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002912 if (ro->cnt == -1)
Walter Dörwald7569dfe2007-05-19 21:49:49 +00002913 return PyUnicode_FromFormat("repeat(%R)", ro->element);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002914 else
Walter Dörwald7569dfe2007-05-19 21:49:49 +00002915 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002916}
2917
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002918static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002919repeat_len(repeatobject *ro)
2920{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002921 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002922 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002923 return NULL;
2924 }
Christian Heimes217cfd12007-12-02 14:31:20 +00002925 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002926}
2927
Armin Rigof5b3e362006-02-11 21:32:43 +00002928PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002929
2930static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00002931 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002932 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002933};
2934
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002935PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002936"repeat(element [,times]) -> create an iterator which returns the element\n\
2937for the specified number of times. If not specified, returns the element\n\
2938endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002939
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002940static PyTypeObject repeat_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002941 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002942 "itertools.repeat", /* tp_name */
2943 sizeof(repeatobject), /* tp_basicsize */
2944 0, /* tp_itemsize */
2945 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002946 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002947 0, /* tp_print */
2948 0, /* tp_getattr */
2949 0, /* tp_setattr */
2950 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002951 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002952 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002953 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002954 0, /* tp_as_mapping */
2955 0, /* tp_hash */
2956 0, /* tp_call */
2957 0, /* tp_str */
2958 PyObject_GenericGetAttr, /* tp_getattro */
2959 0, /* tp_setattro */
2960 0, /* tp_as_buffer */
2961 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2962 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002963 repeat_doc, /* tp_doc */
2964 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002965 0, /* tp_clear */
2966 0, /* tp_richcompare */
2967 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002968 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002969 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002970 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002971 0, /* tp_members */
2972 0, /* tp_getset */
2973 0, /* tp_base */
2974 0, /* tp_dict */
2975 0, /* tp_descr_get */
2976 0, /* tp_descr_set */
2977 0, /* tp_dictoffset */
2978 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002979 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002980 repeat_new, /* tp_new */
2981 PyObject_GC_Del, /* tp_free */
2982};
2983
Thomas Wouterscf297e42007-02-23 15:07:44 +00002984/* iziplongest object ************************************************************/
2985
2986#include "Python.h"
2987
2988typedef struct {
2989 PyObject_HEAD
2990 Py_ssize_t tuplesize;
2991 Py_ssize_t numactive;
2992 PyObject *ittuple; /* tuple of iterators */
2993 PyObject *result;
2994 PyObject *fillvalue;
2995} iziplongestobject;
2996
2997static PyTypeObject iziplongest_type;
2998
2999static PyObject *
3000izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3001{
3002 iziplongestobject *lz;
3003 Py_ssize_t i;
3004 PyObject *ittuple; /* tuple of iterators */
3005 PyObject *result;
3006 PyObject *fillvalue = Py_None;
3007 Py_ssize_t tuplesize = PySequence_Length(args);
3008
3009 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3010 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3011 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3012 PyErr_SetString(PyExc_TypeError,
3013 "izip_longest() got an unexpected keyword argument");
3014 return NULL;
3015 }
3016 }
3017
3018 /* args must be a tuple */
3019 assert(PyTuple_Check(args));
3020
3021 /* obtain iterators */
3022 ittuple = PyTuple_New(tuplesize);
3023 if (ittuple == NULL)
3024 return NULL;
3025 for (i=0; i < tuplesize; ++i) {
3026 PyObject *item = PyTuple_GET_ITEM(args, i);
3027 PyObject *it = PyObject_GetIter(item);
3028 if (it == NULL) {
3029 if (PyErr_ExceptionMatches(PyExc_TypeError))
3030 PyErr_Format(PyExc_TypeError,
3031 "izip_longest argument #%zd must support iteration",
3032 i+1);
3033 Py_DECREF(ittuple);
3034 return NULL;
3035 }
3036 PyTuple_SET_ITEM(ittuple, i, it);
3037 }
3038
3039 /* create a result holder */
3040 result = PyTuple_New(tuplesize);
3041 if (result == NULL) {
3042 Py_DECREF(ittuple);
3043 return NULL;
3044 }
3045 for (i=0 ; i < tuplesize ; i++) {
3046 Py_INCREF(Py_None);
3047 PyTuple_SET_ITEM(result, i, Py_None);
3048 }
3049
3050 /* create iziplongestobject structure */
3051 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3052 if (lz == NULL) {
3053 Py_DECREF(ittuple);
3054 Py_DECREF(result);
3055 return NULL;
3056 }
3057 lz->ittuple = ittuple;
3058 lz->tuplesize = tuplesize;
3059 lz->numactive = tuplesize;
3060 lz->result = result;
3061 Py_INCREF(fillvalue);
3062 lz->fillvalue = fillvalue;
3063 return (PyObject *)lz;
3064}
3065
3066static void
3067izip_longest_dealloc(iziplongestobject *lz)
3068{
3069 PyObject_GC_UnTrack(lz);
3070 Py_XDECREF(lz->ittuple);
3071 Py_XDECREF(lz->result);
3072 Py_XDECREF(lz->fillvalue);
Christian Heimes90aa7642007-12-19 02:45:37 +00003073 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003074}
3075
3076static int
3077izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3078{
3079 Py_VISIT(lz->ittuple);
3080 Py_VISIT(lz->result);
3081 Py_VISIT(lz->fillvalue);
3082 return 0;
3083}
3084
3085static PyObject *
3086izip_longest_next(iziplongestobject *lz)
3087{
3088 Py_ssize_t i;
3089 Py_ssize_t tuplesize = lz->tuplesize;
3090 PyObject *result = lz->result;
3091 PyObject *it;
3092 PyObject *item;
3093 PyObject *olditem;
3094
3095 if (tuplesize == 0)
3096 return NULL;
3097 if (lz->numactive == 0)
3098 return NULL;
Christian Heimes90aa7642007-12-19 02:45:37 +00003099 if (Py_REFCNT(result) == 1) {
Thomas Wouterscf297e42007-02-23 15:07:44 +00003100 Py_INCREF(result);
3101 for (i=0 ; i < tuplesize ; i++) {
3102 it = PyTuple_GET_ITEM(lz->ittuple, i);
3103 if (it == NULL) {
3104 Py_INCREF(lz->fillvalue);
3105 item = lz->fillvalue;
3106 } else {
3107 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00003108 item = (*Py_TYPE(it)->tp_iternext)(it);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003109 if (item == NULL) {
3110 lz->numactive -= 1;
3111 if (lz->numactive == 0) {
3112 Py_DECREF(result);
3113 return NULL;
3114 } else {
3115 Py_INCREF(lz->fillvalue);
3116 item = lz->fillvalue;
3117 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3118 Py_DECREF(it);
3119 }
3120 }
3121 }
3122 olditem = PyTuple_GET_ITEM(result, i);
3123 PyTuple_SET_ITEM(result, i, item);
3124 Py_DECREF(olditem);
3125 }
3126 } else {
3127 result = PyTuple_New(tuplesize);
3128 if (result == NULL)
3129 return NULL;
3130 for (i=0 ; i < tuplesize ; i++) {
3131 it = PyTuple_GET_ITEM(lz->ittuple, i);
3132 if (it == NULL) {
3133 Py_INCREF(lz->fillvalue);
3134 item = lz->fillvalue;
3135 } else {
3136 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00003137 item = (*Py_TYPE(it)->tp_iternext)(it);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003138 if (item == NULL) {
3139 lz->numactive -= 1;
3140 if (lz->numactive == 0) {
3141 Py_DECREF(result);
3142 return NULL;
3143 } else {
3144 Py_INCREF(lz->fillvalue);
3145 item = lz->fillvalue;
3146 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3147 Py_DECREF(it);
3148 }
3149 }
3150 }
3151 PyTuple_SET_ITEM(result, i, item);
3152 }
3153 }
3154 return result;
3155}
3156
3157PyDoc_STRVAR(izip_longest_doc,
3158"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3159\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00003160Return an izip_longest object whose .__next__() method returns a tuple where\n\
3161the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003162method continues until the longest iterable in the argument sequence\n\
3163is exhausted and then it raises StopIteration. When the shorter iterables\n\
3164are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3165defaults to None or can be specified by a keyword argument.\n\
3166");
3167
3168static PyTypeObject iziplongest_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00003169 PyVarObject_HEAD_INIT(NULL, 0)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003170 "itertools.izip_longest", /* tp_name */
3171 sizeof(iziplongestobject), /* tp_basicsize */
3172 0, /* tp_itemsize */
3173 /* methods */
3174 (destructor)izip_longest_dealloc, /* tp_dealloc */
3175 0, /* tp_print */
3176 0, /* tp_getattr */
3177 0, /* tp_setattr */
3178 0, /* tp_compare */
3179 0, /* tp_repr */
3180 0, /* tp_as_number */
3181 0, /* tp_as_sequence */
3182 0, /* tp_as_mapping */
3183 0, /* tp_hash */
3184 0, /* tp_call */
3185 0, /* tp_str */
3186 PyObject_GenericGetAttr, /* tp_getattro */
3187 0, /* tp_setattro */
3188 0, /* tp_as_buffer */
3189 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3190 Py_TPFLAGS_BASETYPE, /* tp_flags */
3191 izip_longest_doc, /* tp_doc */
3192 (traverseproc)izip_longest_traverse, /* tp_traverse */
3193 0, /* tp_clear */
3194 0, /* tp_richcompare */
3195 0, /* tp_weaklistoffset */
3196 PyObject_SelfIter, /* tp_iter */
3197 (iternextfunc)izip_longest_next, /* tp_iternext */
3198 0, /* tp_methods */
3199 0, /* tp_members */
3200 0, /* tp_getset */
3201 0, /* tp_base */
3202 0, /* tp_dict */
3203 0, /* tp_descr_get */
3204 0, /* tp_descr_set */
3205 0, /* tp_dictoffset */
3206 0, /* tp_init */
3207 0, /* tp_alloc */
3208 izip_longest_new, /* tp_new */
3209 PyObject_GC_Del, /* tp_free */
3210};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003211
3212/* module level code ********************************************************/
3213
3214PyDoc_STRVAR(module_doc,
3215"Functional tools for creating and using iterators.\n\
3216\n\
3217Infinite iterators:\n\
3218count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003219cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003220repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003221\n\
3222Iterators terminating on the shortest input sequence:\n\
3223izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003224izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003225ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
3226ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003227islice(seq, [start,] stop [, step]) --> elements from\n\
3228 seq[start:stop:step]\n\
3229imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
3230starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003231tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003232chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003233takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
3234dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003235groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003236");
3237
3238
Raymond Hettingerad983e72003-11-12 14:32:26 +00003239static PyMethodDef module_methods[] = {
3240 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3241 {NULL, NULL} /* sentinel */
3242};
3243
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003244PyMODINIT_FUNC
3245inititertools(void)
3246{
Raymond Hettinger60eca932003-02-09 06:40:58 +00003247 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003248 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003249 char *name;
3250 PyTypeObject *typelist[] = {
Christian Heimes380f7f22008-02-28 11:19:05 +00003251 &combinations_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003252 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003253 &dropwhile_type,
3254 &takewhile_type,
3255 &islice_type,
3256 &starmap_type,
3257 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003258 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003259 &ifilter_type,
3260 &ifilterfalse_type,
3261 &count_type,
3262 &izip_type,
Christian Heimesc3f30c42008-02-22 16:37:40 +00003263 &iziplongest_type,
Christian Heimes380f7f22008-02-28 11:19:05 +00003264 &product_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003265 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003266 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003267 NULL
3268 };
3269
Christian Heimes90aa7642007-12-19 02:45:37 +00003270 Py_TYPE(&teedataobject_type) = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00003271 m = Py_InitModule3("itertools", module_methods, module_doc);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00003272 if (m == NULL)
3273 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003274
Raymond Hettinger60eca932003-02-09 06:40:58 +00003275 for (i=0 ; typelist[i] != NULL ; i++) {
3276 if (PyType_Ready(typelist[i]) < 0)
3277 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003278 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003279 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003280 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003281 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003282 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003283
3284 if (PyType_Ready(&teedataobject_type) < 0)
3285 return;
3286 if (PyType_Ready(&tee_type) < 0)
3287 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003288 if (PyType_Ready(&_grouper_type) < 0)
3289 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003290}