blob: 8125dcb064204d513c5324d54bc2f97397ba4438 [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
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000201 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000202 if (igo == NULL)
203 return NULL;
204 igo->parent = (PyObject *)parent;
205 Py_INCREF(parent);
206 igo->tgtkey = tgtkey;
207 Py_INCREF(tgtkey);
208
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000209 PyObject_GC_Track(igo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000210 return (PyObject *)igo;
211}
212
213static void
214_grouper_dealloc(_grouperobject *igo)
215{
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000216 PyObject_GC_UnTrack(igo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000217 Py_DECREF(igo->parent);
218 Py_DECREF(igo->tgtkey);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000219 PyObject_GC_Del(igo);
220}
221
222static int
223_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
224{
225 Py_VISIT(igo->parent);
226 Py_VISIT(igo->tgtkey);
227 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000228}
229
230static PyObject *
231_grouper_next(_grouperobject *igo)
232{
233 groupbyobject *gbo = (groupbyobject *)igo->parent;
234 PyObject *newvalue, *newkey, *r;
235 int rcmp;
236
237 if (gbo->currvalue == NULL) {
238 newvalue = PyIter_Next(gbo->it);
239 if (newvalue == NULL)
240 return NULL;
241
242 if (gbo->keyfunc == Py_None) {
243 newkey = newvalue;
244 Py_INCREF(newvalue);
245 } else {
246 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
247 newvalue, NULL);
248 if (newkey == NULL) {
249 Py_DECREF(newvalue);
250 return NULL;
251 }
252 }
253
254 assert(gbo->currkey == NULL);
255 gbo->currkey = newkey;
256 gbo->currvalue = newvalue;
257 }
258
259 assert(gbo->currkey != NULL);
260 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
261 if (rcmp <= 0)
262 /* got any error or current group is end */
263 return NULL;
264
265 r = gbo->currvalue;
266 gbo->currvalue = NULL;
Raymond Hettinger75ccea32004-09-01 07:02:44 +0000267 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000268
269 return r;
270}
271
272static PyTypeObject _grouper_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +0000273 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000274 "itertools._grouper", /* tp_name */
275 sizeof(_grouperobject), /* tp_basicsize */
276 0, /* tp_itemsize */
277 /* methods */
278 (destructor)_grouper_dealloc, /* tp_dealloc */
279 0, /* tp_print */
280 0, /* tp_getattr */
281 0, /* tp_setattr */
282 0, /* tp_compare */
283 0, /* tp_repr */
284 0, /* tp_as_number */
285 0, /* tp_as_sequence */
286 0, /* tp_as_mapping */
287 0, /* tp_hash */
288 0, /* tp_call */
289 0, /* tp_str */
290 PyObject_GenericGetAttr, /* tp_getattro */
291 0, /* tp_setattro */
292 0, /* tp_as_buffer */
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000293 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000294 0, /* tp_doc */
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000295 (traverseproc)_grouper_traverse,/* tp_traverse */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000296 0, /* tp_clear */
297 0, /* tp_richcompare */
298 0, /* tp_weaklistoffset */
299 PyObject_SelfIter, /* tp_iter */
300 (iternextfunc)_grouper_next, /* tp_iternext */
301 0, /* tp_methods */
302 0, /* tp_members */
303 0, /* tp_getset */
304 0, /* tp_base */
305 0, /* tp_dict */
306 0, /* tp_descr_get */
307 0, /* tp_descr_set */
308 0, /* tp_dictoffset */
309 0, /* tp_init */
310 0, /* tp_alloc */
311 0, /* tp_new */
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000312 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000313};
314
315
316
Raymond Hettingerad983e72003-11-12 14:32:26 +0000317/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000318
Raymond Hettingerad983e72003-11-12 14:32:26 +0000319/* The teedataobject pre-allocates space for LINKCELLS number of objects.
320 To help the object fit neatly inside cache lines (space for 16 to 32
321 pointers), the value should be a multiple of 16 minus space for
322 the other structure members including PyHEAD overhead. The larger the
323 value, the less memory overhead per object and the less time spent
324 allocating/deallocating new links. The smaller the number, the less
325 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000326*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000327#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000328
329typedef struct {
330 PyObject_HEAD
331 PyObject *it;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000332 int numread;
333 PyObject *nextlink;
334 PyObject *(values[LINKCELLS]);
335} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000336
337typedef struct {
338 PyObject_HEAD
Raymond Hettingerad983e72003-11-12 14:32:26 +0000339 teedataobject *dataobj;
340 int index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000341 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000342} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000343
Raymond Hettingerad983e72003-11-12 14:32:26 +0000344static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000345
346static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000347teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000348{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000349 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000350
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000351 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000352 if (tdo == NULL)
Raymond Hettinger45143692003-10-25 06:37:47 +0000353 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000354
355 tdo->numread = 0;
356 tdo->nextlink = NULL;
357 Py_INCREF(it);
358 tdo->it = it;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000359 PyObject_GC_Track(tdo);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000360 return (PyObject *)tdo;
361}
362
363static PyObject *
364teedataobject_jumplink(teedataobject *tdo)
365{
366 if (tdo->nextlink == NULL)
367 tdo->nextlink = teedataobject_new(tdo->it);
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000368 Py_XINCREF(tdo->nextlink);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000369 return tdo->nextlink;
370}
371
372static PyObject *
373teedataobject_getitem(teedataobject *tdo, int i)
374{
375 PyObject *value;
376
377 assert(i < LINKCELLS);
378 if (i < tdo->numread)
379 value = tdo->values[i];
380 else {
381 /* this is the lead iterator, so fetch more data */
382 assert(i == tdo->numread);
383 value = PyIter_Next(tdo->it);
384 if (value == NULL)
385 return NULL;
386 tdo->numread++;
387 tdo->values[i] = value;
Raymond Hettinger45143692003-10-25 06:37:47 +0000388 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000389 Py_INCREF(value);
390 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000391}
392
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000393static int
394teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
395{
396 int i;
397 Py_VISIT(tdo->it);
398 for (i = 0; i < tdo->numread; i++)
399 Py_VISIT(tdo->values[i]);
400 Py_VISIT(tdo->nextlink);
401 return 0;
402}
403
404static int
405teedataobject_clear(teedataobject *tdo)
406{
407 int i;
408 Py_CLEAR(tdo->it);
409 for (i=0 ; i<tdo->numread ; i++)
410 Py_CLEAR(tdo->values[i]);
411 Py_CLEAR(tdo->nextlink);
412 return 0;
413}
414
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000416teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417{
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000418 PyObject_GC_UnTrack(tdo);
419 teedataobject_clear(tdo);
420 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000421}
422
Raymond Hettingerad983e72003-11-12 14:32:26 +0000423PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000424
Raymond Hettingerad983e72003-11-12 14:32:26 +0000425static PyTypeObject teedataobject_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +0000426 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000427 "itertools.tee_dataobject", /* tp_name */
428 sizeof(teedataobject), /* tp_basicsize */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000429 0, /* tp_itemsize */
430 /* methods */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000431 (destructor)teedataobject_dealloc, /* tp_dealloc */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000432 0, /* tp_print */
433 0, /* tp_getattr */
434 0, /* tp_setattr */
435 0, /* tp_compare */
436 0, /* tp_repr */
437 0, /* tp_as_number */
438 0, /* tp_as_sequence */
439 0, /* tp_as_mapping */
440 0, /* tp_hash */
441 0, /* tp_call */
442 0, /* tp_str */
443 PyObject_GenericGetAttr, /* tp_getattro */
444 0, /* tp_setattro */
445 0, /* tp_as_buffer */
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000446 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000447 teedataobject_doc, /* tp_doc */
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000448 (traverseproc)teedataobject_traverse, /* tp_traverse */
449 (inquiry)teedataobject_clear, /* tp_clear */
450 0, /* tp_richcompare */
451 0, /* tp_weaklistoffset */
452 0, /* tp_iter */
453 0, /* tp_iternext */
454 0, /* tp_methods */
455 0, /* tp_members */
456 0, /* tp_getset */
457 0, /* tp_base */
458 0, /* tp_dict */
459 0, /* tp_descr_get */
460 0, /* tp_descr_set */
461 0, /* tp_dictoffset */
462 0, /* tp_init */
463 0, /* tp_alloc */
464 0, /* tp_new */
465 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000466};
467
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000468
469static PyTypeObject tee_type;
470
471static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000472tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000473{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000474 PyObject *value, *link;
475
476 if (to->index >= LINKCELLS) {
477 link = teedataobject_jumplink(to->dataobj);
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000478 Py_DECREF(to->dataobj);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000479 to->dataobj = (teedataobject *)link;
480 to->index = 0;
481 }
482 value = teedataobject_getitem(to->dataobj, to->index);
483 if (value == NULL)
484 return NULL;
485 to->index++;
486 return value;
487}
488
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000489static int
490tee_traverse(teeobject *to, visitproc visit, void *arg)
491{
492 Py_VISIT((PyObject *)to->dataobj);
493 return 0;
494}
495
Raymond Hettingerad983e72003-11-12 14:32:26 +0000496static PyObject *
497tee_copy(teeobject *to)
498{
499 teeobject *newto;
500
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000501 newto = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000502 if (newto == NULL)
503 return NULL;
504 Py_INCREF(to->dataobj);
505 newto->dataobj = to->dataobj;
506 newto->index = to->index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000507 newto->weakreflist = NULL;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000508 PyObject_GC_Track(newto);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000509 return (PyObject *)newto;
510}
511
512PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
513
514static PyObject *
515tee_fromiterable(PyObject *iterable)
516{
517 teeobject *to;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000518 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000519
520 it = PyObject_GetIter(iterable);
521 if (it == NULL)
522 return NULL;
523 if (PyObject_TypeCheck(it, &tee_type)) {
524 to = (teeobject *)tee_copy((teeobject *)it);
525 goto done;
526 }
527
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000528 to = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000529 if (to == NULL)
530 goto done;
531 to->dataobj = (teedataobject *)teedataobject_new(it);
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000532 if (!to->dataobj) {
533 PyObject_GC_Del(to);
534 to = NULL;
535 goto done;
536 }
537
Raymond Hettingerad983e72003-11-12 14:32:26 +0000538 to->index = 0;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000539 to->weakreflist = NULL;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000540 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000541done:
542 Py_XDECREF(it);
543 return (PyObject *)to;
544}
545
546static PyObject *
547tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
548{
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000549 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000550
551 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
552 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000553 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000554}
555
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000556static int
557tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000558{
Raymond Hettingera9f60922004-10-17 16:40:14 +0000559 if (to->weakreflist != NULL)
560 PyObject_ClearWeakRefs((PyObject *) to);
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000561 Py_CLEAR(to->dataobj);
562 return 0;
563}
564
565static void
566tee_dealloc(teeobject *to)
567{
568 PyObject_GC_UnTrack(to);
569 tee_clear(to);
570 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000571}
572
Raymond Hettingerad983e72003-11-12 14:32:26 +0000573PyDoc_STRVAR(teeobject_doc,
574"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000575
Raymond Hettingerad983e72003-11-12 14:32:26 +0000576static PyMethodDef tee_methods[] = {
577 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
578 {NULL, NULL} /* sentinel */
579};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000580
581static PyTypeObject tee_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +0000582 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000583 "itertools.tee", /* tp_name */
584 sizeof(teeobject), /* tp_basicsize */
585 0, /* tp_itemsize */
586 /* methods */
587 (destructor)tee_dealloc, /* tp_dealloc */
588 0, /* tp_print */
589 0, /* tp_getattr */
590 0, /* tp_setattr */
591 0, /* tp_compare */
592 0, /* tp_repr */
593 0, /* tp_as_number */
594 0, /* tp_as_sequence */
595 0, /* tp_as_mapping */
596 0, /* tp_hash */
597 0, /* tp_call */
598 0, /* tp_str */
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000599 0, /* tp_getattro */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000600 0, /* tp_setattro */
601 0, /* tp_as_buffer */
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000602 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000603 teeobject_doc, /* tp_doc */
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000604 (traverseproc)tee_traverse, /* tp_traverse */
605 (inquiry)tee_clear, /* tp_clear */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000606 0, /* tp_richcompare */
Raymond Hettingera9f60922004-10-17 16:40:14 +0000607 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000608 PyObject_SelfIter, /* tp_iter */
609 (iternextfunc)tee_next, /* tp_iternext */
610 tee_methods, /* tp_methods */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000611 0, /* tp_members */
612 0, /* tp_getset */
613 0, /* tp_base */
614 0, /* tp_dict */
615 0, /* tp_descr_get */
616 0, /* tp_descr_set */
617 0, /* tp_dictoffset */
618 0, /* tp_init */
619 0, /* tp_alloc */
620 tee_new, /* tp_new */
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000621 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000622};
623
Raymond Hettingerad983e72003-11-12 14:32:26 +0000624static PyObject *
625tee(PyObject *self, PyObject *args)
626{
Thomas Wouters89f507f2006-12-13 04:49:30 +0000627 Py_ssize_t i, n=2;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000628 PyObject *it, *iterable, *copyable, *result;
629
Thomas Wouters89f507f2006-12-13 04:49:30 +0000630 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000631 return NULL;
Thomas Wouters89f507f2006-12-13 04:49:30 +0000632 if (n < 0) {
633 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
634 return NULL;
635 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000636 result = PyTuple_New(n);
637 if (result == NULL)
638 return NULL;
639 if (n == 0)
640 return result;
641 it = PyObject_GetIter(iterable);
642 if (it == NULL) {
643 Py_DECREF(result);
644 return NULL;
645 }
646 if (!PyObject_HasAttrString(it, "__copy__")) {
647 copyable = tee_fromiterable(it);
648 Py_DECREF(it);
649 if (copyable == NULL) {
650 Py_DECREF(result);
651 return NULL;
652 }
653 } else
654 copyable = it;
655 PyTuple_SET_ITEM(result, 0, copyable);
656 for (i=1 ; i<n ; i++) {
657 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
658 if (copyable == NULL) {
659 Py_DECREF(result);
660 return NULL;
661 }
662 PyTuple_SET_ITEM(result, i, copyable);
663 }
664 return result;
665}
666
667PyDoc_STRVAR(tee_doc,
668"tee(iterable, n=2) --> tuple of n independent iterators.");
669
670
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000671/* cycle object **********************************************************/
672
673typedef struct {
674 PyObject_HEAD
675 PyObject *it;
676 PyObject *saved;
677 int firstpass;
678} cycleobject;
679
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000680static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000681
682static PyObject *
683cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
684{
685 PyObject *it;
686 PyObject *iterable;
687 PyObject *saved;
688 cycleobject *lz;
689
Thomas Woutersb2137042007-02-01 18:02:27 +0000690 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000691 return NULL;
692
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000693 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
694 return NULL;
695
696 /* Get iterator. */
697 it = PyObject_GetIter(iterable);
698 if (it == NULL)
699 return NULL;
700
701 saved = PyList_New(0);
702 if (saved == NULL) {
703 Py_DECREF(it);
704 return NULL;
705 }
706
707 /* create cycleobject structure */
708 lz = (cycleobject *)type->tp_alloc(type, 0);
709 if (lz == NULL) {
710 Py_DECREF(it);
711 Py_DECREF(saved);
712 return NULL;
713 }
714 lz->it = it;
715 lz->saved = saved;
716 lz->firstpass = 0;
717
718 return (PyObject *)lz;
719}
720
721static void
722cycle_dealloc(cycleobject *lz)
723{
724 PyObject_GC_UnTrack(lz);
725 Py_XDECREF(lz->saved);
726 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +0000727 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000728}
729
730static int
731cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
732{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000733 Py_VISIT(lz->it);
734 Py_VISIT(lz->saved);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000735 return 0;
736}
737
738static PyObject *
739cycle_next(cycleobject *lz)
740{
741 PyObject *item;
742 PyObject *it;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000743 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000744
745 while (1) {
746 item = PyIter_Next(lz->it);
747 if (item != NULL) {
748 if (!lz->firstpass)
749 PyList_Append(lz->saved, item);
750 return item;
751 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000752 if (PyErr_Occurred()) {
753 if (PyErr_ExceptionMatches(PyExc_StopIteration))
754 PyErr_Clear();
755 else
756 return NULL;
757 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000758 if (PyList_Size(lz->saved) == 0)
759 return NULL;
760 it = PyObject_GetIter(lz->saved);
761 if (it == NULL)
762 return NULL;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000763 tmp = lz->it;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000764 lz->it = it;
765 lz->firstpass = 1;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000766 Py_DECREF(tmp);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000767 }
768}
769
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000770PyDoc_STRVAR(cycle_doc,
771"cycle(iterable) --> cycle object\n\
772\n\
773Return elements from the iterable until it is exhausted.\n\
774Then repeat the sequence indefinitely.");
775
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000776static PyTypeObject cycle_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +0000777 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000778 "itertools.cycle", /* tp_name */
779 sizeof(cycleobject), /* tp_basicsize */
780 0, /* tp_itemsize */
781 /* methods */
782 (destructor)cycle_dealloc, /* tp_dealloc */
783 0, /* tp_print */
784 0, /* tp_getattr */
785 0, /* tp_setattr */
786 0, /* tp_compare */
787 0, /* tp_repr */
788 0, /* tp_as_number */
789 0, /* tp_as_sequence */
790 0, /* tp_as_mapping */
791 0, /* tp_hash */
792 0, /* tp_call */
793 0, /* tp_str */
794 PyObject_GenericGetAttr, /* tp_getattro */
795 0, /* tp_setattro */
796 0, /* tp_as_buffer */
797 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
798 Py_TPFLAGS_BASETYPE, /* tp_flags */
799 cycle_doc, /* tp_doc */
800 (traverseproc)cycle_traverse, /* tp_traverse */
801 0, /* tp_clear */
802 0, /* tp_richcompare */
803 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000804 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000805 (iternextfunc)cycle_next, /* tp_iternext */
806 0, /* tp_methods */
807 0, /* tp_members */
808 0, /* tp_getset */
809 0, /* tp_base */
810 0, /* tp_dict */
811 0, /* tp_descr_get */
812 0, /* tp_descr_set */
813 0, /* tp_dictoffset */
814 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000815 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000816 cycle_new, /* tp_new */
817 PyObject_GC_Del, /* tp_free */
818};
819
820
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000821/* dropwhile object **********************************************************/
822
823typedef struct {
824 PyObject_HEAD
825 PyObject *func;
826 PyObject *it;
827 long start;
828} dropwhileobject;
829
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000830static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000831
832static PyObject *
833dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
834{
835 PyObject *func, *seq;
836 PyObject *it;
837 dropwhileobject *lz;
838
Thomas Woutersb2137042007-02-01 18:02:27 +0000839 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000840 return NULL;
841
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000842 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
843 return NULL;
844
845 /* Get iterator. */
846 it = PyObject_GetIter(seq);
847 if (it == NULL)
848 return NULL;
849
850 /* create dropwhileobject structure */
851 lz = (dropwhileobject *)type->tp_alloc(type, 0);
852 if (lz == NULL) {
853 Py_DECREF(it);
854 return NULL;
855 }
856 Py_INCREF(func);
857 lz->func = func;
858 lz->it = it;
859 lz->start = 0;
860
861 return (PyObject *)lz;
862}
863
864static void
865dropwhile_dealloc(dropwhileobject *lz)
866{
867 PyObject_GC_UnTrack(lz);
868 Py_XDECREF(lz->func);
869 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +0000870 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000871}
872
873static int
874dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
875{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000876 Py_VISIT(lz->it);
877 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000878 return 0;
879}
880
881static PyObject *
882dropwhile_next(dropwhileobject *lz)
883{
884 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000885 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000886 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000887 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000888
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000889 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +0000890 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000891 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000892 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000893 if (item == NULL)
894 return NULL;
895 if (lz->start == 1)
896 return item;
897
898 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
899 if (good == NULL) {
900 Py_DECREF(item);
901 return NULL;
902 }
903 ok = PyObject_IsTrue(good);
904 Py_DECREF(good);
905 if (!ok) {
906 lz->start = 1;
907 return item;
908 }
909 Py_DECREF(item);
910 }
911}
912
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000913PyDoc_STRVAR(dropwhile_doc,
914"dropwhile(predicate, iterable) --> dropwhile object\n\
915\n\
916Drop items from the iterable while predicate(item) is true.\n\
917Afterwards, return every element until the iterable is exhausted.");
918
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000919static PyTypeObject dropwhile_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +0000920 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000921 "itertools.dropwhile", /* tp_name */
922 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000923 0, /* tp_itemsize */
924 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000925 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000926 0, /* tp_print */
927 0, /* tp_getattr */
928 0, /* tp_setattr */
929 0, /* tp_compare */
930 0, /* tp_repr */
931 0, /* tp_as_number */
932 0, /* tp_as_sequence */
933 0, /* tp_as_mapping */
934 0, /* tp_hash */
935 0, /* tp_call */
936 0, /* tp_str */
937 PyObject_GenericGetAttr, /* tp_getattro */
938 0, /* tp_setattro */
939 0, /* tp_as_buffer */
940 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
941 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000942 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000943 (traverseproc)dropwhile_traverse, /* tp_traverse */
944 0, /* tp_clear */
945 0, /* tp_richcompare */
946 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000947 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000948 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000949 0, /* tp_methods */
950 0, /* tp_members */
951 0, /* tp_getset */
952 0, /* tp_base */
953 0, /* tp_dict */
954 0, /* tp_descr_get */
955 0, /* tp_descr_set */
956 0, /* tp_dictoffset */
957 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000958 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000959 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000960 PyObject_GC_Del, /* tp_free */
961};
962
963
964/* takewhile object **********************************************************/
965
966typedef struct {
967 PyObject_HEAD
968 PyObject *func;
969 PyObject *it;
970 long stop;
971} takewhileobject;
972
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000973static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000974
975static PyObject *
976takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
977{
978 PyObject *func, *seq;
979 PyObject *it;
980 takewhileobject *lz;
981
Thomas Woutersb2137042007-02-01 18:02:27 +0000982 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000983 return NULL;
984
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000985 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
986 return NULL;
987
988 /* Get iterator. */
989 it = PyObject_GetIter(seq);
990 if (it == NULL)
991 return NULL;
992
993 /* create takewhileobject structure */
994 lz = (takewhileobject *)type->tp_alloc(type, 0);
995 if (lz == NULL) {
996 Py_DECREF(it);
997 return NULL;
998 }
999 Py_INCREF(func);
1000 lz->func = func;
1001 lz->it = it;
1002 lz->stop = 0;
1003
1004 return (PyObject *)lz;
1005}
1006
1007static void
1008takewhile_dealloc(takewhileobject *lz)
1009{
1010 PyObject_GC_UnTrack(lz);
1011 Py_XDECREF(lz->func);
1012 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00001013 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001014}
1015
1016static int
1017takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1018{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001019 Py_VISIT(lz->it);
1020 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001021 return 0;
1022}
1023
1024static PyObject *
1025takewhile_next(takewhileobject *lz)
1026{
1027 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001028 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001029 long ok;
1030
1031 if (lz->stop == 1)
1032 return NULL;
1033
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001034 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00001035 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001036 if (item == NULL)
1037 return NULL;
1038
1039 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1040 if (good == NULL) {
1041 Py_DECREF(item);
1042 return NULL;
1043 }
1044 ok = PyObject_IsTrue(good);
1045 Py_DECREF(good);
1046 if (ok)
1047 return item;
1048 Py_DECREF(item);
1049 lz->stop = 1;
1050 return NULL;
1051}
1052
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053PyDoc_STRVAR(takewhile_doc,
1054"takewhile(predicate, iterable) --> takewhile object\n\
1055\n\
1056Return successive entries from an iterable as long as the \n\
1057predicate evaluates to true for each entry.");
1058
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001059static PyTypeObject takewhile_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001060 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001061 "itertools.takewhile", /* tp_name */
1062 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001063 0, /* tp_itemsize */
1064 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001065 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001066 0, /* tp_print */
1067 0, /* tp_getattr */
1068 0, /* tp_setattr */
1069 0, /* tp_compare */
1070 0, /* tp_repr */
1071 0, /* tp_as_number */
1072 0, /* tp_as_sequence */
1073 0, /* tp_as_mapping */
1074 0, /* tp_hash */
1075 0, /* tp_call */
1076 0, /* tp_str */
1077 PyObject_GenericGetAttr, /* tp_getattro */
1078 0, /* tp_setattro */
1079 0, /* tp_as_buffer */
1080 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1081 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001082 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001083 (traverseproc)takewhile_traverse, /* tp_traverse */
1084 0, /* tp_clear */
1085 0, /* tp_richcompare */
1086 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001087 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001088 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001089 0, /* tp_methods */
1090 0, /* tp_members */
1091 0, /* tp_getset */
1092 0, /* tp_base */
1093 0, /* tp_dict */
1094 0, /* tp_descr_get */
1095 0, /* tp_descr_set */
1096 0, /* tp_dictoffset */
1097 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001098 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001099 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001100 PyObject_GC_Del, /* tp_free */
1101};
1102
1103
1104/* islice object ************************************************************/
1105
1106typedef struct {
1107 PyObject_HEAD
1108 PyObject *it;
Thomas Wouters477c8d52006-05-27 19:21:47 +00001109 Py_ssize_t next;
1110 Py_ssize_t stop;
1111 Py_ssize_t step;
1112 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001113} isliceobject;
1114
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001115static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116
1117static PyObject *
1118islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1119{
1120 PyObject *seq;
Thomas Wouters477c8d52006-05-27 19:21:47 +00001121 Py_ssize_t start=0, stop=-1, step=1;
Raymond Hettingerb2594052004-12-05 09:25:51 +00001122 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001123 Py_ssize_t numargs;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001124 isliceobject *lz;
1125
Thomas Woutersb2137042007-02-01 18:02:27 +00001126 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001127 return NULL;
1128
Raymond Hettingerb2594052004-12-05 09:25:51 +00001129 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001130 return NULL;
1131
Raymond Hettingerb2594052004-12-05 09:25:51 +00001132 numargs = PyTuple_Size(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001133 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001134 if (a1 != Py_None) {
Christian Heimes217cfd12007-12-02 14:31:20 +00001135 stop = PyLong_AsSsize_t(a1);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001136 if (stop == -1) {
1137 if (PyErr_Occurred())
1138 PyErr_Clear();
1139 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001140 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001141 return NULL;
1142 }
1143 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001144 } else {
Raymond Hettingerb2594052004-12-05 09:25:51 +00001145 if (a1 != Py_None)
Christian Heimes217cfd12007-12-02 14:31:20 +00001146 start = PyLong_AsSsize_t(a1);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001147 if (start == -1 && PyErr_Occurred())
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001148 PyErr_Clear();
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001149 if (a2 != Py_None) {
Christian Heimes217cfd12007-12-02 14:31:20 +00001150 stop = PyLong_AsSsize_t(a2);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001151 if (stop == -1) {
1152 if (PyErr_Occurred())
1153 PyErr_Clear();
1154 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001155 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001156 return NULL;
1157 }
1158 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001159 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001160 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001161 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001162 "Indices for islice() must be non-negative integers or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001163 return NULL;
1164 }
1165
Raymond Hettingerb2594052004-12-05 09:25:51 +00001166 if (a3 != NULL) {
1167 if (a3 != Py_None)
Christian Heimes217cfd12007-12-02 14:31:20 +00001168 step = PyLong_AsSsize_t(a3);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001169 if (step == -1 && PyErr_Occurred())
1170 PyErr_Clear();
1171 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001172 if (step<1) {
1173 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001174 "Step for islice() must be a positive integer or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001175 return NULL;
1176 }
1177
1178 /* Get iterator. */
1179 it = PyObject_GetIter(seq);
1180 if (it == NULL)
1181 return NULL;
1182
1183 /* create isliceobject structure */
1184 lz = (isliceobject *)type->tp_alloc(type, 0);
1185 if (lz == NULL) {
1186 Py_DECREF(it);
1187 return NULL;
1188 }
1189 lz->it = it;
1190 lz->next = start;
1191 lz->stop = stop;
1192 lz->step = step;
1193 lz->cnt = 0L;
1194
1195 return (PyObject *)lz;
1196}
1197
1198static void
1199islice_dealloc(isliceobject *lz)
1200{
1201 PyObject_GC_UnTrack(lz);
1202 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00001203 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204}
1205
1206static int
1207islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1208{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001209 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001210 return 0;
1211}
1212
1213static PyObject *
1214islice_next(isliceobject *lz)
1215{
1216 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001217 PyObject *it = lz->it;
Thomas Wouters477c8d52006-05-27 19:21:47 +00001218 Py_ssize_t oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001219 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001220
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001221 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00001222 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001223 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001224 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001225 if (item == NULL)
1226 return NULL;
1227 Py_DECREF(item);
1228 lz->cnt++;
1229 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001230 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001231 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001232 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001233 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001234 if (item == NULL)
1235 return NULL;
1236 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001237 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001238 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001239 if (lz->next < oldnext) /* Check for overflow */
1240 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001241 return item;
1242}
1243
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001244PyDoc_STRVAR(islice_doc,
1245"islice(iterable, [start,] stop [, step]) --> islice object\n\
1246\n\
1247Return an iterator whose next() method returns selected values from an\n\
1248iterable. If start is specified, will skip all preceding elements;\n\
1249otherwise, start defaults to zero. Step defaults to one. If\n\
1250specified as another value, step determines how many values are \n\
1251skipped between successive calls. Works like a slice() on a list\n\
1252but returns an iterator.");
1253
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001254static PyTypeObject islice_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001255 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001256 "itertools.islice", /* tp_name */
1257 sizeof(isliceobject), /* tp_basicsize */
1258 0, /* tp_itemsize */
1259 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001260 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001261 0, /* tp_print */
1262 0, /* tp_getattr */
1263 0, /* tp_setattr */
1264 0, /* tp_compare */
1265 0, /* tp_repr */
1266 0, /* tp_as_number */
1267 0, /* tp_as_sequence */
1268 0, /* tp_as_mapping */
1269 0, /* tp_hash */
1270 0, /* tp_call */
1271 0, /* tp_str */
1272 PyObject_GenericGetAttr, /* tp_getattro */
1273 0, /* tp_setattro */
1274 0, /* tp_as_buffer */
1275 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1276 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001277 islice_doc, /* tp_doc */
1278 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001279 0, /* tp_clear */
1280 0, /* tp_richcompare */
1281 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001282 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001283 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001284 0, /* tp_methods */
1285 0, /* tp_members */
1286 0, /* tp_getset */
1287 0, /* tp_base */
1288 0, /* tp_dict */
1289 0, /* tp_descr_get */
1290 0, /* tp_descr_set */
1291 0, /* tp_dictoffset */
1292 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001293 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001294 islice_new, /* tp_new */
1295 PyObject_GC_Del, /* tp_free */
1296};
1297
1298
1299/* starmap object ************************************************************/
1300
1301typedef struct {
1302 PyObject_HEAD
1303 PyObject *func;
1304 PyObject *it;
1305} starmapobject;
1306
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001307static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001308
1309static PyObject *
1310starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1311{
1312 PyObject *func, *seq;
1313 PyObject *it;
1314 starmapobject *lz;
1315
Thomas Woutersb2137042007-02-01 18:02:27 +00001316 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001317 return NULL;
1318
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001319 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1320 return NULL;
1321
1322 /* Get iterator. */
1323 it = PyObject_GetIter(seq);
1324 if (it == NULL)
1325 return NULL;
1326
1327 /* create starmapobject structure */
1328 lz = (starmapobject *)type->tp_alloc(type, 0);
1329 if (lz == NULL) {
1330 Py_DECREF(it);
1331 return NULL;
1332 }
1333 Py_INCREF(func);
1334 lz->func = func;
1335 lz->it = it;
1336
1337 return (PyObject *)lz;
1338}
1339
1340static void
1341starmap_dealloc(starmapobject *lz)
1342{
1343 PyObject_GC_UnTrack(lz);
1344 Py_XDECREF(lz->func);
1345 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00001346 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001347}
1348
1349static int
1350starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1351{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001352 Py_VISIT(lz->it);
1353 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001354 return 0;
1355}
1356
1357static PyObject *
1358starmap_next(starmapobject *lz)
1359{
1360 PyObject *args;
1361 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001362 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001363
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001364 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00001365 args = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001366 if (args == NULL)
1367 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001368 if (!PyTuple_CheckExact(args)) {
Christian Heimes679db4a2008-01-18 09:56:22 +00001369 PyObject *newargs = PySequence_Tuple(args);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001370 Py_DECREF(args);
Christian Heimes679db4a2008-01-18 09:56:22 +00001371 if (newargs == NULL)
1372 return NULL;
1373 args = newargs;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001374 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001375 result = PyObject_Call(lz->func, args, NULL);
1376 Py_DECREF(args);
1377 return result;
1378}
1379
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001380PyDoc_STRVAR(starmap_doc,
1381"starmap(function, sequence) --> starmap object\n\
1382\n\
1383Return an iterator whose values are returned from the function evaluated\n\
1384with a argument tuple taken from the given sequence.");
1385
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001386static PyTypeObject starmap_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001387 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001388 "itertools.starmap", /* tp_name */
1389 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001390 0, /* tp_itemsize */
1391 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001392 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001393 0, /* tp_print */
1394 0, /* tp_getattr */
1395 0, /* tp_setattr */
1396 0, /* tp_compare */
1397 0, /* tp_repr */
1398 0, /* tp_as_number */
1399 0, /* tp_as_sequence */
1400 0, /* tp_as_mapping */
1401 0, /* tp_hash */
1402 0, /* tp_call */
1403 0, /* tp_str */
1404 PyObject_GenericGetAttr, /* tp_getattro */
1405 0, /* tp_setattro */
1406 0, /* tp_as_buffer */
1407 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1408 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001409 starmap_doc, /* tp_doc */
1410 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001411 0, /* tp_clear */
1412 0, /* tp_richcompare */
1413 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001414 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001415 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001416 0, /* tp_methods */
1417 0, /* tp_members */
1418 0, /* tp_getset */
1419 0, /* tp_base */
1420 0, /* tp_dict */
1421 0, /* tp_descr_get */
1422 0, /* tp_descr_set */
1423 0, /* tp_dictoffset */
1424 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001425 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001426 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001427 PyObject_GC_Del, /* tp_free */
1428};
1429
1430
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001431/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001432
1433typedef struct {
1434 PyObject_HEAD
Christian Heimesf16baeb2008-02-29 14:57:44 +00001435 PyObject *source; /* Iterator over input iterables */
1436 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001437} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001438
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001439static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001440
Christian Heimesf16baeb2008-02-29 14:57:44 +00001441static PyObject *
1442chain_new_internal(PyTypeObject *type, PyObject *source)
1443{
1444 chainobject *lz;
1445
1446 lz = (chainobject *)type->tp_alloc(type, 0);
1447 if (lz == NULL) {
1448 Py_DECREF(source);
1449 return NULL;
1450 }
1451
1452 lz->source = source;
1453 lz->active = NULL;
1454 return (PyObject *)lz;
1455}
1456
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001458chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459{
Christian Heimesf16baeb2008-02-29 14:57:44 +00001460 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001461
Thomas Woutersb2137042007-02-01 18:02:27 +00001462 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001463 return NULL;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001464
1465 source = PyObject_GetIter(args);
1466 if (source == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001467 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001468
Christian Heimesf16baeb2008-02-29 14:57:44 +00001469 return chain_new_internal(type, source);
1470}
1471
1472static PyObject *
1473chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1474{
1475 PyObject *source;
1476
1477 source = PyObject_GetIter(arg);
1478 if (source == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001479 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001480
Christian Heimesf16baeb2008-02-29 14:57:44 +00001481 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001482}
1483
1484static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001485chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001486{
1487 PyObject_GC_UnTrack(lz);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001488 Py_XDECREF(lz->active);
1489 Py_XDECREF(lz->source);
Christian Heimes90aa7642007-12-19 02:45:37 +00001490 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001491}
1492
Raymond Hettinger2012f172003-02-07 05:32:58 +00001493static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001494chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001495{
Christian Heimesf16baeb2008-02-29 14:57:44 +00001496 Py_VISIT(lz->source);
1497 Py_VISIT(lz->active);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001498 return 0;
1499}
1500
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001501static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001502chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001503{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001504 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001505
Christian Heimesf16baeb2008-02-29 14:57:44 +00001506 if (lz->source == NULL)
1507 return NULL; /* already stopped */
1508
1509 if (lz->active == NULL) {
1510 PyObject *iterable = PyIter_Next(lz->source);
1511 if (iterable == NULL) {
1512 Py_CLEAR(lz->source);
1513 return NULL; /* no more input sources */
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001514 }
Christian Heimesf16baeb2008-02-29 14:57:44 +00001515 lz->active = PyObject_GetIter(iterable);
Christian Heimes78644762008-03-04 23:39:23 +00001516 Py_DECREF(iterable);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001517 if (lz->active == NULL) {
Christian Heimesf16baeb2008-02-29 14:57:44 +00001518 Py_CLEAR(lz->source);
1519 return NULL; /* input not iterable */
1520 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001521 }
Christian Heimesf16baeb2008-02-29 14:57:44 +00001522 item = PyIter_Next(lz->active);
1523 if (item != NULL)
1524 return item;
1525 if (PyErr_Occurred()) {
1526 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1527 PyErr_Clear();
1528 else
1529 return NULL; /* input raised an exception */
1530 }
1531 Py_CLEAR(lz->active);
1532 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001533}
1534
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001535PyDoc_STRVAR(chain_doc,
1536"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001537\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001538Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001539first iterable until it is exhausted, then elements from the next\n\
1540iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001541
Christian Heimesf16baeb2008-02-29 14:57:44 +00001542PyDoc_STRVAR(chain_from_iterable_doc,
1543"chain.from_iterable(iterable) --> chain object\n\
1544\n\
1545Alternate chain() contructor taking a single iterable argument\n\
1546that evaluates lazily.");
1547
1548static PyMethodDef chain_methods[] = {
1549 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1550 chain_from_iterable_doc},
1551 {NULL, NULL} /* sentinel */
1552};
1553
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001554static PyTypeObject chain_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001555 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001556 "itertools.chain", /* tp_name */
1557 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001558 0, /* tp_itemsize */
1559 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001560 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001561 0, /* tp_print */
1562 0, /* tp_getattr */
1563 0, /* tp_setattr */
1564 0, /* tp_compare */
1565 0, /* tp_repr */
1566 0, /* tp_as_number */
1567 0, /* tp_as_sequence */
1568 0, /* tp_as_mapping */
1569 0, /* tp_hash */
1570 0, /* tp_call */
1571 0, /* tp_str */
1572 PyObject_GenericGetAttr, /* tp_getattro */
1573 0, /* tp_setattro */
1574 0, /* tp_as_buffer */
1575 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1576 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001577 chain_doc, /* tp_doc */
1578 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001579 0, /* tp_clear */
1580 0, /* tp_richcompare */
1581 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001582 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001583 (iternextfunc)chain_next, /* tp_iternext */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001584 chain_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001585 0, /* tp_members */
1586 0, /* tp_getset */
1587 0, /* tp_base */
1588 0, /* tp_dict */
1589 0, /* tp_descr_get */
1590 0, /* tp_descr_set */
1591 0, /* tp_dictoffset */
1592 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001593 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001594 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001595 PyObject_GC_Del, /* tp_free */
1596};
1597
1598
Christian Heimesc3f30c42008-02-22 16:37:40 +00001599/* product object ************************************************************/
1600
1601typedef struct {
1602 PyObject_HEAD
1603 PyObject *pools; /* tuple of pool tuples */
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001604 Py_ssize_t *indices; /* one index per pool */
1605 PyObject *result; /* most recently returned result tuple */
1606 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001607} productobject;
1608
1609static PyTypeObject product_type;
1610
1611static PyObject *
1612product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1613{
1614 productobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001615 Py_ssize_t nargs, npools, repeat=1;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001616 PyObject *pools = NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001617 Py_ssize_t *indices = NULL;
1618 Py_ssize_t i;
1619
Christian Heimesf16baeb2008-02-29 14:57:44 +00001620 if (kwds != NULL) {
1621 char *kwlist[] = {"repeat", 0};
1622 PyObject *tmpargs = PyTuple_New(0);
1623 if (tmpargs == NULL)
1624 return NULL;
1625 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1626 Py_DECREF(tmpargs);
1627 return NULL;
1628 }
1629 Py_DECREF(tmpargs);
1630 if (repeat < 0) {
1631 PyErr_SetString(PyExc_ValueError,
1632 "repeat argument cannot be negative");
1633 return NULL;
1634 }
1635 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001636
1637 assert(PyTuple_Check(args));
Christian Heimesf16baeb2008-02-29 14:57:44 +00001638 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1639 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001640
Christian Heimesc3f30c42008-02-22 16:37:40 +00001641 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
Christian Heimesb558a2e2008-03-02 22:46:37 +00001642 if (indices == NULL) {
Christian Heimesc3f30c42008-02-22 16:37:40 +00001643 PyErr_NoMemory();
1644 goto error;
1645 }
1646
1647 pools = PyTuple_New(npools);
1648 if (pools == NULL)
1649 goto error;
1650
Christian Heimesf16baeb2008-02-29 14:57:44 +00001651 for (i=0; i < nargs ; ++i) {
Christian Heimesc3f30c42008-02-22 16:37:40 +00001652 PyObject *item = PyTuple_GET_ITEM(args, i);
1653 PyObject *pool = PySequence_Tuple(item);
1654 if (pool == NULL)
1655 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001656 PyTuple_SET_ITEM(pools, i, pool);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001657 indices[i] = 0;
1658 }
Christian Heimesf16baeb2008-02-29 14:57:44 +00001659 for ( ; i < npools; ++i) {
1660 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1661 Py_INCREF(pool);
1662 PyTuple_SET_ITEM(pools, i, pool);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001663 indices[i] = 0;
1664 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001665
1666 /* create productobject structure */
1667 lz = (productobject *)type->tp_alloc(type, 0);
Christian Heimes380f7f22008-02-28 11:19:05 +00001668 if (lz == NULL)
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001669 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001670
1671 lz->pools = pools;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001672 lz->indices = indices;
1673 lz->result = NULL;
1674 lz->stopped = 0;
1675
1676 return (PyObject *)lz;
1677
1678error:
Christian Heimesc3f30c42008-02-22 16:37:40 +00001679 if (indices != NULL)
1680 PyMem_Free(indices);
1681 Py_XDECREF(pools);
1682 return NULL;
1683}
1684
1685static void
1686product_dealloc(productobject *lz)
1687{
1688 PyObject_GC_UnTrack(lz);
1689 Py_XDECREF(lz->pools);
1690 Py_XDECREF(lz->result);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001691 PyMem_Free(lz->indices);
1692 Py_TYPE(lz)->tp_free(lz);
1693}
1694
1695static int
1696product_traverse(productobject *lz, visitproc visit, void *arg)
1697{
1698 Py_VISIT(lz->pools);
1699 Py_VISIT(lz->result);
1700 return 0;
1701}
1702
1703static PyObject *
1704product_next(productobject *lz)
1705{
1706 PyObject *pool;
1707 PyObject *elem;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001708 PyObject *oldelem;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001709 PyObject *pools = lz->pools;
1710 PyObject *result = lz->result;
1711 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1712 Py_ssize_t i;
1713
1714 if (lz->stopped)
1715 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001716
Christian Heimesc3f30c42008-02-22 16:37:40 +00001717 if (result == NULL) {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001718 /* On the first pass, return an initial tuple filled with the
Christian Heimes78644762008-03-04 23:39:23 +00001719 first element from each pool. */
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001720 result = PyTuple_New(npools);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001721 if (result == NULL)
1722 goto empty;
1723 lz->result = result;
1724 for (i=0; i < npools; i++) {
1725 pool = PyTuple_GET_ITEM(pools, i);
1726 if (PyTuple_GET_SIZE(pool) == 0)
1727 goto empty;
1728 elem = PyTuple_GET_ITEM(pool, 0);
1729 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001730 PyTuple_SET_ITEM(result, i, elem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001731 }
1732 } else {
1733 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001734
1735 /* Copy the previous result tuple or re-use it if available */
1736 if (Py_REFCNT(result) > 1) {
1737 PyObject *old_result = result;
1738 result = PyTuple_New(npools);
1739 if (result == NULL)
1740 goto empty;
1741 lz->result = result;
1742 for (i=0; i < npools; i++) {
1743 elem = PyTuple_GET_ITEM(old_result, i);
1744 Py_INCREF(elem);
1745 PyTuple_SET_ITEM(result, i, elem);
1746 }
1747 Py_DECREF(old_result);
1748 }
1749 /* Now, we've got the only copy so we can update it in-place */
Christian Heimesb558a2e2008-03-02 22:46:37 +00001750 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001751
1752 /* Update the pool indices right-to-left. Only advance to the
1753 next pool when the previous one rolls-over */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001754 for (i=npools-1 ; i >= 0 ; i--) {
1755 pool = PyTuple_GET_ITEM(pools, i);
1756 indices[i]++;
Christian Heimesb558a2e2008-03-02 22:46:37 +00001757 if (indices[i] == PyTuple_GET_SIZE(pool)) {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001758 /* Roll-over and advance to next pool */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001759 indices[i] = 0;
1760 elem = PyTuple_GET_ITEM(pool, 0);
1761 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001762 oldelem = PyTuple_GET_ITEM(result, i);
1763 PyTuple_SET_ITEM(result, i, elem);
1764 Py_DECREF(oldelem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001765 } else {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001766 /* No rollover. Just increment and stop here. */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001767 elem = PyTuple_GET_ITEM(pool, indices[i]);
1768 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001769 oldelem = PyTuple_GET_ITEM(result, i);
1770 PyTuple_SET_ITEM(result, i, elem);
1771 Py_DECREF(oldelem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001772 break;
1773 }
1774 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001775
1776 /* If i is negative, then the indices have all rolled-over
1777 and we're done. */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001778 if (i < 0)
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001779 goto empty;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001780 }
1781
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001782 Py_INCREF(result);
1783 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001784
1785empty:
1786 lz->stopped = 1;
1787 return NULL;
1788}
1789
1790PyDoc_STRVAR(product_doc,
1791"product(*iterables) --> product object\n\
1792\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001793Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001794For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1795The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1796cycle in a manner similar to an odometer (with the rightmost element changing\n\
1797on every iteration).\n\n\
1798product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1799product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1800
1801static PyTypeObject product_type = {
1802 PyVarObject_HEAD_INIT(NULL, 0)
1803 "itertools.product", /* tp_name */
1804 sizeof(productobject), /* tp_basicsize */
1805 0, /* tp_itemsize */
1806 /* methods */
1807 (destructor)product_dealloc, /* tp_dealloc */
1808 0, /* tp_print */
1809 0, /* tp_getattr */
1810 0, /* tp_setattr */
1811 0, /* tp_compare */
1812 0, /* tp_repr */
1813 0, /* tp_as_number */
1814 0, /* tp_as_sequence */
1815 0, /* tp_as_mapping */
1816 0, /* tp_hash */
1817 0, /* tp_call */
1818 0, /* tp_str */
1819 PyObject_GenericGetAttr, /* tp_getattro */
1820 0, /* tp_setattro */
1821 0, /* tp_as_buffer */
1822 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1823 Py_TPFLAGS_BASETYPE, /* tp_flags */
1824 product_doc, /* tp_doc */
1825 (traverseproc)product_traverse, /* tp_traverse */
1826 0, /* tp_clear */
1827 0, /* tp_richcompare */
1828 0, /* tp_weaklistoffset */
1829 PyObject_SelfIter, /* tp_iter */
1830 (iternextfunc)product_next, /* tp_iternext */
1831 0, /* tp_methods */
1832 0, /* tp_members */
1833 0, /* tp_getset */
1834 0, /* tp_base */
1835 0, /* tp_dict */
1836 0, /* tp_descr_get */
1837 0, /* tp_descr_set */
1838 0, /* tp_dictoffset */
1839 0, /* tp_init */
1840 0, /* tp_alloc */
1841 product_new, /* tp_new */
1842 PyObject_GC_Del, /* tp_free */
1843};
1844
1845
Christian Heimes380f7f22008-02-28 11:19:05 +00001846/* combinations object ************************************************************/
1847
1848typedef struct {
1849 PyObject_HEAD
1850 PyObject *pool; /* input converted to a tuple */
1851 Py_ssize_t *indices; /* one index per result element */
1852 PyObject *result; /* most recently returned result tuple */
1853 Py_ssize_t r; /* size of result tuple */
1854 int stopped; /* set to 1 when the combinations iterator is exhausted */
1855} combinationsobject;
1856
1857static PyTypeObject combinations_type;
1858
1859static PyObject *
1860combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1861{
1862 combinationsobject *co;
1863 Py_ssize_t n;
1864 Py_ssize_t r;
1865 PyObject *pool = NULL;
1866 PyObject *iterable = NULL;
1867 Py_ssize_t *indices = NULL;
1868 Py_ssize_t i;
1869 static char *kwargs[] = {"iterable", "r", NULL};
1870
1871 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
1872 &iterable, &r))
1873 return NULL;
1874
1875 pool = PySequence_Tuple(iterable);
1876 if (pool == NULL)
1877 goto error;
1878 n = PyTuple_GET_SIZE(pool);
1879 if (r < 0) {
1880 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
1881 goto error;
1882 }
Christian Heimes380f7f22008-02-28 11:19:05 +00001883
1884 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
1885 if (indices == NULL) {
1886 PyErr_NoMemory();
1887 goto error;
1888 }
1889
1890 for (i=0 ; i<r ; i++)
1891 indices[i] = i;
1892
1893 /* create combinationsobject structure */
1894 co = (combinationsobject *)type->tp_alloc(type, 0);
1895 if (co == NULL)
1896 goto error;
1897
1898 co->pool = pool;
1899 co->indices = indices;
1900 co->result = NULL;
1901 co->r = r;
Raymond Hettinger5bad41e2009-01-08 21:01:54 +00001902 co->stopped = r > n ? 1 : 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00001903
1904 return (PyObject *)co;
1905
1906error:
1907 if (indices != NULL)
1908 PyMem_Free(indices);
1909 Py_XDECREF(pool);
1910 return NULL;
1911}
1912
1913static void
1914combinations_dealloc(combinationsobject *co)
1915{
1916 PyObject_GC_UnTrack(co);
1917 Py_XDECREF(co->pool);
1918 Py_XDECREF(co->result);
1919 PyMem_Free(co->indices);
1920 Py_TYPE(co)->tp_free(co);
1921}
1922
1923static int
1924combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
1925{
1926 Py_VISIT(co->pool);
1927 Py_VISIT(co->result);
1928 return 0;
1929}
1930
1931static PyObject *
1932combinations_next(combinationsobject *co)
1933{
1934 PyObject *elem;
1935 PyObject *oldelem;
1936 PyObject *pool = co->pool;
1937 Py_ssize_t *indices = co->indices;
1938 PyObject *result = co->result;
1939 Py_ssize_t n = PyTuple_GET_SIZE(pool);
1940 Py_ssize_t r = co->r;
1941 Py_ssize_t i, j, index;
1942
1943 if (co->stopped)
1944 return NULL;
1945
1946 if (result == NULL) {
1947 /* On the first pass, initialize result tuple using the indices */
1948 result = PyTuple_New(r);
1949 if (result == NULL)
1950 goto empty;
1951 co->result = result;
1952 for (i=0; i<r ; i++) {
1953 index = indices[i];
1954 elem = PyTuple_GET_ITEM(pool, index);
1955 Py_INCREF(elem);
1956 PyTuple_SET_ITEM(result, i, elem);
1957 }
1958 } else {
1959 /* Copy the previous result tuple or re-use it if available */
1960 if (Py_REFCNT(result) > 1) {
1961 PyObject *old_result = result;
1962 result = PyTuple_New(r);
1963 if (result == NULL)
1964 goto empty;
1965 co->result = result;
1966 for (i=0; i<r ; i++) {
1967 elem = PyTuple_GET_ITEM(old_result, i);
1968 Py_INCREF(elem);
1969 PyTuple_SET_ITEM(result, i, elem);
1970 }
1971 Py_DECREF(old_result);
1972 }
1973 /* Now, we've got the only copy so we can update it in-place
1974 * CPython's empty tuple is a singleton and cached in
1975 * PyTuple's freelist.
1976 */
1977 assert(r == 0 || Py_REFCNT(result) == 1);
1978
1979 /* Scan indices right-to-left until finding one that is not
1980 at its maximum (i + n - r). */
1981 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
1982 ;
1983
1984 /* If i is negative, then the indices are all at
1985 their maximum value and we're done. */
1986 if (i < 0)
1987 goto empty;
1988
1989 /* Increment the current index which we know is not at its
1990 maximum. Then move back to the right setting each index
1991 to its lowest possible value (one higher than the index
1992 to its left -- this maintains the sort order invariant). */
1993 indices[i]++;
1994 for (j=i+1 ; j<r ; j++)
1995 indices[j] = indices[j-1] + 1;
1996
1997 /* Update the result tuple for the new indices
1998 starting with i, the leftmost index that changed */
1999 for ( ; i<r ; i++) {
2000 index = indices[i];
2001 elem = PyTuple_GET_ITEM(pool, index);
2002 Py_INCREF(elem);
2003 oldelem = PyTuple_GET_ITEM(result, i);
2004 PyTuple_SET_ITEM(result, i, elem);
2005 Py_DECREF(oldelem);
2006 }
2007 }
2008
2009 Py_INCREF(result);
2010 return result;
2011
2012empty:
2013 co->stopped = 1;
2014 return NULL;
2015}
2016
2017PyDoc_STRVAR(combinations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00002018"combinations(iterable[, r]) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002019\n\
2020Return successive r-length combinations of elements in the iterable.\n\n\
2021combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2022
2023static PyTypeObject combinations_type = {
2024 PyVarObject_HEAD_INIT(NULL, 0)
2025 "itertools.combinations", /* tp_name */
2026 sizeof(combinationsobject), /* tp_basicsize */
2027 0, /* tp_itemsize */
2028 /* methods */
2029 (destructor)combinations_dealloc, /* tp_dealloc */
2030 0, /* tp_print */
2031 0, /* tp_getattr */
2032 0, /* tp_setattr */
2033 0, /* tp_compare */
2034 0, /* tp_repr */
2035 0, /* tp_as_number */
2036 0, /* tp_as_sequence */
2037 0, /* tp_as_mapping */
2038 0, /* tp_hash */
2039 0, /* tp_call */
2040 0, /* tp_str */
2041 PyObject_GenericGetAttr, /* tp_getattro */
2042 0, /* tp_setattro */
2043 0, /* tp_as_buffer */
2044 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2045 Py_TPFLAGS_BASETYPE, /* tp_flags */
2046 combinations_doc, /* tp_doc */
2047 (traverseproc)combinations_traverse, /* tp_traverse */
2048 0, /* tp_clear */
2049 0, /* tp_richcompare */
2050 0, /* tp_weaklistoffset */
2051 PyObject_SelfIter, /* tp_iter */
2052 (iternextfunc)combinations_next, /* tp_iternext */
2053 0, /* tp_methods */
2054 0, /* tp_members */
2055 0, /* tp_getset */
2056 0, /* tp_base */
2057 0, /* tp_dict */
2058 0, /* tp_descr_get */
2059 0, /* tp_descr_set */
2060 0, /* tp_dictoffset */
2061 0, /* tp_init */
2062 0, /* tp_alloc */
2063 combinations_new, /* tp_new */
2064 PyObject_GC_Del, /* tp_free */
2065};
2066
2067
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002068/* permutations object ************************************************************
2069
2070def permutations(iterable, r=None):
2071 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2072 pool = tuple(iterable)
2073 n = len(pool)
2074 r = n if r is None else r
2075 indices = range(n)
2076 cycles = range(n-r+1, n+1)[::-1]
2077 yield tuple(pool[i] for i in indices[:r])
2078 while n:
2079 for i in reversed(range(r)):
2080 cycles[i] -= 1
2081 if cycles[i] == 0:
2082 indices[i:] = indices[i+1:] + indices[i:i+1]
2083 cycles[i] = n - i
2084 else:
2085 j = cycles[i]
2086 indices[i], indices[-j] = indices[-j], indices[i]
2087 yield tuple(pool[i] for i in indices[:r])
2088 break
2089 else:
2090 return
2091*/
2092
2093typedef struct {
2094 PyObject_HEAD
2095 PyObject *pool; /* input converted to a tuple */
2096 Py_ssize_t *indices; /* one index per element in the pool */
2097 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2098 PyObject *result; /* most recently returned result tuple */
2099 Py_ssize_t r; /* size of result tuple */
2100 int stopped; /* set to 1 when the permutations iterator is exhausted */
2101} permutationsobject;
2102
2103static PyTypeObject permutations_type;
2104
2105static PyObject *
2106permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2107{
2108 permutationsobject *po;
2109 Py_ssize_t n;
2110 Py_ssize_t r;
2111 PyObject *robj = Py_None;
2112 PyObject *pool = NULL;
2113 PyObject *iterable = NULL;
2114 Py_ssize_t *indices = NULL;
2115 Py_ssize_t *cycles = NULL;
2116 Py_ssize_t i;
2117 static char *kwargs[] = {"iterable", "r", NULL};
2118
2119 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2120 &iterable, &robj))
2121 return NULL;
2122
2123 pool = PySequence_Tuple(iterable);
2124 if (pool == NULL)
2125 goto error;
2126 n = PyTuple_GET_SIZE(pool);
2127
2128 r = n;
2129 if (robj != Py_None) {
2130 if (!PyLong_Check(robj)) {
2131 PyErr_SetString(PyExc_TypeError, "Expected int as r");
Neal Norwitzd2bee322008-04-01 07:37:58 +00002132 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002133 }
2134 r = PyLong_AsSsize_t(robj);
2135 if (r == -1 && PyErr_Occurred())
2136 goto error;
2137 }
2138 if (r < 0) {
2139 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2140 goto error;
2141 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002142
2143 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2144 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2145 if (indices == NULL || cycles == NULL) {
2146 PyErr_NoMemory();
2147 goto error;
2148 }
2149
2150 for (i=0 ; i<n ; i++)
2151 indices[i] = i;
2152 for (i=0 ; i<r ; i++)
2153 cycles[i] = n - i;
2154
2155 /* create permutationsobject structure */
2156 po = (permutationsobject *)type->tp_alloc(type, 0);
2157 if (po == NULL)
2158 goto error;
2159
2160 po->pool = pool;
2161 po->indices = indices;
2162 po->cycles = cycles;
2163 po->result = NULL;
2164 po->r = r;
Raymond Hettinger5bad41e2009-01-08 21:01:54 +00002165 po->stopped = r > n ? 1 : 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002166
2167 return (PyObject *)po;
2168
2169error:
2170 if (indices != NULL)
2171 PyMem_Free(indices);
2172 if (cycles != NULL)
2173 PyMem_Free(cycles);
2174 Py_XDECREF(pool);
2175 return NULL;
2176}
2177
2178static void
2179permutations_dealloc(permutationsobject *po)
2180{
2181 PyObject_GC_UnTrack(po);
2182 Py_XDECREF(po->pool);
2183 Py_XDECREF(po->result);
2184 PyMem_Free(po->indices);
2185 PyMem_Free(po->cycles);
2186 Py_TYPE(po)->tp_free(po);
2187}
2188
2189static int
2190permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2191{
2192 Py_VISIT(po->pool);
2193 Py_VISIT(po->result);
2194 return 0;
2195}
2196
2197static PyObject *
2198permutations_next(permutationsobject *po)
2199{
2200 PyObject *elem;
2201 PyObject *oldelem;
2202 PyObject *pool = po->pool;
2203 Py_ssize_t *indices = po->indices;
2204 Py_ssize_t *cycles = po->cycles;
2205 PyObject *result = po->result;
2206 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2207 Py_ssize_t r = po->r;
2208 Py_ssize_t i, j, k, index;
2209
2210 if (po->stopped)
2211 return NULL;
2212
2213 if (result == NULL) {
2214 /* On the first pass, initialize result tuple using the indices */
2215 result = PyTuple_New(r);
2216 if (result == NULL)
2217 goto empty;
2218 po->result = result;
2219 for (i=0; i<r ; i++) {
2220 index = indices[i];
2221 elem = PyTuple_GET_ITEM(pool, index);
2222 Py_INCREF(elem);
2223 PyTuple_SET_ITEM(result, i, elem);
2224 }
2225 } else {
2226 if (n == 0)
2227 goto empty;
2228
2229 /* Copy the previous result tuple or re-use it if available */
2230 if (Py_REFCNT(result) > 1) {
2231 PyObject *old_result = result;
2232 result = PyTuple_New(r);
2233 if (result == NULL)
2234 goto empty;
2235 po->result = result;
2236 for (i=0; i<r ; i++) {
2237 elem = PyTuple_GET_ITEM(old_result, i);
2238 Py_INCREF(elem);
2239 PyTuple_SET_ITEM(result, i, elem);
2240 }
2241 Py_DECREF(old_result);
2242 }
2243 /* Now, we've got the only copy so we can update it in-place */
2244 assert(r == 0 || Py_REFCNT(result) == 1);
2245
2246 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2247 for (i=r-1 ; i>=0 ; i--) {
2248 cycles[i] -= 1;
2249 if (cycles[i] == 0) {
2250 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2251 index = indices[i];
2252 for (j=i ; j<n-1 ; j++)
2253 indices[j] = indices[j+1];
2254 indices[n-1] = index;
2255 cycles[i] = n - i;
2256 } else {
2257 j = cycles[i];
2258 index = indices[i];
2259 indices[i] = indices[n-j];
2260 indices[n-j] = index;
2261
2262 for (k=i; k<r ; k++) {
2263 /* start with i, the leftmost element that changed */
2264 /* yield tuple(pool[k] for k in indices[:r]) */
2265 index = indices[k];
2266 elem = PyTuple_GET_ITEM(pool, index);
2267 Py_INCREF(elem);
2268 oldelem = PyTuple_GET_ITEM(result, k);
2269 PyTuple_SET_ITEM(result, k, elem);
2270 Py_DECREF(oldelem);
2271 }
2272 break;
2273 }
2274 }
2275 /* If i is negative, then the cycles have all
2276 rolled-over and we're done. */
2277 if (i < 0)
2278 goto empty;
2279 }
2280 Py_INCREF(result);
2281 return result;
2282
2283empty:
2284 po->stopped = 1;
2285 return NULL;
2286}
2287
2288PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00002289"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002290\n\
2291Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00002292permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002293
2294static PyTypeObject permutations_type = {
2295 PyVarObject_HEAD_INIT(NULL, 0)
2296 "itertools.permutations", /* tp_name */
2297 sizeof(permutationsobject), /* tp_basicsize */
2298 0, /* tp_itemsize */
2299 /* methods */
2300 (destructor)permutations_dealloc, /* tp_dealloc */
2301 0, /* tp_print */
2302 0, /* tp_getattr */
2303 0, /* tp_setattr */
2304 0, /* tp_compare */
2305 0, /* tp_repr */
2306 0, /* tp_as_number */
2307 0, /* tp_as_sequence */
2308 0, /* tp_as_mapping */
2309 0, /* tp_hash */
2310 0, /* tp_call */
2311 0, /* tp_str */
2312 PyObject_GenericGetAttr, /* tp_getattro */
2313 0, /* tp_setattro */
2314 0, /* tp_as_buffer */
2315 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2316 Py_TPFLAGS_BASETYPE, /* tp_flags */
2317 permutations_doc, /* tp_doc */
2318 (traverseproc)permutations_traverse, /* tp_traverse */
2319 0, /* tp_clear */
2320 0, /* tp_richcompare */
2321 0, /* tp_weaklistoffset */
2322 PyObject_SelfIter, /* tp_iter */
2323 (iternextfunc)permutations_next, /* tp_iternext */
2324 0, /* tp_methods */
2325 0, /* tp_members */
2326 0, /* tp_getset */
2327 0, /* tp_base */
2328 0, /* tp_dict */
2329 0, /* tp_descr_get */
2330 0, /* tp_descr_set */
2331 0, /* tp_dictoffset */
2332 0, /* tp_init */
2333 0, /* tp_alloc */
2334 permutations_new, /* tp_new */
2335 PyObject_GC_Del, /* tp_free */
2336};
2337
2338
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002339/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00002340
2341typedef struct {
2342 PyObject_HEAD
2343 PyObject *func;
2344 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002345} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002346
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002347static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002348
2349static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002350filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002351{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00002352 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002353 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002354 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002355
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002356 if (type == &filterfalse_type &&
2357 !_PyArg_NoKeywords("filterfalse()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002358 return NULL;
2359
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002360 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
Raymond Hettinger60eca932003-02-09 06:40:58 +00002361 return NULL;
2362
2363 /* Get iterator. */
2364 it = PyObject_GetIter(seq);
2365 if (it == NULL)
2366 return NULL;
2367
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002368 /* create filterfalseobject structure */
2369 lz = (filterfalseobject *)type->tp_alloc(type, 0);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002370 if (lz == NULL) {
2371 Py_DECREF(it);
2372 return NULL;
2373 }
2374 Py_INCREF(func);
2375 lz->func = func;
2376 lz->it = it;
2377
2378 return (PyObject *)lz;
2379}
2380
2381static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002382filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002383{
2384 PyObject_GC_UnTrack(lz);
2385 Py_XDECREF(lz->func);
2386 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00002387 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002388}
2389
2390static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002391filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002392{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002393 Py_VISIT(lz->it);
2394 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002395 return 0;
2396}
2397
2398static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002399filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002400{
2401 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002402 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002403 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002404 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002405
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002406 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00002407 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002408 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002409 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002410 if (item == NULL)
2411 return NULL;
2412
Christian Heimes836baa52008-02-26 08:18:30 +00002413 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
Raymond Hettinger60eca932003-02-09 06:40:58 +00002414 ok = PyObject_IsTrue(item);
2415 } else {
2416 PyObject *good;
2417 good = PyObject_CallFunctionObjArgs(lz->func,
2418 item, NULL);
2419 if (good == NULL) {
2420 Py_DECREF(item);
2421 return NULL;
2422 }
2423 ok = PyObject_IsTrue(good);
2424 Py_DECREF(good);
2425 }
2426 if (!ok)
2427 return item;
2428 Py_DECREF(item);
2429 }
2430}
2431
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002432PyDoc_STRVAR(filterfalse_doc,
2433"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002434\n\
2435Return those items of sequence for which function(item) is false.\n\
2436If function is None, return the items that are false.");
2437
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002438static PyTypeObject filterfalse_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002439 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002440 "itertools.filterfalse", /* tp_name */
2441 sizeof(filterfalseobject), /* tp_basicsize */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002442 0, /* tp_itemsize */
2443 /* methods */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002444 (destructor)filterfalse_dealloc, /* tp_dealloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002445 0, /* tp_print */
2446 0, /* tp_getattr */
2447 0, /* tp_setattr */
2448 0, /* tp_compare */
2449 0, /* tp_repr */
2450 0, /* tp_as_number */
2451 0, /* tp_as_sequence */
2452 0, /* tp_as_mapping */
2453 0, /* tp_hash */
2454 0, /* tp_call */
2455 0, /* tp_str */
2456 PyObject_GenericGetAttr, /* tp_getattro */
2457 0, /* tp_setattro */
2458 0, /* tp_as_buffer */
2459 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2460 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002461 filterfalse_doc, /* tp_doc */
2462 (traverseproc)filterfalse_traverse, /* tp_traverse */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002463 0, /* tp_clear */
2464 0, /* tp_richcompare */
2465 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002466 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002467 (iternextfunc)filterfalse_next, /* tp_iternext */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002468 0, /* tp_methods */
2469 0, /* tp_members */
2470 0, /* tp_getset */
2471 0, /* tp_base */
2472 0, /* tp_dict */
2473 0, /* tp_descr_get */
2474 0, /* tp_descr_set */
2475 0, /* tp_dictoffset */
2476 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002477 0, /* tp_alloc */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002478 filterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002479 PyObject_GC_Del, /* tp_free */
2480};
2481
2482
2483/* count object ************************************************************/
2484
2485typedef struct {
2486 PyObject_HEAD
Thomas Wouters477c8d52006-05-27 19:21:47 +00002487 Py_ssize_t cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002488 PyObject *long_cnt; /* Arbitrarily large count when cnt >= PY_SSIZE_T_MAX */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002489} countobject;
2490
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002491static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002492
2493static PyObject *
2494count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2495{
2496 countobject *lz;
Thomas Wouters477c8d52006-05-27 19:21:47 +00002497 Py_ssize_t cnt = 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002498 PyObject *cnt_arg = NULL;
2499 PyObject *long_cnt = NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002500
Thomas Woutersb2137042007-02-01 18:02:27 +00002501 if (type == &count_type && !_PyArg_NoKeywords("count()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002502 return NULL;
2503
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002504 if (!PyArg_UnpackTuple(args, "count", 0, 1, &cnt_arg))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002505 return NULL;
2506
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002507 if (cnt_arg != NULL) {
Christian Heimes217cfd12007-12-02 14:31:20 +00002508 cnt = PyLong_AsSsize_t(cnt_arg);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002509 if (cnt == -1 && PyErr_Occurred()) {
2510 PyErr_Clear();
2511 if (!PyLong_Check(cnt_arg)) {
2512 PyErr_SetString(PyExc_TypeError, "an integer is required");
2513 return NULL;
2514 }
2515 long_cnt = cnt_arg;
2516 Py_INCREF(long_cnt);
2517 cnt = PY_SSIZE_T_MAX;
2518 }
2519 }
2520
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002521 /* create countobject structure */
2522 lz = (countobject *)PyObject_New(countobject, &count_type);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002523 if (lz == NULL) {
2524 Py_XDECREF(long_cnt);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002525 return NULL;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002526 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002527 lz->cnt = cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002528 lz->long_cnt = long_cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002529
2530 return (PyObject *)lz;
2531}
2532
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002533static void
2534count_dealloc(countobject *lz)
2535{
2536 Py_XDECREF(lz->long_cnt);
2537 PyObject_Del(lz);
2538}
2539
2540static PyObject *
2541count_nextlong(countobject *lz)
2542{
2543 static PyObject *one = NULL;
2544 PyObject *cnt;
2545 PyObject *stepped_up;
2546
2547 if (lz->long_cnt == NULL) {
Christian Heimes217cfd12007-12-02 14:31:20 +00002548 lz->long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002549 if (lz->long_cnt == NULL)
2550 return NULL;
2551 }
2552 if (one == NULL) {
Christian Heimes217cfd12007-12-02 14:31:20 +00002553 one = PyLong_FromLong(1);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002554 if (one == NULL)
2555 return NULL;
2556 }
2557 cnt = lz->long_cnt;
2558 assert(cnt != NULL);
2559 stepped_up = PyNumber_Add(cnt, one);
2560 if (stepped_up == NULL)
2561 return NULL;
2562 lz->long_cnt = stepped_up;
2563 return cnt;
2564}
2565
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002566static PyObject *
2567count_next(countobject *lz)
2568{
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002569 if (lz->cnt == PY_SSIZE_T_MAX)
2570 return count_nextlong(lz);
Christian Heimes217cfd12007-12-02 14:31:20 +00002571 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002572}
2573
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002574static PyObject *
2575count_repr(countobject *lz)
2576{
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002577 if (lz->cnt != PY_SSIZE_T_MAX)
2578 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
2579
2580 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002581}
2582
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002583PyDoc_STRVAR(count_doc,
2584"count([firstval]) --> count object\n\
2585\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002586Return a count object whose .__next__() method returns consecutive\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002587integers starting from zero or, if specified, from firstval.");
2588
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002589static PyTypeObject count_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002590 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002591 "itertools.count", /* tp_name */
2592 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002593 0, /* tp_itemsize */
2594 /* methods */
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002595 (destructor)count_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002596 0, /* tp_print */
2597 0, /* tp_getattr */
2598 0, /* tp_setattr */
2599 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002600 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002601 0, /* tp_as_number */
2602 0, /* tp_as_sequence */
2603 0, /* tp_as_mapping */
2604 0, /* tp_hash */
2605 0, /* tp_call */
2606 0, /* tp_str */
2607 PyObject_GenericGetAttr, /* tp_getattro */
2608 0, /* tp_setattro */
2609 0, /* tp_as_buffer */
2610 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002611 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002612 0, /* tp_traverse */
2613 0, /* tp_clear */
2614 0, /* tp_richcompare */
2615 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002616 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002617 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002618 0, /* tp_methods */
2619 0, /* tp_members */
2620 0, /* tp_getset */
2621 0, /* tp_base */
2622 0, /* tp_dict */
2623 0, /* tp_descr_get */
2624 0, /* tp_descr_set */
2625 0, /* tp_dictoffset */
2626 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002627 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002628 count_new, /* tp_new */
2629};
2630
2631
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002632/* repeat object ************************************************************/
2633
2634typedef struct {
2635 PyObject_HEAD
2636 PyObject *element;
Thomas Wouters477c8d52006-05-27 19:21:47 +00002637 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002638} repeatobject;
2639
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002640static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002641
2642static PyObject *
2643repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2644{
2645 repeatobject *ro;
2646 PyObject *element;
Thomas Wouters477c8d52006-05-27 19:21:47 +00002647 Py_ssize_t cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002648
Thomas Woutersb2137042007-02-01 18:02:27 +00002649 if (type == &repeat_type && !_PyArg_NoKeywords("repeat()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002650 return NULL;
2651
Thomas Wouters477c8d52006-05-27 19:21:47 +00002652 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002653 return NULL;
2654
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002655 if (PyTuple_Size(args) == 2 && cnt < 0)
2656 cnt = 0;
2657
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002658 ro = (repeatobject *)type->tp_alloc(type, 0);
2659 if (ro == NULL)
2660 return NULL;
2661 Py_INCREF(element);
2662 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002663 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002664 return (PyObject *)ro;
2665}
2666
2667static void
2668repeat_dealloc(repeatobject *ro)
2669{
2670 PyObject_GC_UnTrack(ro);
2671 Py_XDECREF(ro->element);
Christian Heimes90aa7642007-12-19 02:45:37 +00002672 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002673}
2674
2675static int
2676repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2677{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002678 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002679 return 0;
2680}
2681
2682static PyObject *
2683repeat_next(repeatobject *ro)
2684{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002685 if (ro->cnt == 0)
2686 return NULL;
2687 if (ro->cnt > 0)
2688 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002689 Py_INCREF(ro->element);
2690 return ro->element;
2691}
2692
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002693static PyObject *
2694repeat_repr(repeatobject *ro)
2695{
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002696 if (ro->cnt == -1)
Walter Dörwald7569dfe2007-05-19 21:49:49 +00002697 return PyUnicode_FromFormat("repeat(%R)", ro->element);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002698 else
Walter Dörwald7569dfe2007-05-19 21:49:49 +00002699 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002700}
2701
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002702static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002703repeat_len(repeatobject *ro)
2704{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002705 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002706 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002707 return NULL;
2708 }
Christian Heimes217cfd12007-12-02 14:31:20 +00002709 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002710}
2711
Armin Rigof5b3e362006-02-11 21:32:43 +00002712PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002713
2714static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00002715 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002716 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002717};
2718
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002719PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002720"repeat(element [,times]) -> create an iterator which returns the element\n\
2721for the specified number of times. If not specified, returns the element\n\
2722endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002723
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002724static PyTypeObject repeat_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002725 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002726 "itertools.repeat", /* tp_name */
2727 sizeof(repeatobject), /* tp_basicsize */
2728 0, /* tp_itemsize */
2729 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002730 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002731 0, /* tp_print */
2732 0, /* tp_getattr */
2733 0, /* tp_setattr */
2734 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002735 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002736 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002737 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002738 0, /* tp_as_mapping */
2739 0, /* tp_hash */
2740 0, /* tp_call */
2741 0, /* tp_str */
2742 PyObject_GenericGetAttr, /* tp_getattro */
2743 0, /* tp_setattro */
2744 0, /* tp_as_buffer */
2745 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2746 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002747 repeat_doc, /* tp_doc */
2748 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002749 0, /* tp_clear */
2750 0, /* tp_richcompare */
2751 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002752 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002753 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002754 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002755 0, /* tp_members */
2756 0, /* tp_getset */
2757 0, /* tp_base */
2758 0, /* tp_dict */
2759 0, /* tp_descr_get */
2760 0, /* tp_descr_set */
2761 0, /* tp_dictoffset */
2762 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002763 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002764 repeat_new, /* tp_new */
2765 PyObject_GC_Del, /* tp_free */
2766};
2767
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002768/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00002769
2770#include "Python.h"
2771
2772typedef struct {
2773 PyObject_HEAD
2774 Py_ssize_t tuplesize;
2775 Py_ssize_t numactive;
2776 PyObject *ittuple; /* tuple of iterators */
2777 PyObject *result;
2778 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002779} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00002780
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002781static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00002782
2783static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002784zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00002785{
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002786 ziplongestobject *lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00002787 Py_ssize_t i;
2788 PyObject *ittuple; /* tuple of iterators */
2789 PyObject *result;
2790 PyObject *fillvalue = Py_None;
2791 Py_ssize_t tuplesize = PySequence_Length(args);
2792
2793 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
2794 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
2795 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
2796 PyErr_SetString(PyExc_TypeError,
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002797 "zip_longest() got an unexpected keyword argument");
Thomas Wouterscf297e42007-02-23 15:07:44 +00002798 return NULL;
2799 }
2800 }
2801
2802 /* args must be a tuple */
2803 assert(PyTuple_Check(args));
2804
2805 /* obtain iterators */
2806 ittuple = PyTuple_New(tuplesize);
2807 if (ittuple == NULL)
2808 return NULL;
2809 for (i=0; i < tuplesize; ++i) {
2810 PyObject *item = PyTuple_GET_ITEM(args, i);
2811 PyObject *it = PyObject_GetIter(item);
2812 if (it == NULL) {
2813 if (PyErr_ExceptionMatches(PyExc_TypeError))
2814 PyErr_Format(PyExc_TypeError,
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002815 "zip_longest argument #%zd must support iteration",
Thomas Wouterscf297e42007-02-23 15:07:44 +00002816 i+1);
2817 Py_DECREF(ittuple);
2818 return NULL;
2819 }
2820 PyTuple_SET_ITEM(ittuple, i, it);
2821 }
2822
2823 /* create a result holder */
2824 result = PyTuple_New(tuplesize);
2825 if (result == NULL) {
2826 Py_DECREF(ittuple);
2827 return NULL;
2828 }
2829 for (i=0 ; i < tuplesize ; i++) {
2830 Py_INCREF(Py_None);
2831 PyTuple_SET_ITEM(result, i, Py_None);
2832 }
2833
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002834 /* create ziplongestobject structure */
2835 lz = (ziplongestobject *)type->tp_alloc(type, 0);
Thomas Wouterscf297e42007-02-23 15:07:44 +00002836 if (lz == NULL) {
2837 Py_DECREF(ittuple);
2838 Py_DECREF(result);
2839 return NULL;
2840 }
2841 lz->ittuple = ittuple;
2842 lz->tuplesize = tuplesize;
2843 lz->numactive = tuplesize;
2844 lz->result = result;
2845 Py_INCREF(fillvalue);
2846 lz->fillvalue = fillvalue;
2847 return (PyObject *)lz;
2848}
2849
2850static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002851zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00002852{
2853 PyObject_GC_UnTrack(lz);
2854 Py_XDECREF(lz->ittuple);
2855 Py_XDECREF(lz->result);
2856 Py_XDECREF(lz->fillvalue);
Christian Heimes90aa7642007-12-19 02:45:37 +00002857 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00002858}
2859
2860static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002861zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00002862{
2863 Py_VISIT(lz->ittuple);
2864 Py_VISIT(lz->result);
2865 Py_VISIT(lz->fillvalue);
2866 return 0;
2867}
2868
2869static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002870zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00002871{
2872 Py_ssize_t i;
2873 Py_ssize_t tuplesize = lz->tuplesize;
2874 PyObject *result = lz->result;
2875 PyObject *it;
2876 PyObject *item;
2877 PyObject *olditem;
2878
2879 if (tuplesize == 0)
2880 return NULL;
2881 if (lz->numactive == 0)
2882 return NULL;
Christian Heimes90aa7642007-12-19 02:45:37 +00002883 if (Py_REFCNT(result) == 1) {
Thomas Wouterscf297e42007-02-23 15:07:44 +00002884 Py_INCREF(result);
2885 for (i=0 ; i < tuplesize ; i++) {
2886 it = PyTuple_GET_ITEM(lz->ittuple, i);
2887 if (it == NULL) {
2888 Py_INCREF(lz->fillvalue);
2889 item = lz->fillvalue;
2890 } else {
2891 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00002892 item = (*Py_TYPE(it)->tp_iternext)(it);
Thomas Wouterscf297e42007-02-23 15:07:44 +00002893 if (item == NULL) {
2894 lz->numactive -= 1;
2895 if (lz->numactive == 0) {
2896 Py_DECREF(result);
2897 return NULL;
2898 } else {
2899 Py_INCREF(lz->fillvalue);
2900 item = lz->fillvalue;
2901 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
2902 Py_DECREF(it);
2903 }
2904 }
2905 }
2906 olditem = PyTuple_GET_ITEM(result, i);
2907 PyTuple_SET_ITEM(result, i, item);
2908 Py_DECREF(olditem);
2909 }
2910 } else {
2911 result = PyTuple_New(tuplesize);
2912 if (result == NULL)
2913 return NULL;
2914 for (i=0 ; i < tuplesize ; i++) {
2915 it = PyTuple_GET_ITEM(lz->ittuple, i);
2916 if (it == NULL) {
2917 Py_INCREF(lz->fillvalue);
2918 item = lz->fillvalue;
2919 } else {
2920 assert(PyIter_Check(it));
Christian Heimes90aa7642007-12-19 02:45:37 +00002921 item = (*Py_TYPE(it)->tp_iternext)(it);
Thomas Wouterscf297e42007-02-23 15:07:44 +00002922 if (item == NULL) {
2923 lz->numactive -= 1;
2924 if (lz->numactive == 0) {
2925 Py_DECREF(result);
2926 return NULL;
2927 } else {
2928 Py_INCREF(lz->fillvalue);
2929 item = lz->fillvalue;
2930 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
2931 Py_DECREF(it);
2932 }
2933 }
2934 }
2935 PyTuple_SET_ITEM(result, i, item);
2936 }
2937 }
2938 return result;
2939}
2940
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002941PyDoc_STRVAR(zip_longest_doc,
2942"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00002943\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002944Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002945the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00002946method continues until the longest iterable in the argument sequence\n\
2947is exhausted and then it raises StopIteration. When the shorter iterables\n\
2948are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
2949defaults to None or can be specified by a keyword argument.\n\
2950");
2951
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002952static PyTypeObject ziplongest_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002953 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002954 "itertools.zip_longest", /* tp_name */
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002955 sizeof(ziplongestobject), /* tp_basicsize */
Thomas Wouterscf297e42007-02-23 15:07:44 +00002956 0, /* tp_itemsize */
2957 /* methods */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002958 (destructor)zip_longest_dealloc, /* tp_dealloc */
Thomas Wouterscf297e42007-02-23 15:07:44 +00002959 0, /* tp_print */
2960 0, /* tp_getattr */
2961 0, /* tp_setattr */
2962 0, /* tp_compare */
2963 0, /* tp_repr */
2964 0, /* tp_as_number */
2965 0, /* tp_as_sequence */
2966 0, /* tp_as_mapping */
2967 0, /* tp_hash */
2968 0, /* tp_call */
2969 0, /* tp_str */
2970 PyObject_GenericGetAttr, /* tp_getattro */
2971 0, /* tp_setattro */
2972 0, /* tp_as_buffer */
2973 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2974 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002975 zip_longest_doc, /* tp_doc */
2976 (traverseproc)zip_longest_traverse, /* tp_traverse */
Thomas Wouterscf297e42007-02-23 15:07:44 +00002977 0, /* tp_clear */
2978 0, /* tp_richcompare */
2979 0, /* tp_weaklistoffset */
2980 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002981 (iternextfunc)zip_longest_next, /* tp_iternext */
Thomas Wouterscf297e42007-02-23 15:07:44 +00002982 0, /* tp_methods */
2983 0, /* tp_members */
2984 0, /* tp_getset */
2985 0, /* tp_base */
2986 0, /* tp_dict */
2987 0, /* tp_descr_get */
2988 0, /* tp_descr_set */
2989 0, /* tp_dictoffset */
2990 0, /* tp_init */
2991 0, /* tp_alloc */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002992 zip_longest_new, /* tp_new */
Thomas Wouterscf297e42007-02-23 15:07:44 +00002993 PyObject_GC_Del, /* tp_free */
2994};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002995
2996/* module level code ********************************************************/
2997
2998PyDoc_STRVAR(module_doc,
2999"Functional tools for creating and using iterators.\n\
3000\n\
3001Infinite iterators:\n\
3002count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003003cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003004repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003005\n\
3006Iterators terminating on the shortest input sequence:\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003007zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
3008filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003009islice(seq, [start,] stop [, step]) --> elements from\n\
3010 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003011starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003012tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003013chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003014takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
3015dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003016groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003017");
3018
3019
Raymond Hettingerad983e72003-11-12 14:32:26 +00003020static PyMethodDef module_methods[] = {
3021 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3022 {NULL, NULL} /* sentinel */
3023};
3024
Martin v. Löwis1a214512008-06-11 05:26:20 +00003025
3026static struct PyModuleDef itertoolsmodule = {
3027 PyModuleDef_HEAD_INIT,
3028 "itertools",
3029 module_doc,
3030 -1,
3031 module_methods,
3032 NULL,
3033 NULL,
3034 NULL,
3035 NULL
3036};
3037
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003038PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00003039PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003040{
Raymond Hettinger60eca932003-02-09 06:40:58 +00003041 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003042 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003043 char *name;
3044 PyTypeObject *typelist[] = {
Christian Heimes380f7f22008-02-28 11:19:05 +00003045 &combinations_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003046 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003047 &dropwhile_type,
3048 &takewhile_type,
3049 &islice_type,
3050 &starmap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003051 &chain_type,
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003052 &filterfalse_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003053 &count_type,
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003054 &ziplongest_type,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003055 &permutations_type,
Christian Heimes380f7f22008-02-28 11:19:05 +00003056 &product_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003057 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003058 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003059 NULL
3060 };
3061
Christian Heimes90aa7642007-12-19 02:45:37 +00003062 Py_TYPE(&teedataobject_type) = &PyType_Type;
Martin v. Löwis1a214512008-06-11 05:26:20 +00003063 m = PyModule_Create(&itertoolsmodule);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00003064 if (m == NULL)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003065 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003066
Raymond Hettinger60eca932003-02-09 06:40:58 +00003067 for (i=0 ; typelist[i] != NULL ; i++) {
3068 if (PyType_Ready(typelist[i]) < 0)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003069 return NULL;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003070 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003071 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003072 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003073 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003074 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003075
3076 if (PyType_Ready(&teedataobject_type) < 0)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003077 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +00003078 if (PyType_Ready(&tee_type) < 0)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003079 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003080 if (PyType_Ready(&_grouper_type) < 0)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003081 return NULL;
3082 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003083}