blob: 283d8483d63b998b971b5653da9f020280cc37a9 [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
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001421/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001422
1423typedef struct {
1424 PyObject_HEAD
Christian Heimesf16baeb2008-02-29 14:57:44 +00001425 PyObject *source; /* Iterator over input iterables */
1426 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001427} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001428
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001429static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001430
Christian Heimesf16baeb2008-02-29 14:57:44 +00001431static PyObject *
1432chain_new_internal(PyTypeObject *type, PyObject *source)
1433{
1434 chainobject *lz;
1435
1436 lz = (chainobject *)type->tp_alloc(type, 0);
1437 if (lz == NULL) {
1438 Py_DECREF(source);
1439 return NULL;
1440 }
1441
1442 lz->source = source;
1443 lz->active = NULL;
1444 return (PyObject *)lz;
1445}
1446
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001447static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001448chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001449{
Christian Heimesf16baeb2008-02-29 14:57:44 +00001450 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001451
Thomas Woutersb2137042007-02-01 18:02:27 +00001452 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001453 return NULL;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001454
1455 source = PyObject_GetIter(args);
1456 if (source == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001458
Christian Heimesf16baeb2008-02-29 14:57:44 +00001459 return chain_new_internal(type, source);
1460}
1461
1462static PyObject *
1463chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1464{
1465 PyObject *source;
1466
1467 source = PyObject_GetIter(arg);
1468 if (source == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001470
Christian Heimesf16baeb2008-02-29 14:57:44 +00001471 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001472}
1473
1474static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001475chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001476{
1477 PyObject_GC_UnTrack(lz);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001478 Py_XDECREF(lz->active);
1479 Py_XDECREF(lz->source);
Christian Heimes90aa7642007-12-19 02:45:37 +00001480 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001481}
1482
Raymond Hettinger2012f172003-02-07 05:32:58 +00001483static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001484chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001485{
Christian Heimesf16baeb2008-02-29 14:57:44 +00001486 Py_VISIT(lz->source);
1487 Py_VISIT(lz->active);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001488 return 0;
1489}
1490
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001491static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001492chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001493{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001494 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001495
Christian Heimesf16baeb2008-02-29 14:57:44 +00001496 if (lz->source == NULL)
1497 return NULL; /* already stopped */
1498
1499 if (lz->active == NULL) {
1500 PyObject *iterable = PyIter_Next(lz->source);
1501 if (iterable == NULL) {
1502 Py_CLEAR(lz->source);
1503 return NULL; /* no more input sources */
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001504 }
Christian Heimesf16baeb2008-02-29 14:57:44 +00001505 lz->active = PyObject_GetIter(iterable);
Christian Heimes78644762008-03-04 23:39:23 +00001506 Py_DECREF(iterable);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001507 if (lz->active == NULL) {
Christian Heimesf16baeb2008-02-29 14:57:44 +00001508 Py_CLEAR(lz->source);
1509 return NULL; /* input not iterable */
1510 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001511 }
Christian Heimesf16baeb2008-02-29 14:57:44 +00001512 item = PyIter_Next(lz->active);
1513 if (item != NULL)
1514 return item;
1515 if (PyErr_Occurred()) {
1516 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1517 PyErr_Clear();
1518 else
1519 return NULL; /* input raised an exception */
1520 }
1521 Py_CLEAR(lz->active);
1522 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001523}
1524
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001525PyDoc_STRVAR(chain_doc,
1526"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001527\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001528Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001529first iterable until it is exhausted, then elements from the next\n\
1530iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001531
Christian Heimesf16baeb2008-02-29 14:57:44 +00001532PyDoc_STRVAR(chain_from_iterable_doc,
1533"chain.from_iterable(iterable) --> chain object\n\
1534\n\
1535Alternate chain() contructor taking a single iterable argument\n\
1536that evaluates lazily.");
1537
1538static PyMethodDef chain_methods[] = {
1539 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1540 chain_from_iterable_doc},
1541 {NULL, NULL} /* sentinel */
1542};
1543
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001544static PyTypeObject chain_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001545 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001546 "itertools.chain", /* tp_name */
1547 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001548 0, /* tp_itemsize */
1549 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001550 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001551 0, /* tp_print */
1552 0, /* tp_getattr */
1553 0, /* tp_setattr */
1554 0, /* tp_compare */
1555 0, /* tp_repr */
1556 0, /* tp_as_number */
1557 0, /* tp_as_sequence */
1558 0, /* tp_as_mapping */
1559 0, /* tp_hash */
1560 0, /* tp_call */
1561 0, /* tp_str */
1562 PyObject_GenericGetAttr, /* tp_getattro */
1563 0, /* tp_setattro */
1564 0, /* tp_as_buffer */
1565 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1566 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001567 chain_doc, /* tp_doc */
1568 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001569 0, /* tp_clear */
1570 0, /* tp_richcompare */
1571 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001572 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001573 (iternextfunc)chain_next, /* tp_iternext */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001574 chain_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001575 0, /* tp_members */
1576 0, /* tp_getset */
1577 0, /* tp_base */
1578 0, /* tp_dict */
1579 0, /* tp_descr_get */
1580 0, /* tp_descr_set */
1581 0, /* tp_dictoffset */
1582 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001583 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001584 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001585 PyObject_GC_Del, /* tp_free */
1586};
1587
1588
Christian Heimesc3f30c42008-02-22 16:37:40 +00001589/* product object ************************************************************/
1590
1591typedef struct {
1592 PyObject_HEAD
1593 PyObject *pools; /* tuple of pool tuples */
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001594 Py_ssize_t *indices; /* one index per pool */
1595 PyObject *result; /* most recently returned result tuple */
1596 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001597} productobject;
1598
1599static PyTypeObject product_type;
1600
1601static PyObject *
1602product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1603{
1604 productobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001605 Py_ssize_t nargs, npools, repeat=1;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001606 PyObject *pools = NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001607 Py_ssize_t *indices = NULL;
1608 Py_ssize_t i;
1609
Christian Heimesf16baeb2008-02-29 14:57:44 +00001610 if (kwds != NULL) {
1611 char *kwlist[] = {"repeat", 0};
1612 PyObject *tmpargs = PyTuple_New(0);
1613 if (tmpargs == NULL)
1614 return NULL;
1615 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1616 Py_DECREF(tmpargs);
1617 return NULL;
1618 }
1619 Py_DECREF(tmpargs);
1620 if (repeat < 0) {
1621 PyErr_SetString(PyExc_ValueError,
1622 "repeat argument cannot be negative");
1623 return NULL;
1624 }
1625 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001626
1627 assert(PyTuple_Check(args));
Christian Heimesf16baeb2008-02-29 14:57:44 +00001628 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1629 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001630
Christian Heimesc3f30c42008-02-22 16:37:40 +00001631 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
Christian Heimesb558a2e2008-03-02 22:46:37 +00001632 if (indices == NULL) {
Christian Heimesc3f30c42008-02-22 16:37:40 +00001633 PyErr_NoMemory();
1634 goto error;
1635 }
1636
1637 pools = PyTuple_New(npools);
1638 if (pools == NULL)
1639 goto error;
1640
Christian Heimesf16baeb2008-02-29 14:57:44 +00001641 for (i=0; i < nargs ; ++i) {
Christian Heimesc3f30c42008-02-22 16:37:40 +00001642 PyObject *item = PyTuple_GET_ITEM(args, i);
1643 PyObject *pool = PySequence_Tuple(item);
1644 if (pool == NULL)
1645 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001646 PyTuple_SET_ITEM(pools, i, pool);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001647 indices[i] = 0;
1648 }
Christian Heimesf16baeb2008-02-29 14:57:44 +00001649 for ( ; i < npools; ++i) {
1650 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1651 Py_INCREF(pool);
1652 PyTuple_SET_ITEM(pools, i, pool);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001653 indices[i] = 0;
1654 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001655
1656 /* create productobject structure */
1657 lz = (productobject *)type->tp_alloc(type, 0);
Christian Heimes380f7f22008-02-28 11:19:05 +00001658 if (lz == NULL)
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001659 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001660
1661 lz->pools = pools;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001662 lz->indices = indices;
1663 lz->result = NULL;
1664 lz->stopped = 0;
1665
1666 return (PyObject *)lz;
1667
1668error:
Christian Heimesc3f30c42008-02-22 16:37:40 +00001669 if (indices != NULL)
1670 PyMem_Free(indices);
1671 Py_XDECREF(pools);
1672 return NULL;
1673}
1674
1675static void
1676product_dealloc(productobject *lz)
1677{
1678 PyObject_GC_UnTrack(lz);
1679 Py_XDECREF(lz->pools);
1680 Py_XDECREF(lz->result);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001681 PyMem_Free(lz->indices);
1682 Py_TYPE(lz)->tp_free(lz);
1683}
1684
1685static int
1686product_traverse(productobject *lz, visitproc visit, void *arg)
1687{
1688 Py_VISIT(lz->pools);
1689 Py_VISIT(lz->result);
1690 return 0;
1691}
1692
1693static PyObject *
1694product_next(productobject *lz)
1695{
1696 PyObject *pool;
1697 PyObject *elem;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001698 PyObject *oldelem;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001699 PyObject *pools = lz->pools;
1700 PyObject *result = lz->result;
1701 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1702 Py_ssize_t i;
1703
1704 if (lz->stopped)
1705 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001706
Christian Heimesc3f30c42008-02-22 16:37:40 +00001707 if (result == NULL) {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001708 /* On the first pass, return an initial tuple filled with the
Christian Heimes78644762008-03-04 23:39:23 +00001709 first element from each pool. */
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001710 result = PyTuple_New(npools);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001711 if (result == NULL)
1712 goto empty;
1713 lz->result = result;
1714 for (i=0; i < npools; i++) {
1715 pool = PyTuple_GET_ITEM(pools, i);
1716 if (PyTuple_GET_SIZE(pool) == 0)
1717 goto empty;
1718 elem = PyTuple_GET_ITEM(pool, 0);
1719 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001720 PyTuple_SET_ITEM(result, i, elem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001721 }
1722 } else {
1723 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001724
1725 /* Copy the previous result tuple or re-use it if available */
1726 if (Py_REFCNT(result) > 1) {
1727 PyObject *old_result = result;
1728 result = PyTuple_New(npools);
1729 if (result == NULL)
1730 goto empty;
1731 lz->result = result;
1732 for (i=0; i < npools; i++) {
1733 elem = PyTuple_GET_ITEM(old_result, i);
1734 Py_INCREF(elem);
1735 PyTuple_SET_ITEM(result, i, elem);
1736 }
1737 Py_DECREF(old_result);
1738 }
1739 /* Now, we've got the only copy so we can update it in-place */
Christian Heimesb558a2e2008-03-02 22:46:37 +00001740 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001741
1742 /* Update the pool indices right-to-left. Only advance to the
1743 next pool when the previous one rolls-over */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001744 for (i=npools-1 ; i >= 0 ; i--) {
1745 pool = PyTuple_GET_ITEM(pools, i);
1746 indices[i]++;
Christian Heimesb558a2e2008-03-02 22:46:37 +00001747 if (indices[i] == PyTuple_GET_SIZE(pool)) {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001748 /* Roll-over and advance to next pool */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001749 indices[i] = 0;
1750 elem = PyTuple_GET_ITEM(pool, 0);
1751 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001752 oldelem = PyTuple_GET_ITEM(result, i);
1753 PyTuple_SET_ITEM(result, i, elem);
1754 Py_DECREF(oldelem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001755 } else {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001756 /* No rollover. Just increment and stop here. */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001757 elem = PyTuple_GET_ITEM(pool, indices[i]);
1758 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001759 oldelem = PyTuple_GET_ITEM(result, i);
1760 PyTuple_SET_ITEM(result, i, elem);
1761 Py_DECREF(oldelem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001762 break;
1763 }
1764 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001765
1766 /* If i is negative, then the indices have all rolled-over
1767 and we're done. */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001768 if (i < 0)
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001769 goto empty;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001770 }
1771
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001772 Py_INCREF(result);
1773 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001774
1775empty:
1776 lz->stopped = 1;
1777 return NULL;
1778}
1779
1780PyDoc_STRVAR(product_doc,
1781"product(*iterables) --> product object\n\
1782\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001783Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001784For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1785The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1786cycle in a manner similar to an odometer (with the rightmost element changing\n\
1787on every iteration).\n\n\
1788product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1789product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1790
1791static PyTypeObject product_type = {
1792 PyVarObject_HEAD_INIT(NULL, 0)
1793 "itertools.product", /* tp_name */
1794 sizeof(productobject), /* tp_basicsize */
1795 0, /* tp_itemsize */
1796 /* methods */
1797 (destructor)product_dealloc, /* tp_dealloc */
1798 0, /* tp_print */
1799 0, /* tp_getattr */
1800 0, /* tp_setattr */
1801 0, /* tp_compare */
1802 0, /* tp_repr */
1803 0, /* tp_as_number */
1804 0, /* tp_as_sequence */
1805 0, /* tp_as_mapping */
1806 0, /* tp_hash */
1807 0, /* tp_call */
1808 0, /* tp_str */
1809 PyObject_GenericGetAttr, /* tp_getattro */
1810 0, /* tp_setattro */
1811 0, /* tp_as_buffer */
1812 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1813 Py_TPFLAGS_BASETYPE, /* tp_flags */
1814 product_doc, /* tp_doc */
1815 (traverseproc)product_traverse, /* tp_traverse */
1816 0, /* tp_clear */
1817 0, /* tp_richcompare */
1818 0, /* tp_weaklistoffset */
1819 PyObject_SelfIter, /* tp_iter */
1820 (iternextfunc)product_next, /* tp_iternext */
1821 0, /* tp_methods */
1822 0, /* tp_members */
1823 0, /* tp_getset */
1824 0, /* tp_base */
1825 0, /* tp_dict */
1826 0, /* tp_descr_get */
1827 0, /* tp_descr_set */
1828 0, /* tp_dictoffset */
1829 0, /* tp_init */
1830 0, /* tp_alloc */
1831 product_new, /* tp_new */
1832 PyObject_GC_Del, /* tp_free */
1833};
1834
1835
Christian Heimes380f7f22008-02-28 11:19:05 +00001836/* combinations object ************************************************************/
1837
1838typedef struct {
1839 PyObject_HEAD
1840 PyObject *pool; /* input converted to a tuple */
1841 Py_ssize_t *indices; /* one index per result element */
1842 PyObject *result; /* most recently returned result tuple */
1843 Py_ssize_t r; /* size of result tuple */
1844 int stopped; /* set to 1 when the combinations iterator is exhausted */
1845} combinationsobject;
1846
1847static PyTypeObject combinations_type;
1848
1849static PyObject *
1850combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1851{
1852 combinationsobject *co;
1853 Py_ssize_t n;
1854 Py_ssize_t r;
1855 PyObject *pool = NULL;
1856 PyObject *iterable = NULL;
1857 Py_ssize_t *indices = NULL;
1858 Py_ssize_t i;
1859 static char *kwargs[] = {"iterable", "r", NULL};
1860
1861 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
1862 &iterable, &r))
1863 return NULL;
1864
1865 pool = PySequence_Tuple(iterable);
1866 if (pool == NULL)
1867 goto error;
1868 n = PyTuple_GET_SIZE(pool);
1869 if (r < 0) {
1870 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
1871 goto error;
1872 }
1873 if (r > n) {
1874 PyErr_SetString(PyExc_ValueError, "r cannot be bigger than the iterable");
1875 goto error;
1876 }
1877
1878 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
1879 if (indices == NULL) {
1880 PyErr_NoMemory();
1881 goto error;
1882 }
1883
1884 for (i=0 ; i<r ; i++)
1885 indices[i] = i;
1886
1887 /* create combinationsobject structure */
1888 co = (combinationsobject *)type->tp_alloc(type, 0);
1889 if (co == NULL)
1890 goto error;
1891
1892 co->pool = pool;
1893 co->indices = indices;
1894 co->result = NULL;
1895 co->r = r;
1896 co->stopped = 0;
1897
1898 return (PyObject *)co;
1899
1900error:
1901 if (indices != NULL)
1902 PyMem_Free(indices);
1903 Py_XDECREF(pool);
1904 return NULL;
1905}
1906
1907static void
1908combinations_dealloc(combinationsobject *co)
1909{
1910 PyObject_GC_UnTrack(co);
1911 Py_XDECREF(co->pool);
1912 Py_XDECREF(co->result);
1913 PyMem_Free(co->indices);
1914 Py_TYPE(co)->tp_free(co);
1915}
1916
1917static int
1918combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
1919{
1920 Py_VISIT(co->pool);
1921 Py_VISIT(co->result);
1922 return 0;
1923}
1924
1925static PyObject *
1926combinations_next(combinationsobject *co)
1927{
1928 PyObject *elem;
1929 PyObject *oldelem;
1930 PyObject *pool = co->pool;
1931 Py_ssize_t *indices = co->indices;
1932 PyObject *result = co->result;
1933 Py_ssize_t n = PyTuple_GET_SIZE(pool);
1934 Py_ssize_t r = co->r;
1935 Py_ssize_t i, j, index;
1936
1937 if (co->stopped)
1938 return NULL;
1939
1940 if (result == NULL) {
1941 /* On the first pass, initialize result tuple using the indices */
1942 result = PyTuple_New(r);
1943 if (result == NULL)
1944 goto empty;
1945 co->result = result;
1946 for (i=0; i<r ; i++) {
1947 index = indices[i];
1948 elem = PyTuple_GET_ITEM(pool, index);
1949 Py_INCREF(elem);
1950 PyTuple_SET_ITEM(result, i, elem);
1951 }
1952 } else {
1953 /* Copy the previous result tuple or re-use it if available */
1954 if (Py_REFCNT(result) > 1) {
1955 PyObject *old_result = result;
1956 result = PyTuple_New(r);
1957 if (result == NULL)
1958 goto empty;
1959 co->result = result;
1960 for (i=0; i<r ; i++) {
1961 elem = PyTuple_GET_ITEM(old_result, i);
1962 Py_INCREF(elem);
1963 PyTuple_SET_ITEM(result, i, elem);
1964 }
1965 Py_DECREF(old_result);
1966 }
1967 /* Now, we've got the only copy so we can update it in-place
1968 * CPython's empty tuple is a singleton and cached in
1969 * PyTuple's freelist.
1970 */
1971 assert(r == 0 || Py_REFCNT(result) == 1);
1972
1973 /* Scan indices right-to-left until finding one that is not
1974 at its maximum (i + n - r). */
1975 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
1976 ;
1977
1978 /* If i is negative, then the indices are all at
1979 their maximum value and we're done. */
1980 if (i < 0)
1981 goto empty;
1982
1983 /* Increment the current index which we know is not at its
1984 maximum. Then move back to the right setting each index
1985 to its lowest possible value (one higher than the index
1986 to its left -- this maintains the sort order invariant). */
1987 indices[i]++;
1988 for (j=i+1 ; j<r ; j++)
1989 indices[j] = indices[j-1] + 1;
1990
1991 /* Update the result tuple for the new indices
1992 starting with i, the leftmost index that changed */
1993 for ( ; i<r ; i++) {
1994 index = indices[i];
1995 elem = PyTuple_GET_ITEM(pool, index);
1996 Py_INCREF(elem);
1997 oldelem = PyTuple_GET_ITEM(result, i);
1998 PyTuple_SET_ITEM(result, i, elem);
1999 Py_DECREF(oldelem);
2000 }
2001 }
2002
2003 Py_INCREF(result);
2004 return result;
2005
2006empty:
2007 co->stopped = 1;
2008 return NULL;
2009}
2010
2011PyDoc_STRVAR(combinations_doc,
2012"combinations(iterables) --> combinations object\n\
2013\n\
2014Return successive r-length combinations of elements in the iterable.\n\n\
2015combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2016
2017static PyTypeObject combinations_type = {
2018 PyVarObject_HEAD_INIT(NULL, 0)
2019 "itertools.combinations", /* tp_name */
2020 sizeof(combinationsobject), /* tp_basicsize */
2021 0, /* tp_itemsize */
2022 /* methods */
2023 (destructor)combinations_dealloc, /* tp_dealloc */
2024 0, /* tp_print */
2025 0, /* tp_getattr */
2026 0, /* tp_setattr */
2027 0, /* tp_compare */
2028 0, /* tp_repr */
2029 0, /* tp_as_number */
2030 0, /* tp_as_sequence */
2031 0, /* tp_as_mapping */
2032 0, /* tp_hash */
2033 0, /* tp_call */
2034 0, /* tp_str */
2035 PyObject_GenericGetAttr, /* tp_getattro */
2036 0, /* tp_setattro */
2037 0, /* tp_as_buffer */
2038 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2039 Py_TPFLAGS_BASETYPE, /* tp_flags */
2040 combinations_doc, /* tp_doc */
2041 (traverseproc)combinations_traverse, /* tp_traverse */
2042 0, /* tp_clear */
2043 0, /* tp_richcompare */
2044 0, /* tp_weaklistoffset */
2045 PyObject_SelfIter, /* tp_iter */
2046 (iternextfunc)combinations_next, /* tp_iternext */
2047 0, /* tp_methods */
2048 0, /* tp_members */
2049 0, /* tp_getset */
2050 0, /* tp_base */
2051 0, /* tp_dict */
2052 0, /* tp_descr_get */
2053 0, /* tp_descr_set */
2054 0, /* tp_dictoffset */
2055 0, /* tp_init */
2056 0, /* tp_alloc */
2057 combinations_new, /* tp_new */
2058 PyObject_GC_Del, /* tp_free */
2059};
2060
2061
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002062/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00002063
2064typedef struct {
2065 PyObject_HEAD
2066 PyObject *func;
2067 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002068} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002069
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002070static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002071
2072static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002073filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002074{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00002075 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002076 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002077 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002078
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002079 if (type == &filterfalse_type &&
2080 !_PyArg_NoKeywords("filterfalse()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002081 return NULL;
2082
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002083 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
Raymond Hettinger60eca932003-02-09 06:40:58 +00002084 return NULL;
2085
2086 /* Get iterator. */
2087 it = PyObject_GetIter(seq);
2088 if (it == NULL)
2089 return NULL;
2090
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002091 /* create filterfalseobject structure */
2092 lz = (filterfalseobject *)type->tp_alloc(type, 0);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002093 if (lz == NULL) {
2094 Py_DECREF(it);
2095 return NULL;
2096 }
2097 Py_INCREF(func);
2098 lz->func = func;
2099 lz->it = it;
2100
2101 return (PyObject *)lz;
2102}
2103
2104static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002105filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002106{
2107 PyObject_GC_UnTrack(lz);
2108 Py_XDECREF(lz->func);
2109 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00002110 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002111}
2112
2113static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002114filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002115{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002116 Py_VISIT(lz->it);
2117 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002118 return 0;
2119}
2120
2121static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002122filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002123{
2124 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002125 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002126 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002127 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002128
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002129 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00002130 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002131 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002132 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002133 if (item == NULL)
2134 return NULL;
2135
Christian Heimes836baa52008-02-26 08:18:30 +00002136 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
Raymond Hettinger60eca932003-02-09 06:40:58 +00002137 ok = PyObject_IsTrue(item);
2138 } else {
2139 PyObject *good;
2140 good = PyObject_CallFunctionObjArgs(lz->func,
2141 item, NULL);
2142 if (good == NULL) {
2143 Py_DECREF(item);
2144 return NULL;
2145 }
2146 ok = PyObject_IsTrue(good);
2147 Py_DECREF(good);
2148 }
2149 if (!ok)
2150 return item;
2151 Py_DECREF(item);
2152 }
2153}
2154
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002155PyDoc_STRVAR(filterfalse_doc,
2156"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002157\n\
2158Return those items of sequence for which function(item) is false.\n\
2159If function is None, return the items that are false.");
2160
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002161static PyTypeObject filterfalse_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002162 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002163 "itertools.filterfalse", /* tp_name */
2164 sizeof(filterfalseobject), /* tp_basicsize */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002165 0, /* tp_itemsize */
2166 /* methods */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002167 (destructor)filterfalse_dealloc, /* tp_dealloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002168 0, /* tp_print */
2169 0, /* tp_getattr */
2170 0, /* tp_setattr */
2171 0, /* tp_compare */
2172 0, /* tp_repr */
2173 0, /* tp_as_number */
2174 0, /* tp_as_sequence */
2175 0, /* tp_as_mapping */
2176 0, /* tp_hash */
2177 0, /* tp_call */
2178 0, /* tp_str */
2179 PyObject_GenericGetAttr, /* tp_getattro */
2180 0, /* tp_setattro */
2181 0, /* tp_as_buffer */
2182 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2183 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002184 filterfalse_doc, /* tp_doc */
2185 (traverseproc)filterfalse_traverse, /* tp_traverse */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002186 0, /* tp_clear */
2187 0, /* tp_richcompare */
2188 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002189 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002190 (iternextfunc)filterfalse_next, /* tp_iternext */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002191 0, /* tp_methods */
2192 0, /* tp_members */
2193 0, /* tp_getset */
2194 0, /* tp_base */
2195 0, /* tp_dict */
2196 0, /* tp_descr_get */
2197 0, /* tp_descr_set */
2198 0, /* tp_dictoffset */
2199 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002200 0, /* tp_alloc */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002201 filterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002202 PyObject_GC_Del, /* tp_free */
2203};
2204
2205
2206/* count object ************************************************************/
2207
2208typedef struct {
2209 PyObject_HEAD
Thomas Wouters477c8d52006-05-27 19:21:47 +00002210 Py_ssize_t cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002211 PyObject *long_cnt; /* Arbitrarily large count when cnt >= PY_SSIZE_T_MAX */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002212} countobject;
2213
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002214static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002215
2216static PyObject *
2217count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2218{
2219 countobject *lz;
Thomas Wouters477c8d52006-05-27 19:21:47 +00002220 Py_ssize_t cnt = 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002221 PyObject *cnt_arg = NULL;
2222 PyObject *long_cnt = NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002223
Thomas Woutersb2137042007-02-01 18:02:27 +00002224 if (type == &count_type && !_PyArg_NoKeywords("count()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002225 return NULL;
2226
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002227 if (!PyArg_UnpackTuple(args, "count", 0, 1, &cnt_arg))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002228 return NULL;
2229
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002230 if (cnt_arg != NULL) {
Christian Heimes217cfd12007-12-02 14:31:20 +00002231 cnt = PyLong_AsSsize_t(cnt_arg);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002232 if (cnt == -1 && PyErr_Occurred()) {
2233 PyErr_Clear();
2234 if (!PyLong_Check(cnt_arg)) {
2235 PyErr_SetString(PyExc_TypeError, "an integer is required");
2236 return NULL;
2237 }
2238 long_cnt = cnt_arg;
2239 Py_INCREF(long_cnt);
2240 cnt = PY_SSIZE_T_MAX;
2241 }
2242 }
2243
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002244 /* create countobject structure */
2245 lz = (countobject *)PyObject_New(countobject, &count_type);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002246 if (lz == NULL) {
2247 Py_XDECREF(long_cnt);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002248 return NULL;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002249 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002250 lz->cnt = cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002251 lz->long_cnt = long_cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002252
2253 return (PyObject *)lz;
2254}
2255
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002256static void
2257count_dealloc(countobject *lz)
2258{
2259 Py_XDECREF(lz->long_cnt);
2260 PyObject_Del(lz);
2261}
2262
2263static PyObject *
2264count_nextlong(countobject *lz)
2265{
2266 static PyObject *one = NULL;
2267 PyObject *cnt;
2268 PyObject *stepped_up;
2269
2270 if (lz->long_cnt == NULL) {
Christian Heimes217cfd12007-12-02 14:31:20 +00002271 lz->long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002272 if (lz->long_cnt == NULL)
2273 return NULL;
2274 }
2275 if (one == NULL) {
Christian Heimes217cfd12007-12-02 14:31:20 +00002276 one = PyLong_FromLong(1);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002277 if (one == NULL)
2278 return NULL;
2279 }
2280 cnt = lz->long_cnt;
2281 assert(cnt != NULL);
2282 stepped_up = PyNumber_Add(cnt, one);
2283 if (stepped_up == NULL)
2284 return NULL;
2285 lz->long_cnt = stepped_up;
2286 return cnt;
2287}
2288
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002289static PyObject *
2290count_next(countobject *lz)
2291{
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002292 if (lz->cnt == PY_SSIZE_T_MAX)
2293 return count_nextlong(lz);
Christian Heimes217cfd12007-12-02 14:31:20 +00002294 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002295}
2296
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002297static PyObject *
2298count_repr(countobject *lz)
2299{
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002300 if (lz->cnt != PY_SSIZE_T_MAX)
2301 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
2302
2303 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002304}
2305
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002306PyDoc_STRVAR(count_doc,
2307"count([firstval]) --> count object\n\
2308\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002309Return a count object whose .__next__() method returns consecutive\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002310integers starting from zero or, if specified, from firstval.");
2311
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002312static PyTypeObject count_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002313 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002314 "itertools.count", /* tp_name */
2315 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002316 0, /* tp_itemsize */
2317 /* methods */
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002318 (destructor)count_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002319 0, /* tp_print */
2320 0, /* tp_getattr */
2321 0, /* tp_setattr */
2322 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002323 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002324 0, /* tp_as_number */
2325 0, /* tp_as_sequence */
2326 0, /* tp_as_mapping */
2327 0, /* tp_hash */
2328 0, /* tp_call */
2329 0, /* tp_str */
2330 PyObject_GenericGetAttr, /* tp_getattro */
2331 0, /* tp_setattro */
2332 0, /* tp_as_buffer */
2333 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002334 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002335 0, /* tp_traverse */
2336 0, /* tp_clear */
2337 0, /* tp_richcompare */
2338 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002339 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002340 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002341 0, /* tp_methods */
2342 0, /* tp_members */
2343 0, /* tp_getset */
2344 0, /* tp_base */
2345 0, /* tp_dict */
2346 0, /* tp_descr_get */
2347 0, /* tp_descr_set */
2348 0, /* tp_dictoffset */
2349 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002350 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002351 count_new, /* tp_new */
2352};
2353
2354
2355/* izip object ************************************************************/
2356
2357#include "Python.h"
2358
2359typedef struct {
2360 PyObject_HEAD
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002361 Py_ssize_t tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002362 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002363 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002364} izipobject;
2365
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002366static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002367
2368static PyObject *
2369izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2370{
2371 izipobject *lz;
Thomas Wouters477c8d52006-05-27 19:21:47 +00002372 Py_ssize_t i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002373 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002374 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002375 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002376
Thomas Woutersb2137042007-02-01 18:02:27 +00002377 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002378 return NULL;
2379
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002380 /* args must be a tuple */
2381 assert(PyTuple_Check(args));
2382
2383 /* obtain iterators */
2384 ittuple = PyTuple_New(tuplesize);
Thomas Wouters477c8d52006-05-27 19:21:47 +00002385 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002386 return NULL;
2387 for (i=0; i < tuplesize; ++i) {
2388 PyObject *item = PyTuple_GET_ITEM(args, i);
2389 PyObject *it = PyObject_GetIter(item);
2390 if (it == NULL) {
2391 if (PyErr_ExceptionMatches(PyExc_TypeError))
2392 PyErr_Format(PyExc_TypeError,
Thomas Wouters477c8d52006-05-27 19:21:47 +00002393 "izip argument #%zd must support iteration",
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002394 i+1);
2395 Py_DECREF(ittuple);
2396 return NULL;
2397 }
2398 PyTuple_SET_ITEM(ittuple, i, it);
2399 }
2400
Raymond Hettinger2012f172003-02-07 05:32:58 +00002401 /* create a result holder */
2402 result = PyTuple_New(tuplesize);
2403 if (result == NULL) {
2404 Py_DECREF(ittuple);
2405 return NULL;
2406 }
2407 for (i=0 ; i < tuplesize ; i++) {
2408 Py_INCREF(Py_None);
2409 PyTuple_SET_ITEM(result, i, Py_None);
2410 }
2411
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002412 /* create izipobject structure */
2413 lz = (izipobject *)type->tp_alloc(type, 0);
2414 if (lz == NULL) {
2415 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002416 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002417 return NULL;
2418 }
2419 lz->ittuple = ittuple;
2420 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002421 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002422
2423 return (PyObject *)lz;
2424}
2425
2426static void
2427izip_dealloc(izipobject *lz)
2428{
2429 PyObject_GC_UnTrack(lz);
2430 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002431 Py_XDECREF(lz->result);
Christian Heimes90aa7642007-12-19 02:45:37 +00002432 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002433}
2434
2435static int
2436izip_traverse(izipobject *lz, visitproc visit, void *arg)
2437{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002438 Py_VISIT(lz->ittuple);
2439 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002440 return 0;
2441}
2442
2443static PyObject *
2444izip_next(izipobject *lz)
2445{
Thomas Wouters477c8d52006-05-27 19:21:47 +00002446 Py_ssize_t i;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002447 Py_ssize_t tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002448 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002449 PyObject *it;
2450 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002451 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002452
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002453 if (tuplesize == 0)
2454 return NULL;
Christian Heimes90aa7642007-12-19 02:45:37 +00002455 if (Py_REFCNT(result) == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002456 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002457 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002458 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002459 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00002460 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002461 if (item == NULL) {
2462 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002463 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002464 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002465 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002466 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002467 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002468 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002469 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002470 result = PyTuple_New(tuplesize);
2471 if (result == NULL)
2472 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002473 for (i=0 ; i < tuplesize ; i++) {
2474 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002475 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00002476 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002477 if (item == NULL) {
2478 Py_DECREF(result);
2479 return NULL;
2480 }
2481 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002482 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002483 }
2484 return result;
2485}
2486
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002487PyDoc_STRVAR(izip_doc,
2488"izip(iter1 [,iter2 [...]]) --> izip object\n\
2489\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002490Return a izip object whose .__next__() method returns a tuple where\n\
2491the i-th element comes from the i-th iterable argument. The .__next__()\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002492method continues until the shortest iterable in the argument sequence\n\
2493is exhausted and then it raises StopIteration. Works like the zip()\n\
2494function but consumes less memory by returning an iterator instead of\n\
2495a list.");
2496
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002497static PyTypeObject izip_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002498 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002499 "itertools.izip", /* tp_name */
2500 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002501 0, /* tp_itemsize */
2502 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002503 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002504 0, /* tp_print */
2505 0, /* tp_getattr */
2506 0, /* tp_setattr */
2507 0, /* tp_compare */
2508 0, /* tp_repr */
2509 0, /* tp_as_number */
2510 0, /* tp_as_sequence */
2511 0, /* tp_as_mapping */
2512 0, /* tp_hash */
2513 0, /* tp_call */
2514 0, /* tp_str */
2515 PyObject_GenericGetAttr, /* tp_getattro */
2516 0, /* tp_setattro */
2517 0, /* tp_as_buffer */
2518 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2519 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002520 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002521 (traverseproc)izip_traverse, /* tp_traverse */
2522 0, /* tp_clear */
2523 0, /* tp_richcompare */
2524 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002525 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002526 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002527 0, /* tp_methods */
2528 0, /* tp_members */
2529 0, /* tp_getset */
2530 0, /* tp_base */
2531 0, /* tp_dict */
2532 0, /* tp_descr_get */
2533 0, /* tp_descr_set */
2534 0, /* tp_dictoffset */
2535 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002536 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002537 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002538 PyObject_GC_Del, /* tp_free */
2539};
2540
2541
2542/* repeat object ************************************************************/
2543
2544typedef struct {
2545 PyObject_HEAD
2546 PyObject *element;
Thomas Wouters477c8d52006-05-27 19:21:47 +00002547 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002548} repeatobject;
2549
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002550static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002551
2552static PyObject *
2553repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2554{
2555 repeatobject *ro;
2556 PyObject *element;
Thomas Wouters477c8d52006-05-27 19:21:47 +00002557 Py_ssize_t cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002558
Thomas Woutersb2137042007-02-01 18:02:27 +00002559 if (type == &repeat_type && !_PyArg_NoKeywords("repeat()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002560 return NULL;
2561
Thomas Wouters477c8d52006-05-27 19:21:47 +00002562 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002563 return NULL;
2564
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002565 if (PyTuple_Size(args) == 2 && cnt < 0)
2566 cnt = 0;
2567
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002568 ro = (repeatobject *)type->tp_alloc(type, 0);
2569 if (ro == NULL)
2570 return NULL;
2571 Py_INCREF(element);
2572 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002573 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002574 return (PyObject *)ro;
2575}
2576
2577static void
2578repeat_dealloc(repeatobject *ro)
2579{
2580 PyObject_GC_UnTrack(ro);
2581 Py_XDECREF(ro->element);
Christian Heimes90aa7642007-12-19 02:45:37 +00002582 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002583}
2584
2585static int
2586repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2587{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002588 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002589 return 0;
2590}
2591
2592static PyObject *
2593repeat_next(repeatobject *ro)
2594{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002595 if (ro->cnt == 0)
2596 return NULL;
2597 if (ro->cnt > 0)
2598 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002599 Py_INCREF(ro->element);
2600 return ro->element;
2601}
2602
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002603static PyObject *
2604repeat_repr(repeatobject *ro)
2605{
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002606 if (ro->cnt == -1)
Walter Dörwald7569dfe2007-05-19 21:49:49 +00002607 return PyUnicode_FromFormat("repeat(%R)", ro->element);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002608 else
Walter Dörwald7569dfe2007-05-19 21:49:49 +00002609 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002610}
2611
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002612static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002613repeat_len(repeatobject *ro)
2614{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002615 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002616 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002617 return NULL;
2618 }
Christian Heimes217cfd12007-12-02 14:31:20 +00002619 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002620}
2621
Armin Rigof5b3e362006-02-11 21:32:43 +00002622PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002623
2624static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00002625 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002626 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002627};
2628
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002629PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002630"repeat(element [,times]) -> create an iterator which returns the element\n\
2631for the specified number of times. If not specified, returns the element\n\
2632endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002633
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002634static PyTypeObject repeat_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002635 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002636 "itertools.repeat", /* tp_name */
2637 sizeof(repeatobject), /* tp_basicsize */
2638 0, /* tp_itemsize */
2639 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002640 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002641 0, /* tp_print */
2642 0, /* tp_getattr */
2643 0, /* tp_setattr */
2644 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002645 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002646 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002647 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002648 0, /* tp_as_mapping */
2649 0, /* tp_hash */
2650 0, /* tp_call */
2651 0, /* tp_str */
2652 PyObject_GenericGetAttr, /* tp_getattro */
2653 0, /* tp_setattro */
2654 0, /* tp_as_buffer */
2655 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2656 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002657 repeat_doc, /* tp_doc */
2658 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002659 0, /* tp_clear */
2660 0, /* tp_richcompare */
2661 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002662 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002663 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002664 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002665 0, /* tp_members */
2666 0, /* tp_getset */
2667 0, /* tp_base */
2668 0, /* tp_dict */
2669 0, /* tp_descr_get */
2670 0, /* tp_descr_set */
2671 0, /* tp_dictoffset */
2672 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002673 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002674 repeat_new, /* tp_new */
2675 PyObject_GC_Del, /* tp_free */
2676};
2677
Thomas Wouterscf297e42007-02-23 15:07:44 +00002678/* iziplongest object ************************************************************/
2679
2680#include "Python.h"
2681
2682typedef struct {
2683 PyObject_HEAD
2684 Py_ssize_t tuplesize;
2685 Py_ssize_t numactive;
2686 PyObject *ittuple; /* tuple of iterators */
2687 PyObject *result;
2688 PyObject *fillvalue;
2689} iziplongestobject;
2690
2691static PyTypeObject iziplongest_type;
2692
2693static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002694zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00002695{
2696 iziplongestobject *lz;
2697 Py_ssize_t i;
2698 PyObject *ittuple; /* tuple of iterators */
2699 PyObject *result;
2700 PyObject *fillvalue = Py_None;
2701 Py_ssize_t tuplesize = PySequence_Length(args);
2702
2703 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
2704 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
2705 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
2706 PyErr_SetString(PyExc_TypeError,
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002707 "zip_longest() got an unexpected keyword argument");
Thomas Wouterscf297e42007-02-23 15:07:44 +00002708 return NULL;
2709 }
2710 }
2711
2712 /* args must be a tuple */
2713 assert(PyTuple_Check(args));
2714
2715 /* obtain iterators */
2716 ittuple = PyTuple_New(tuplesize);
2717 if (ittuple == NULL)
2718 return NULL;
2719 for (i=0; i < tuplesize; ++i) {
2720 PyObject *item = PyTuple_GET_ITEM(args, i);
2721 PyObject *it = PyObject_GetIter(item);
2722 if (it == NULL) {
2723 if (PyErr_ExceptionMatches(PyExc_TypeError))
2724 PyErr_Format(PyExc_TypeError,
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002725 "zip_longest argument #%zd must support iteration",
Thomas Wouterscf297e42007-02-23 15:07:44 +00002726 i+1);
2727 Py_DECREF(ittuple);
2728 return NULL;
2729 }
2730 PyTuple_SET_ITEM(ittuple, i, it);
2731 }
2732
2733 /* create a result holder */
2734 result = PyTuple_New(tuplesize);
2735 if (result == NULL) {
2736 Py_DECREF(ittuple);
2737 return NULL;
2738 }
2739 for (i=0 ; i < tuplesize ; i++) {
2740 Py_INCREF(Py_None);
2741 PyTuple_SET_ITEM(result, i, Py_None);
2742 }
2743
2744 /* create iziplongestobject structure */
2745 lz = (iziplongestobject *)type->tp_alloc(type, 0);
2746 if (lz == NULL) {
2747 Py_DECREF(ittuple);
2748 Py_DECREF(result);
2749 return NULL;
2750 }
2751 lz->ittuple = ittuple;
2752 lz->tuplesize = tuplesize;
2753 lz->numactive = tuplesize;
2754 lz->result = result;
2755 Py_INCREF(fillvalue);
2756 lz->fillvalue = fillvalue;
2757 return (PyObject *)lz;
2758}
2759
2760static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002761zip_longest_dealloc(iziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00002762{
2763 PyObject_GC_UnTrack(lz);
2764 Py_XDECREF(lz->ittuple);
2765 Py_XDECREF(lz->result);
2766 Py_XDECREF(lz->fillvalue);
Christian Heimes90aa7642007-12-19 02:45:37 +00002767 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00002768}
2769
2770static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002771zip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00002772{
2773 Py_VISIT(lz->ittuple);
2774 Py_VISIT(lz->result);
2775 Py_VISIT(lz->fillvalue);
2776 return 0;
2777}
2778
2779static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002780zip_longest_next(iziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00002781{
2782 Py_ssize_t i;
2783 Py_ssize_t tuplesize = lz->tuplesize;
2784 PyObject *result = lz->result;
2785 PyObject *it;
2786 PyObject *item;
2787 PyObject *olditem;
2788
2789 if (tuplesize == 0)
2790 return NULL;
2791 if (lz->numactive == 0)
2792 return NULL;
Christian Heimes90aa7642007-12-19 02:45:37 +00002793 if (Py_REFCNT(result) == 1) {
Thomas Wouterscf297e42007-02-23 15:07:44 +00002794 Py_INCREF(result);
2795 for (i=0 ; i < tuplesize ; i++) {
2796 it = PyTuple_GET_ITEM(lz->ittuple, i);
2797 if (it == NULL) {
2798 Py_INCREF(lz->fillvalue);
2799 item = lz->fillvalue;
2800 } else {
2801 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00002802 item = (*Py_TYPE(it)->tp_iternext)(it);
Thomas Wouterscf297e42007-02-23 15:07:44 +00002803 if (item == NULL) {
2804 lz->numactive -= 1;
2805 if (lz->numactive == 0) {
2806 Py_DECREF(result);
2807 return NULL;
2808 } else {
2809 Py_INCREF(lz->fillvalue);
2810 item = lz->fillvalue;
2811 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
2812 Py_DECREF(it);
2813 }
2814 }
2815 }
2816 olditem = PyTuple_GET_ITEM(result, i);
2817 PyTuple_SET_ITEM(result, i, item);
2818 Py_DECREF(olditem);
2819 }
2820 } else {
2821 result = PyTuple_New(tuplesize);
2822 if (result == NULL)
2823 return NULL;
2824 for (i=0 ; i < tuplesize ; i++) {
2825 it = PyTuple_GET_ITEM(lz->ittuple, i);
2826 if (it == NULL) {
2827 Py_INCREF(lz->fillvalue);
2828 item = lz->fillvalue;
2829 } else {
2830 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00002831 item = (*Py_TYPE(it)->tp_iternext)(it);
Thomas Wouterscf297e42007-02-23 15:07:44 +00002832 if (item == NULL) {
2833 lz->numactive -= 1;
2834 if (lz->numactive == 0) {
2835 Py_DECREF(result);
2836 return NULL;
2837 } else {
2838 Py_INCREF(lz->fillvalue);
2839 item = lz->fillvalue;
2840 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
2841 Py_DECREF(it);
2842 }
2843 }
2844 }
2845 PyTuple_SET_ITEM(result, i, item);
2846 }
2847 }
2848 return result;
2849}
2850
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002851PyDoc_STRVAR(zip_longest_doc,
2852"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00002853\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002854Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002855the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00002856method continues until the longest iterable in the argument sequence\n\
2857is exhausted and then it raises StopIteration. When the shorter iterables\n\
2858are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
2859defaults to None or can be specified by a keyword argument.\n\
2860");
2861
2862static PyTypeObject iziplongest_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002863 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002864 "itertools.zip_longest", /* tp_name */
Thomas Wouterscf297e42007-02-23 15:07:44 +00002865 sizeof(iziplongestobject), /* tp_basicsize */
2866 0, /* tp_itemsize */
2867 /* methods */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002868 (destructor)zip_longest_dealloc, /* tp_dealloc */
Thomas Wouterscf297e42007-02-23 15:07:44 +00002869 0, /* tp_print */
2870 0, /* tp_getattr */
2871 0, /* tp_setattr */
2872 0, /* tp_compare */
2873 0, /* tp_repr */
2874 0, /* tp_as_number */
2875 0, /* tp_as_sequence */
2876 0, /* tp_as_mapping */
2877 0, /* tp_hash */
2878 0, /* tp_call */
2879 0, /* tp_str */
2880 PyObject_GenericGetAttr, /* tp_getattro */
2881 0, /* tp_setattro */
2882 0, /* tp_as_buffer */
2883 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2884 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002885 zip_longest_doc, /* tp_doc */
2886 (traverseproc)zip_longest_traverse, /* tp_traverse */
Thomas Wouterscf297e42007-02-23 15:07:44 +00002887 0, /* tp_clear */
2888 0, /* tp_richcompare */
2889 0, /* tp_weaklistoffset */
2890 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002891 (iternextfunc)zip_longest_next, /* tp_iternext */
Thomas Wouterscf297e42007-02-23 15:07:44 +00002892 0, /* tp_methods */
2893 0, /* tp_members */
2894 0, /* tp_getset */
2895 0, /* tp_base */
2896 0, /* tp_dict */
2897 0, /* tp_descr_get */
2898 0, /* tp_descr_set */
2899 0, /* tp_dictoffset */
2900 0, /* tp_init */
2901 0, /* tp_alloc */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002902 zip_longest_new, /* tp_new */
Thomas Wouterscf297e42007-02-23 15:07:44 +00002903 PyObject_GC_Del, /* tp_free */
2904};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002905
2906/* module level code ********************************************************/
2907
2908PyDoc_STRVAR(module_doc,
2909"Functional tools for creating and using iterators.\n\
2910\n\
2911Infinite iterators:\n\
2912count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002913cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002914repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002915\n\
2916Iterators terminating on the shortest input sequence:\n\
2917izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002918zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
2919filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002920islice(seq, [start,] stop [, step]) --> elements from\n\
2921 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002922starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00002923tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002924chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002925takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2926dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002927groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002928");
2929
2930
Raymond Hettingerad983e72003-11-12 14:32:26 +00002931static PyMethodDef module_methods[] = {
2932 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2933 {NULL, NULL} /* sentinel */
2934};
2935
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002936PyMODINIT_FUNC
2937inititertools(void)
2938{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002939 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002940 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002941 char *name;
2942 PyTypeObject *typelist[] = {
Christian Heimes380f7f22008-02-28 11:19:05 +00002943 &combinations_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002944 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002945 &dropwhile_type,
2946 &takewhile_type,
2947 &islice_type,
2948 &starmap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002949 &chain_type,
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002950 &filterfalse_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002951 &count_type,
2952 &izip_type,
Christian Heimesc3f30c42008-02-22 16:37:40 +00002953 &iziplongest_type,
Christian Heimes380f7f22008-02-28 11:19:05 +00002954 &product_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002955 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002956 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002957 NULL
2958 };
2959
Christian Heimes90aa7642007-12-19 02:45:37 +00002960 Py_TYPE(&teedataobject_type) = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00002961 m = Py_InitModule3("itertools", module_methods, module_doc);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00002962 if (m == NULL)
2963 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002964
Raymond Hettinger60eca932003-02-09 06:40:58 +00002965 for (i=0 ; typelist[i] != NULL ; i++) {
2966 if (PyType_Ready(typelist[i]) < 0)
2967 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002968 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002969 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002970 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002971 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002972 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00002973
2974 if (PyType_Ready(&teedataobject_type) < 0)
2975 return;
2976 if (PyType_Ready(&tee_type) < 0)
2977 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002978 if (PyType_Ready(&_grouper_type) < 0)
2979 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002980}