blob: 3cec4212918111eec74191786add4db113dff84b [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 */
Mark Dickinsone94c6792009-02-02 20:36:42 +0000151 0, /* tp_reserved */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000152 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 */
Mark Dickinsone94c6792009-02-02 20:36:42 +0000282 0, /* tp_reserved */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000283 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 */
Mark Dickinsone94c6792009-02-02 20:36:42 +0000435 0, /* tp_reserved */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000436 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 */
Mark Dickinsone94c6792009-02-02 20:36:42 +0000591 0, /* tp_reserved */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000592 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 */
Mark Dickinsone94c6792009-02-02 20:36:42 +0000786 0, /* tp_reserved */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000787 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
Christian Heimes90aa7642007-12-19 02:45:37 +0000889 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000890 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000891 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000892 if (item == NULL)
893 return NULL;
894 if (lz->start == 1)
895 return item;
896
897 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
898 if (good == NULL) {
899 Py_DECREF(item);
900 return NULL;
901 }
902 ok = PyObject_IsTrue(good);
903 Py_DECREF(good);
904 if (!ok) {
905 lz->start = 1;
906 return item;
907 }
908 Py_DECREF(item);
909 }
910}
911
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000912PyDoc_STRVAR(dropwhile_doc,
913"dropwhile(predicate, iterable) --> dropwhile object\n\
914\n\
915Drop items from the iterable while predicate(item) is true.\n\
916Afterwards, return every element until the iterable is exhausted.");
917
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000918static PyTypeObject dropwhile_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +0000919 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000920 "itertools.dropwhile", /* tp_name */
921 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000922 0, /* tp_itemsize */
923 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000924 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000925 0, /* tp_print */
926 0, /* tp_getattr */
927 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +0000928 0, /* tp_reserved */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000929 0, /* tp_repr */
930 0, /* tp_as_number */
931 0, /* tp_as_sequence */
932 0, /* tp_as_mapping */
933 0, /* tp_hash */
934 0, /* tp_call */
935 0, /* tp_str */
936 PyObject_GenericGetAttr, /* tp_getattro */
937 0, /* tp_setattro */
938 0, /* tp_as_buffer */
939 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
940 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000941 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000942 (traverseproc)dropwhile_traverse, /* tp_traverse */
943 0, /* tp_clear */
944 0, /* tp_richcompare */
945 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000946 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000947 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000948 0, /* tp_methods */
949 0, /* tp_members */
950 0, /* tp_getset */
951 0, /* tp_base */
952 0, /* tp_dict */
953 0, /* tp_descr_get */
954 0, /* tp_descr_set */
955 0, /* tp_dictoffset */
956 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000957 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000958 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000959 PyObject_GC_Del, /* tp_free */
960};
961
962
963/* takewhile object **********************************************************/
964
965typedef struct {
966 PyObject_HEAD
967 PyObject *func;
968 PyObject *it;
969 long stop;
970} takewhileobject;
971
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000972static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000973
974static PyObject *
975takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
976{
977 PyObject *func, *seq;
978 PyObject *it;
979 takewhileobject *lz;
980
Thomas Woutersb2137042007-02-01 18:02:27 +0000981 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000982 return NULL;
983
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000984 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
985 return NULL;
986
987 /* Get iterator. */
988 it = PyObject_GetIter(seq);
989 if (it == NULL)
990 return NULL;
991
992 /* create takewhileobject structure */
993 lz = (takewhileobject *)type->tp_alloc(type, 0);
994 if (lz == NULL) {
995 Py_DECREF(it);
996 return NULL;
997 }
998 Py_INCREF(func);
999 lz->func = func;
1000 lz->it = it;
1001 lz->stop = 0;
1002
1003 return (PyObject *)lz;
1004}
1005
1006static void
1007takewhile_dealloc(takewhileobject *lz)
1008{
1009 PyObject_GC_UnTrack(lz);
1010 Py_XDECREF(lz->func);
1011 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00001012 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001013}
1014
1015static int
1016takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1017{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001018 Py_VISIT(lz->it);
1019 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001020 return 0;
1021}
1022
1023static PyObject *
1024takewhile_next(takewhileobject *lz)
1025{
1026 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001027 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001028 long ok;
1029
1030 if (lz->stop == 1)
1031 return NULL;
1032
Christian Heimes90aa7642007-12-19 02:45:37 +00001033 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001034 if (item == NULL)
1035 return NULL;
1036
1037 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1038 if (good == NULL) {
1039 Py_DECREF(item);
1040 return NULL;
1041 }
1042 ok = PyObject_IsTrue(good);
1043 Py_DECREF(good);
1044 if (ok)
1045 return item;
1046 Py_DECREF(item);
1047 lz->stop = 1;
1048 return NULL;
1049}
1050
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001051PyDoc_STRVAR(takewhile_doc,
1052"takewhile(predicate, iterable) --> takewhile object\n\
1053\n\
1054Return successive entries from an iterable as long as the \n\
1055predicate evaluates to true for each entry.");
1056
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001057static PyTypeObject takewhile_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001058 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001059 "itertools.takewhile", /* tp_name */
1060 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001061 0, /* tp_itemsize */
1062 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001063 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001064 0, /* tp_print */
1065 0, /* tp_getattr */
1066 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00001067 0, /* tp_reserved */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001068 0, /* tp_repr */
1069 0, /* tp_as_number */
1070 0, /* tp_as_sequence */
1071 0, /* tp_as_mapping */
1072 0, /* tp_hash */
1073 0, /* tp_call */
1074 0, /* tp_str */
1075 PyObject_GenericGetAttr, /* tp_getattro */
1076 0, /* tp_setattro */
1077 0, /* tp_as_buffer */
1078 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1079 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001080 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001081 (traverseproc)takewhile_traverse, /* tp_traverse */
1082 0, /* tp_clear */
1083 0, /* tp_richcompare */
1084 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001085 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001086 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001087 0, /* tp_methods */
1088 0, /* tp_members */
1089 0, /* tp_getset */
1090 0, /* tp_base */
1091 0, /* tp_dict */
1092 0, /* tp_descr_get */
1093 0, /* tp_descr_set */
1094 0, /* tp_dictoffset */
1095 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001096 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001097 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001098 PyObject_GC_Del, /* tp_free */
1099};
1100
1101
1102/* islice object ************************************************************/
1103
1104typedef struct {
1105 PyObject_HEAD
1106 PyObject *it;
Thomas Wouters477c8d52006-05-27 19:21:47 +00001107 Py_ssize_t next;
1108 Py_ssize_t stop;
1109 Py_ssize_t step;
1110 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001111} isliceobject;
1112
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001113static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001114
1115static PyObject *
1116islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1117{
1118 PyObject *seq;
Thomas Wouters477c8d52006-05-27 19:21:47 +00001119 Py_ssize_t start=0, stop=-1, step=1;
Raymond Hettingerb2594052004-12-05 09:25:51 +00001120 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001121 Py_ssize_t numargs;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001122 isliceobject *lz;
1123
Thomas Woutersb2137042007-02-01 18:02:27 +00001124 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001125 return NULL;
1126
Raymond Hettingerb2594052004-12-05 09:25:51 +00001127 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001128 return NULL;
1129
Raymond Hettingerb2594052004-12-05 09:25:51 +00001130 numargs = PyTuple_Size(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001131 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001132 if (a1 != Py_None) {
Christian Heimes217cfd12007-12-02 14:31:20 +00001133 stop = PyLong_AsSsize_t(a1);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001134 if (stop == -1) {
1135 if (PyErr_Occurred())
1136 PyErr_Clear();
1137 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001138 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001139 return NULL;
1140 }
1141 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001142 } else {
Raymond Hettingerb2594052004-12-05 09:25:51 +00001143 if (a1 != Py_None)
Christian Heimes217cfd12007-12-02 14:31:20 +00001144 start = PyLong_AsSsize_t(a1);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001145 if (start == -1 && PyErr_Occurred())
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001146 PyErr_Clear();
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001147 if (a2 != Py_None) {
Christian Heimes217cfd12007-12-02 14:31:20 +00001148 stop = PyLong_AsSsize_t(a2);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001149 if (stop == -1) {
1150 if (PyErr_Occurred())
1151 PyErr_Clear();
1152 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001153 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001154 return NULL;
1155 }
1156 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001157 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001158 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001159 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001160 "Indices for islice() must be non-negative integers or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001161 return NULL;
1162 }
1163
Raymond Hettingerb2594052004-12-05 09:25:51 +00001164 if (a3 != NULL) {
1165 if (a3 != Py_None)
Christian Heimes217cfd12007-12-02 14:31:20 +00001166 step = PyLong_AsSsize_t(a3);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001167 if (step == -1 && PyErr_Occurred())
1168 PyErr_Clear();
1169 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001170 if (step<1) {
1171 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001172 "Step for islice() must be a positive integer or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001173 return NULL;
1174 }
1175
1176 /* Get iterator. */
1177 it = PyObject_GetIter(seq);
1178 if (it == NULL)
1179 return NULL;
1180
1181 /* create isliceobject structure */
1182 lz = (isliceobject *)type->tp_alloc(type, 0);
1183 if (lz == NULL) {
1184 Py_DECREF(it);
1185 return NULL;
1186 }
1187 lz->it = it;
1188 lz->next = start;
1189 lz->stop = stop;
1190 lz->step = step;
1191 lz->cnt = 0L;
1192
1193 return (PyObject *)lz;
1194}
1195
1196static void
1197islice_dealloc(isliceobject *lz)
1198{
1199 PyObject_GC_UnTrack(lz);
1200 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00001201 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001202}
1203
1204static int
1205islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1206{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001207 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001208 return 0;
1209}
1210
1211static PyObject *
1212islice_next(isliceobject *lz)
1213{
1214 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001215 PyObject *it = lz->it;
Thomas Wouters477c8d52006-05-27 19:21:47 +00001216 Py_ssize_t oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001217 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001218
Christian Heimes90aa7642007-12-19 02:45:37 +00001219 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001220 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001221 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001222 if (item == NULL)
1223 return NULL;
1224 Py_DECREF(item);
1225 lz->cnt++;
1226 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001227 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001228 return NULL;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001229 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001230 if (item == NULL)
1231 return NULL;
1232 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001233 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001234 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001235 if (lz->next < oldnext) /* Check for overflow */
1236 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001237 return item;
1238}
1239
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001240PyDoc_STRVAR(islice_doc,
1241"islice(iterable, [start,] stop [, step]) --> islice object\n\
1242\n\
1243Return an iterator whose next() method returns selected values from an\n\
1244iterable. If start is specified, will skip all preceding elements;\n\
1245otherwise, start defaults to zero. Step defaults to one. If\n\
1246specified as another value, step determines how many values are \n\
1247skipped between successive calls. Works like a slice() on a list\n\
1248but returns an iterator.");
1249
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001250static PyTypeObject islice_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001251 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001252 "itertools.islice", /* tp_name */
1253 sizeof(isliceobject), /* tp_basicsize */
1254 0, /* tp_itemsize */
1255 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001256 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001257 0, /* tp_print */
1258 0, /* tp_getattr */
1259 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00001260 0, /* tp_reserved */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001261 0, /* tp_repr */
1262 0, /* tp_as_number */
1263 0, /* tp_as_sequence */
1264 0, /* tp_as_mapping */
1265 0, /* tp_hash */
1266 0, /* tp_call */
1267 0, /* tp_str */
1268 PyObject_GenericGetAttr, /* tp_getattro */
1269 0, /* tp_setattro */
1270 0, /* tp_as_buffer */
1271 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1272 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001273 islice_doc, /* tp_doc */
1274 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001275 0, /* tp_clear */
1276 0, /* tp_richcompare */
1277 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001278 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001279 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001280 0, /* tp_methods */
1281 0, /* tp_members */
1282 0, /* tp_getset */
1283 0, /* tp_base */
1284 0, /* tp_dict */
1285 0, /* tp_descr_get */
1286 0, /* tp_descr_set */
1287 0, /* tp_dictoffset */
1288 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001289 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001290 islice_new, /* tp_new */
1291 PyObject_GC_Del, /* tp_free */
1292};
1293
1294
1295/* starmap object ************************************************************/
1296
1297typedef struct {
1298 PyObject_HEAD
1299 PyObject *func;
1300 PyObject *it;
1301} starmapobject;
1302
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001303static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001304
1305static PyObject *
1306starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1307{
1308 PyObject *func, *seq;
1309 PyObject *it;
1310 starmapobject *lz;
1311
Thomas Woutersb2137042007-02-01 18:02:27 +00001312 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001313 return NULL;
1314
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001315 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1316 return NULL;
1317
1318 /* Get iterator. */
1319 it = PyObject_GetIter(seq);
1320 if (it == NULL)
1321 return NULL;
1322
1323 /* create starmapobject structure */
1324 lz = (starmapobject *)type->tp_alloc(type, 0);
1325 if (lz == NULL) {
1326 Py_DECREF(it);
1327 return NULL;
1328 }
1329 Py_INCREF(func);
1330 lz->func = func;
1331 lz->it = it;
1332
1333 return (PyObject *)lz;
1334}
1335
1336static void
1337starmap_dealloc(starmapobject *lz)
1338{
1339 PyObject_GC_UnTrack(lz);
1340 Py_XDECREF(lz->func);
1341 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00001342 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001343}
1344
1345static int
1346starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1347{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001348 Py_VISIT(lz->it);
1349 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001350 return 0;
1351}
1352
1353static PyObject *
1354starmap_next(starmapobject *lz)
1355{
1356 PyObject *args;
1357 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001358 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001359
Christian Heimes90aa7642007-12-19 02:45:37 +00001360 args = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001361 if (args == NULL)
1362 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001363 if (!PyTuple_CheckExact(args)) {
Christian Heimes679db4a2008-01-18 09:56:22 +00001364 PyObject *newargs = PySequence_Tuple(args);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001365 Py_DECREF(args);
Christian Heimes679db4a2008-01-18 09:56:22 +00001366 if (newargs == NULL)
1367 return NULL;
1368 args = newargs;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001369 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001370 result = PyObject_Call(lz->func, args, NULL);
1371 Py_DECREF(args);
1372 return result;
1373}
1374
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001375PyDoc_STRVAR(starmap_doc,
1376"starmap(function, sequence) --> starmap object\n\
1377\n\
1378Return an iterator whose values are returned from the function evaluated\n\
1379with a argument tuple taken from the given sequence.");
1380
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001381static PyTypeObject starmap_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001382 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001383 "itertools.starmap", /* tp_name */
1384 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001385 0, /* tp_itemsize */
1386 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001387 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001388 0, /* tp_print */
1389 0, /* tp_getattr */
1390 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00001391 0, /* tp_reserved */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001392 0, /* tp_repr */
1393 0, /* tp_as_number */
1394 0, /* tp_as_sequence */
1395 0, /* tp_as_mapping */
1396 0, /* tp_hash */
1397 0, /* tp_call */
1398 0, /* tp_str */
1399 PyObject_GenericGetAttr, /* tp_getattro */
1400 0, /* tp_setattro */
1401 0, /* tp_as_buffer */
1402 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1403 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001404 starmap_doc, /* tp_doc */
1405 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001406 0, /* tp_clear */
1407 0, /* tp_richcompare */
1408 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001409 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001410 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001411 0, /* tp_methods */
1412 0, /* tp_members */
1413 0, /* tp_getset */
1414 0, /* tp_base */
1415 0, /* tp_dict */
1416 0, /* tp_descr_get */
1417 0, /* tp_descr_set */
1418 0, /* tp_dictoffset */
1419 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001420 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001421 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001422 PyObject_GC_Del, /* tp_free */
1423};
1424
1425
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001426/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001427
1428typedef struct {
1429 PyObject_HEAD
Christian Heimesf16baeb2008-02-29 14:57:44 +00001430 PyObject *source; /* Iterator over input iterables */
1431 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001432} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001433
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001434static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001435
Christian Heimesf16baeb2008-02-29 14:57:44 +00001436static PyObject *
1437chain_new_internal(PyTypeObject *type, PyObject *source)
1438{
1439 chainobject *lz;
1440
1441 lz = (chainobject *)type->tp_alloc(type, 0);
1442 if (lz == NULL) {
1443 Py_DECREF(source);
1444 return NULL;
1445 }
1446
1447 lz->source = source;
1448 lz->active = NULL;
1449 return (PyObject *)lz;
1450}
1451
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001452static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001453chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001454{
Christian Heimesf16baeb2008-02-29 14:57:44 +00001455 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001456
Thomas Woutersb2137042007-02-01 18:02:27 +00001457 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001458 return NULL;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001459
1460 source = PyObject_GetIter(args);
1461 if (source == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001462 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001463
Christian Heimesf16baeb2008-02-29 14:57:44 +00001464 return chain_new_internal(type, source);
1465}
1466
1467static PyObject *
1468chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1469{
1470 PyObject *source;
1471
1472 source = PyObject_GetIter(arg);
1473 if (source == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001474 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001475
Christian Heimesf16baeb2008-02-29 14:57:44 +00001476 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001477}
1478
1479static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001480chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001481{
1482 PyObject_GC_UnTrack(lz);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001483 Py_XDECREF(lz->active);
1484 Py_XDECREF(lz->source);
Christian Heimes90aa7642007-12-19 02:45:37 +00001485 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001486}
1487
Raymond Hettinger2012f172003-02-07 05:32:58 +00001488static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001489chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001490{
Christian Heimesf16baeb2008-02-29 14:57:44 +00001491 Py_VISIT(lz->source);
1492 Py_VISIT(lz->active);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001493 return 0;
1494}
1495
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001496static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001497chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001498{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001499 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001500
Christian Heimesf16baeb2008-02-29 14:57:44 +00001501 if (lz->source == NULL)
1502 return NULL; /* already stopped */
1503
1504 if (lz->active == NULL) {
1505 PyObject *iterable = PyIter_Next(lz->source);
1506 if (iterable == NULL) {
1507 Py_CLEAR(lz->source);
1508 return NULL; /* no more input sources */
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001509 }
Christian Heimesf16baeb2008-02-29 14:57:44 +00001510 lz->active = PyObject_GetIter(iterable);
Christian Heimes78644762008-03-04 23:39:23 +00001511 Py_DECREF(iterable);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001512 if (lz->active == NULL) {
Christian Heimesf16baeb2008-02-29 14:57:44 +00001513 Py_CLEAR(lz->source);
1514 return NULL; /* input not iterable */
1515 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001516 }
Christian Heimesf16baeb2008-02-29 14:57:44 +00001517 item = PyIter_Next(lz->active);
1518 if (item != NULL)
1519 return item;
1520 if (PyErr_Occurred()) {
1521 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1522 PyErr_Clear();
1523 else
1524 return NULL; /* input raised an exception */
1525 }
1526 Py_CLEAR(lz->active);
1527 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001528}
1529
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001530PyDoc_STRVAR(chain_doc,
1531"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001532\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001533Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001534first iterable until it is exhausted, then elements from the next\n\
1535iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001536
Christian Heimesf16baeb2008-02-29 14:57:44 +00001537PyDoc_STRVAR(chain_from_iterable_doc,
1538"chain.from_iterable(iterable) --> chain object\n\
1539\n\
1540Alternate chain() contructor taking a single iterable argument\n\
1541that evaluates lazily.");
1542
1543static PyMethodDef chain_methods[] = {
1544 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1545 chain_from_iterable_doc},
1546 {NULL, NULL} /* sentinel */
1547};
1548
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001549static PyTypeObject chain_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001550 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001551 "itertools.chain", /* tp_name */
1552 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001553 0, /* tp_itemsize */
1554 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001555 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001556 0, /* tp_print */
1557 0, /* tp_getattr */
1558 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00001559 0, /* tp_reserved */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001560 0, /* tp_repr */
1561 0, /* tp_as_number */
1562 0, /* tp_as_sequence */
1563 0, /* tp_as_mapping */
1564 0, /* tp_hash */
1565 0, /* tp_call */
1566 0, /* tp_str */
1567 PyObject_GenericGetAttr, /* tp_getattro */
1568 0, /* tp_setattro */
1569 0, /* tp_as_buffer */
1570 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1571 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001572 chain_doc, /* tp_doc */
1573 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001574 0, /* tp_clear */
1575 0, /* tp_richcompare */
1576 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001577 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001578 (iternextfunc)chain_next, /* tp_iternext */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001579 chain_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001580 0, /* tp_members */
1581 0, /* tp_getset */
1582 0, /* tp_base */
1583 0, /* tp_dict */
1584 0, /* tp_descr_get */
1585 0, /* tp_descr_set */
1586 0, /* tp_dictoffset */
1587 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001588 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001589 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001590 PyObject_GC_Del, /* tp_free */
1591};
1592
1593
Christian Heimesc3f30c42008-02-22 16:37:40 +00001594/* product object ************************************************************/
1595
1596typedef struct {
1597 PyObject_HEAD
1598 PyObject *pools; /* tuple of pool tuples */
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001599 Py_ssize_t *indices; /* one index per pool */
1600 PyObject *result; /* most recently returned result tuple */
1601 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001602} productobject;
1603
1604static PyTypeObject product_type;
1605
1606static PyObject *
1607product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1608{
1609 productobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001610 Py_ssize_t nargs, npools, repeat=1;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001611 PyObject *pools = NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001612 Py_ssize_t *indices = NULL;
1613 Py_ssize_t i;
1614
Christian Heimesf16baeb2008-02-29 14:57:44 +00001615 if (kwds != NULL) {
1616 char *kwlist[] = {"repeat", 0};
1617 PyObject *tmpargs = PyTuple_New(0);
1618 if (tmpargs == NULL)
1619 return NULL;
1620 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1621 Py_DECREF(tmpargs);
1622 return NULL;
1623 }
1624 Py_DECREF(tmpargs);
1625 if (repeat < 0) {
1626 PyErr_SetString(PyExc_ValueError,
1627 "repeat argument cannot be negative");
1628 return NULL;
1629 }
1630 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001631
1632 assert(PyTuple_Check(args));
Christian Heimesf16baeb2008-02-29 14:57:44 +00001633 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1634 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001635
Christian Heimesc3f30c42008-02-22 16:37:40 +00001636 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
Christian Heimesb558a2e2008-03-02 22:46:37 +00001637 if (indices == NULL) {
Christian Heimesc3f30c42008-02-22 16:37:40 +00001638 PyErr_NoMemory();
1639 goto error;
1640 }
1641
1642 pools = PyTuple_New(npools);
1643 if (pools == NULL)
1644 goto error;
1645
Christian Heimesf16baeb2008-02-29 14:57:44 +00001646 for (i=0; i < nargs ; ++i) {
Christian Heimesc3f30c42008-02-22 16:37:40 +00001647 PyObject *item = PyTuple_GET_ITEM(args, i);
1648 PyObject *pool = PySequence_Tuple(item);
1649 if (pool == NULL)
1650 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001651 PyTuple_SET_ITEM(pools, i, pool);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001652 indices[i] = 0;
1653 }
Christian Heimesf16baeb2008-02-29 14:57:44 +00001654 for ( ; i < npools; ++i) {
1655 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1656 Py_INCREF(pool);
1657 PyTuple_SET_ITEM(pools, i, pool);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001658 indices[i] = 0;
1659 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001660
1661 /* create productobject structure */
1662 lz = (productobject *)type->tp_alloc(type, 0);
Christian Heimes380f7f22008-02-28 11:19:05 +00001663 if (lz == NULL)
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001664 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001665
1666 lz->pools = pools;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001667 lz->indices = indices;
1668 lz->result = NULL;
1669 lz->stopped = 0;
1670
1671 return (PyObject *)lz;
1672
1673error:
Christian Heimesc3f30c42008-02-22 16:37:40 +00001674 if (indices != NULL)
1675 PyMem_Free(indices);
1676 Py_XDECREF(pools);
1677 return NULL;
1678}
1679
1680static void
1681product_dealloc(productobject *lz)
1682{
1683 PyObject_GC_UnTrack(lz);
1684 Py_XDECREF(lz->pools);
1685 Py_XDECREF(lz->result);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001686 if (lz->indices != NULL)
1687 PyMem_Free(lz->indices);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001688 Py_TYPE(lz)->tp_free(lz);
1689}
1690
1691static int
1692product_traverse(productobject *lz, visitproc visit, void *arg)
1693{
1694 Py_VISIT(lz->pools);
1695 Py_VISIT(lz->result);
1696 return 0;
1697}
1698
1699static PyObject *
1700product_next(productobject *lz)
1701{
1702 PyObject *pool;
1703 PyObject *elem;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001704 PyObject *oldelem;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001705 PyObject *pools = lz->pools;
1706 PyObject *result = lz->result;
1707 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1708 Py_ssize_t i;
1709
1710 if (lz->stopped)
1711 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001712
Christian Heimesc3f30c42008-02-22 16:37:40 +00001713 if (result == NULL) {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001714 /* On the first pass, return an initial tuple filled with the
Christian Heimes78644762008-03-04 23:39:23 +00001715 first element from each pool. */
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001716 result = PyTuple_New(npools);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001717 if (result == NULL)
1718 goto empty;
1719 lz->result = result;
1720 for (i=0; i < npools; i++) {
1721 pool = PyTuple_GET_ITEM(pools, i);
1722 if (PyTuple_GET_SIZE(pool) == 0)
1723 goto empty;
1724 elem = PyTuple_GET_ITEM(pool, 0);
1725 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001726 PyTuple_SET_ITEM(result, i, elem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001727 }
1728 } else {
1729 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001730
1731 /* Copy the previous result tuple or re-use it if available */
1732 if (Py_REFCNT(result) > 1) {
1733 PyObject *old_result = result;
1734 result = PyTuple_New(npools);
1735 if (result == NULL)
1736 goto empty;
1737 lz->result = result;
1738 for (i=0; i < npools; i++) {
1739 elem = PyTuple_GET_ITEM(old_result, i);
1740 Py_INCREF(elem);
1741 PyTuple_SET_ITEM(result, i, elem);
1742 }
1743 Py_DECREF(old_result);
1744 }
1745 /* Now, we've got the only copy so we can update it in-place */
Christian Heimesb558a2e2008-03-02 22:46:37 +00001746 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001747
1748 /* Update the pool indices right-to-left. Only advance to the
1749 next pool when the previous one rolls-over */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001750 for (i=npools-1 ; i >= 0 ; i--) {
1751 pool = PyTuple_GET_ITEM(pools, i);
1752 indices[i]++;
Christian Heimesb558a2e2008-03-02 22:46:37 +00001753 if (indices[i] == PyTuple_GET_SIZE(pool)) {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001754 /* Roll-over and advance to next pool */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001755 indices[i] = 0;
1756 elem = PyTuple_GET_ITEM(pool, 0);
1757 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001758 oldelem = PyTuple_GET_ITEM(result, i);
1759 PyTuple_SET_ITEM(result, i, elem);
1760 Py_DECREF(oldelem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001761 } else {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001762 /* No rollover. Just increment and stop here. */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001763 elem = PyTuple_GET_ITEM(pool, indices[i]);
1764 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001765 oldelem = PyTuple_GET_ITEM(result, i);
1766 PyTuple_SET_ITEM(result, i, elem);
1767 Py_DECREF(oldelem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001768 break;
1769 }
1770 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001771
1772 /* If i is negative, then the indices have all rolled-over
1773 and we're done. */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001774 if (i < 0)
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001775 goto empty;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001776 }
1777
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001778 Py_INCREF(result);
1779 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001780
1781empty:
1782 lz->stopped = 1;
1783 return NULL;
1784}
1785
1786PyDoc_STRVAR(product_doc,
1787"product(*iterables) --> product object\n\
1788\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001789Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001790For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1791The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1792cycle in a manner similar to an odometer (with the rightmost element changing\n\
1793on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00001794To compute the product of an iterable with itself, specify the number\n\
1795of repetitions with the optional repeat keyword argument. For example,\n\
1796product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001797product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1798product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1799
1800static PyTypeObject product_type = {
1801 PyVarObject_HEAD_INIT(NULL, 0)
1802 "itertools.product", /* tp_name */
1803 sizeof(productobject), /* tp_basicsize */
1804 0, /* tp_itemsize */
1805 /* methods */
1806 (destructor)product_dealloc, /* tp_dealloc */
1807 0, /* tp_print */
1808 0, /* tp_getattr */
1809 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00001810 0, /* tp_reserved */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001811 0, /* tp_repr */
1812 0, /* tp_as_number */
1813 0, /* tp_as_sequence */
1814 0, /* tp_as_mapping */
1815 0, /* tp_hash */
1816 0, /* tp_call */
1817 0, /* tp_str */
1818 PyObject_GenericGetAttr, /* tp_getattro */
1819 0, /* tp_setattro */
1820 0, /* tp_as_buffer */
1821 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1822 Py_TPFLAGS_BASETYPE, /* tp_flags */
1823 product_doc, /* tp_doc */
1824 (traverseproc)product_traverse, /* tp_traverse */
1825 0, /* tp_clear */
1826 0, /* tp_richcompare */
1827 0, /* tp_weaklistoffset */
1828 PyObject_SelfIter, /* tp_iter */
1829 (iternextfunc)product_next, /* tp_iternext */
1830 0, /* tp_methods */
1831 0, /* tp_members */
1832 0, /* tp_getset */
1833 0, /* tp_base */
1834 0, /* tp_dict */
1835 0, /* tp_descr_get */
1836 0, /* tp_descr_set */
1837 0, /* tp_dictoffset */
1838 0, /* tp_init */
1839 0, /* tp_alloc */
1840 product_new, /* tp_new */
1841 PyObject_GC_Del, /* tp_free */
1842};
1843
1844
Christian Heimes380f7f22008-02-28 11:19:05 +00001845/* combinations object ************************************************************/
1846
1847typedef struct {
1848 PyObject_HEAD
1849 PyObject *pool; /* input converted to a tuple */
1850 Py_ssize_t *indices; /* one index per result element */
1851 PyObject *result; /* most recently returned result tuple */
1852 Py_ssize_t r; /* size of result tuple */
1853 int stopped; /* set to 1 when the combinations iterator is exhausted */
1854} combinationsobject;
1855
1856static PyTypeObject combinations_type;
1857
1858static PyObject *
1859combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1860{
1861 combinationsobject *co;
1862 Py_ssize_t n;
1863 Py_ssize_t r;
1864 PyObject *pool = NULL;
1865 PyObject *iterable = NULL;
1866 Py_ssize_t *indices = NULL;
1867 Py_ssize_t i;
1868 static char *kwargs[] = {"iterable", "r", NULL};
1869
1870 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
1871 &iterable, &r))
1872 return NULL;
1873
1874 pool = PySequence_Tuple(iterable);
1875 if (pool == NULL)
1876 goto error;
1877 n = PyTuple_GET_SIZE(pool);
1878 if (r < 0) {
1879 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
1880 goto error;
1881 }
Christian Heimes380f7f22008-02-28 11:19:05 +00001882
1883 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
1884 if (indices == NULL) {
1885 PyErr_NoMemory();
1886 goto error;
1887 }
1888
1889 for (i=0 ; i<r ; i++)
1890 indices[i] = i;
1891
1892 /* create combinationsobject structure */
1893 co = (combinationsobject *)type->tp_alloc(type, 0);
1894 if (co == NULL)
1895 goto error;
1896
1897 co->pool = pool;
1898 co->indices = indices;
1899 co->result = NULL;
1900 co->r = r;
Raymond Hettinger5bad41e2009-01-08 21:01:54 +00001901 co->stopped = r > n ? 1 : 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00001902
1903 return (PyObject *)co;
1904
1905error:
1906 if (indices != NULL)
1907 PyMem_Free(indices);
1908 Py_XDECREF(pool);
1909 return NULL;
1910}
1911
1912static void
1913combinations_dealloc(combinationsobject *co)
1914{
1915 PyObject_GC_UnTrack(co);
1916 Py_XDECREF(co->pool);
1917 Py_XDECREF(co->result);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001918 if (co->indices != NULL)
1919 PyMem_Free(co->indices);
Christian Heimes380f7f22008-02-28 11:19:05 +00001920 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 */
Mark Dickinsone94c6792009-02-02 20:36:42 +00002033 0, /* tp_reserved */
Christian Heimes380f7f22008-02-28 11:19:05 +00002034 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
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002068/* combinations with replacement object *******************************************/
2069
2070/* Equivalent to:
2071
2072 def combinations_with_replacement(iterable, r):
2073 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2074 # number items returned: (n+r-1)! / r! / (n-1)!
2075 pool = tuple(iterable)
2076 n = len(pool)
2077 indices = [0] * r
2078 yield tuple(pool[i] for i in indices)
2079 while 1:
2080 for i in reversed(range(r)):
2081 if indices[i] != n - 1:
2082 break
2083 else:
2084 return
2085 indices[i:] = [indices[i] + 1] * (r - i)
2086 yield tuple(pool[i] for i in indices)
2087
2088 def combinations_with_replacement2(iterable, r):
2089 'Alternate version that filters from product()'
2090 pool = tuple(iterable)
2091 n = len(pool)
2092 for indices in product(range(n), repeat=r):
2093 if sorted(indices) == list(indices):
2094 yield tuple(pool[i] for i in indices)
2095*/
2096typedef struct {
2097 PyObject_HEAD
2098 PyObject *pool; /* input converted to a tuple */
2099 Py_ssize_t *indices; /* one index per result element */
2100 PyObject *result; /* most recently returned result tuple */
2101 Py_ssize_t r; /* size of result tuple */
2102 int stopped; /* set to 1 when the cwr iterator is exhausted */
2103} cwrobject;
2104
2105static PyTypeObject cwr_type;
2106
2107static PyObject *
2108cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2109{
2110 cwrobject *co;
2111 Py_ssize_t n;
2112 Py_ssize_t r;
2113 PyObject *pool = NULL;
2114 PyObject *iterable = NULL;
2115 Py_ssize_t *indices = NULL;
2116 Py_ssize_t i;
2117 static char *kwargs[] = {"iterable", "r", NULL};
2118
2119 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2120 &iterable, &r))
2121 return NULL;
2122
2123 pool = PySequence_Tuple(iterable);
2124 if (pool == NULL)
2125 goto error;
2126 n = PyTuple_GET_SIZE(pool);
2127 if (r < 0) {
2128 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2129 goto error;
2130 }
2131
2132 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2133 if (indices == NULL) {
2134 PyErr_NoMemory();
2135 goto error;
2136 }
2137
2138 for (i=0 ; i<r ; i++)
2139 indices[i] = 0;
2140
2141 /* create cwrobject structure */
2142 co = (cwrobject *)type->tp_alloc(type, 0);
2143 if (co == NULL)
2144 goto error;
2145
2146 co->pool = pool;
2147 co->indices = indices;
2148 co->result = NULL;
2149 co->r = r;
2150 co->stopped = !n && r;
2151
2152 return (PyObject *)co;
2153
2154error:
2155 if (indices != NULL)
2156 PyMem_Free(indices);
2157 Py_XDECREF(pool);
2158 return NULL;
2159}
2160
2161static void
2162cwr_dealloc(cwrobject *co)
2163{
2164 PyObject_GC_UnTrack(co);
2165 Py_XDECREF(co->pool);
2166 Py_XDECREF(co->result);
2167 if (co->indices != NULL)
2168 PyMem_Free(co->indices);
2169 Py_TYPE(co)->tp_free(co);
2170}
2171
2172static int
2173cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2174{
2175 Py_VISIT(co->pool);
2176 Py_VISIT(co->result);
2177 return 0;
2178}
2179
2180static PyObject *
2181cwr_next(cwrobject *co)
2182{
2183 PyObject *elem;
2184 PyObject *oldelem;
2185 PyObject *pool = co->pool;
2186 Py_ssize_t *indices = co->indices;
2187 PyObject *result = co->result;
2188 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2189 Py_ssize_t r = co->r;
2190 Py_ssize_t i, j, index;
2191
2192 if (co->stopped)
2193 return NULL;
2194
2195 if (result == NULL) {
2196 /* On the first pass, initialize result tuple using the indices */
2197 result = PyTuple_New(r);
2198 if (result == NULL)
2199 goto empty;
2200 co->result = result;
2201 for (i=0; i<r ; i++) {
2202 index = indices[i];
2203 elem = PyTuple_GET_ITEM(pool, index);
2204 Py_INCREF(elem);
2205 PyTuple_SET_ITEM(result, i, elem);
2206 }
2207 } else {
2208 /* Copy the previous result tuple or re-use it if available */
2209 if (Py_REFCNT(result) > 1) {
2210 PyObject *old_result = result;
2211 result = PyTuple_New(r);
2212 if (result == NULL)
2213 goto empty;
2214 co->result = result;
2215 for (i=0; i<r ; i++) {
2216 elem = PyTuple_GET_ITEM(old_result, i);
2217 Py_INCREF(elem);
2218 PyTuple_SET_ITEM(result, i, elem);
2219 }
2220 Py_DECREF(old_result);
2221 }
2222 /* Now, we've got the only copy so we can update it in-place CPython's
2223 empty tuple is a singleton and cached in PyTuple's freelist. */
2224 assert(r == 0 || Py_REFCNT(result) == 1);
2225
2226 /* Scan indices right-to-left until finding one that is not
2227 * at its maximum (n-1). */
2228 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2229 ;
2230
2231 /* If i is negative, then the indices are all at
2232 their maximum value and we're done. */
2233 if (i < 0)
2234 goto empty;
2235
2236 /* Increment the current index which we know is not at its
2237 maximum. Then set all to the right to the same value. */
2238 indices[i]++;
2239 for (j=i+1 ; j<r ; j++)
2240 indices[j] = indices[j-1];
2241
2242 /* Update the result tuple for the new indices
2243 starting with i, the leftmost index that changed */
2244 for ( ; i<r ; i++) {
2245 index = indices[i];
2246 elem = PyTuple_GET_ITEM(pool, index);
2247 Py_INCREF(elem);
2248 oldelem = PyTuple_GET_ITEM(result, i);
2249 PyTuple_SET_ITEM(result, i, elem);
2250 Py_DECREF(oldelem);
2251 }
2252 }
2253
2254 Py_INCREF(result);
2255 return result;
2256
2257empty:
2258 co->stopped = 1;
2259 return NULL;
2260}
2261
2262PyDoc_STRVAR(cwr_doc,
2263"combinations_with_replacement(iterable[, r]) --> combinations_with_replacement object\n\
2264\n\
2265Return successive r-length combinations of elements in the iterable\n\
2266allowing individual elements to have successive repeats.\n\
2267combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2268
2269static PyTypeObject cwr_type = {
2270 PyVarObject_HEAD_INIT(NULL, 0)
2271 "itertools.combinations_with_replacement", /* tp_name */
2272 sizeof(cwrobject), /* tp_basicsize */
2273 0, /* tp_itemsize */
2274 /* methods */
2275 (destructor)cwr_dealloc, /* tp_dealloc */
2276 0, /* tp_print */
2277 0, /* tp_getattr */
2278 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00002279 0, /* tp_reserved */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002280 0, /* tp_repr */
2281 0, /* tp_as_number */
2282 0, /* tp_as_sequence */
2283 0, /* tp_as_mapping */
2284 0, /* tp_hash */
2285 0, /* tp_call */
2286 0, /* tp_str */
2287 PyObject_GenericGetAttr, /* tp_getattro */
2288 0, /* tp_setattro */
2289 0, /* tp_as_buffer */
2290 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2291 Py_TPFLAGS_BASETYPE, /* tp_flags */
2292 cwr_doc, /* tp_doc */
2293 (traverseproc)cwr_traverse, /* tp_traverse */
2294 0, /* tp_clear */
2295 0, /* tp_richcompare */
2296 0, /* tp_weaklistoffset */
2297 PyObject_SelfIter, /* tp_iter */
2298 (iternextfunc)cwr_next, /* tp_iternext */
2299 0, /* tp_methods */
2300 0, /* tp_members */
2301 0, /* tp_getset */
2302 0, /* tp_base */
2303 0, /* tp_dict */
2304 0, /* tp_descr_get */
2305 0, /* tp_descr_set */
2306 0, /* tp_dictoffset */
2307 0, /* tp_init */
2308 0, /* tp_alloc */
2309 cwr_new, /* tp_new */
2310 PyObject_GC_Del, /* tp_free */
2311};
2312
2313
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002314/* permutations object ************************************************************
2315
2316def permutations(iterable, r=None):
2317 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2318 pool = tuple(iterable)
2319 n = len(pool)
2320 r = n if r is None else r
2321 indices = range(n)
2322 cycles = range(n-r+1, n+1)[::-1]
2323 yield tuple(pool[i] for i in indices[:r])
2324 while n:
2325 for i in reversed(range(r)):
2326 cycles[i] -= 1
2327 if cycles[i] == 0:
2328 indices[i:] = indices[i+1:] + indices[i:i+1]
2329 cycles[i] = n - i
2330 else:
2331 j = cycles[i]
2332 indices[i], indices[-j] = indices[-j], indices[i]
2333 yield tuple(pool[i] for i in indices[:r])
2334 break
2335 else:
2336 return
2337*/
2338
2339typedef struct {
2340 PyObject_HEAD
2341 PyObject *pool; /* input converted to a tuple */
2342 Py_ssize_t *indices; /* one index per element in the pool */
2343 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2344 PyObject *result; /* most recently returned result tuple */
2345 Py_ssize_t r; /* size of result tuple */
2346 int stopped; /* set to 1 when the permutations iterator is exhausted */
2347} permutationsobject;
2348
2349static PyTypeObject permutations_type;
2350
2351static PyObject *
2352permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2353{
2354 permutationsobject *po;
2355 Py_ssize_t n;
2356 Py_ssize_t r;
2357 PyObject *robj = Py_None;
2358 PyObject *pool = NULL;
2359 PyObject *iterable = NULL;
2360 Py_ssize_t *indices = NULL;
2361 Py_ssize_t *cycles = NULL;
2362 Py_ssize_t i;
2363 static char *kwargs[] = {"iterable", "r", NULL};
2364
2365 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2366 &iterable, &robj))
2367 return NULL;
2368
2369 pool = PySequence_Tuple(iterable);
2370 if (pool == NULL)
2371 goto error;
2372 n = PyTuple_GET_SIZE(pool);
2373
2374 r = n;
2375 if (robj != Py_None) {
2376 if (!PyLong_Check(robj)) {
2377 PyErr_SetString(PyExc_TypeError, "Expected int as r");
Neal Norwitzd2bee322008-04-01 07:37:58 +00002378 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002379 }
2380 r = PyLong_AsSsize_t(robj);
2381 if (r == -1 && PyErr_Occurred())
2382 goto error;
2383 }
2384 if (r < 0) {
2385 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2386 goto error;
2387 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002388
2389 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2390 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2391 if (indices == NULL || cycles == NULL) {
2392 PyErr_NoMemory();
2393 goto error;
2394 }
2395
2396 for (i=0 ; i<n ; i++)
2397 indices[i] = i;
2398 for (i=0 ; i<r ; i++)
2399 cycles[i] = n - i;
2400
2401 /* create permutationsobject structure */
2402 po = (permutationsobject *)type->tp_alloc(type, 0);
2403 if (po == NULL)
2404 goto error;
2405
2406 po->pool = pool;
2407 po->indices = indices;
2408 po->cycles = cycles;
2409 po->result = NULL;
2410 po->r = r;
Raymond Hettinger5bad41e2009-01-08 21:01:54 +00002411 po->stopped = r > n ? 1 : 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002412
2413 return (PyObject *)po;
2414
2415error:
2416 if (indices != NULL)
2417 PyMem_Free(indices);
2418 if (cycles != NULL)
2419 PyMem_Free(cycles);
2420 Py_XDECREF(pool);
2421 return NULL;
2422}
2423
2424static void
2425permutations_dealloc(permutationsobject *po)
2426{
2427 PyObject_GC_UnTrack(po);
2428 Py_XDECREF(po->pool);
2429 Py_XDECREF(po->result);
2430 PyMem_Free(po->indices);
2431 PyMem_Free(po->cycles);
2432 Py_TYPE(po)->tp_free(po);
2433}
2434
2435static int
2436permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2437{
2438 Py_VISIT(po->pool);
2439 Py_VISIT(po->result);
2440 return 0;
2441}
2442
2443static PyObject *
2444permutations_next(permutationsobject *po)
2445{
2446 PyObject *elem;
2447 PyObject *oldelem;
2448 PyObject *pool = po->pool;
2449 Py_ssize_t *indices = po->indices;
2450 Py_ssize_t *cycles = po->cycles;
2451 PyObject *result = po->result;
2452 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2453 Py_ssize_t r = po->r;
2454 Py_ssize_t i, j, k, index;
2455
2456 if (po->stopped)
2457 return NULL;
2458
2459 if (result == NULL) {
2460 /* On the first pass, initialize result tuple using the indices */
2461 result = PyTuple_New(r);
2462 if (result == NULL)
2463 goto empty;
2464 po->result = result;
2465 for (i=0; i<r ; i++) {
2466 index = indices[i];
2467 elem = PyTuple_GET_ITEM(pool, index);
2468 Py_INCREF(elem);
2469 PyTuple_SET_ITEM(result, i, elem);
2470 }
2471 } else {
2472 if (n == 0)
2473 goto empty;
2474
2475 /* Copy the previous result tuple or re-use it if available */
2476 if (Py_REFCNT(result) > 1) {
2477 PyObject *old_result = result;
2478 result = PyTuple_New(r);
2479 if (result == NULL)
2480 goto empty;
2481 po->result = result;
2482 for (i=0; i<r ; i++) {
2483 elem = PyTuple_GET_ITEM(old_result, i);
2484 Py_INCREF(elem);
2485 PyTuple_SET_ITEM(result, i, elem);
2486 }
2487 Py_DECREF(old_result);
2488 }
2489 /* Now, we've got the only copy so we can update it in-place */
2490 assert(r == 0 || Py_REFCNT(result) == 1);
2491
2492 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2493 for (i=r-1 ; i>=0 ; i--) {
2494 cycles[i] -= 1;
2495 if (cycles[i] == 0) {
2496 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2497 index = indices[i];
2498 for (j=i ; j<n-1 ; j++)
2499 indices[j] = indices[j+1];
2500 indices[n-1] = index;
2501 cycles[i] = n - i;
2502 } else {
2503 j = cycles[i];
2504 index = indices[i];
2505 indices[i] = indices[n-j];
2506 indices[n-j] = index;
2507
2508 for (k=i; k<r ; k++) {
2509 /* start with i, the leftmost element that changed */
2510 /* yield tuple(pool[k] for k in indices[:r]) */
2511 index = indices[k];
2512 elem = PyTuple_GET_ITEM(pool, index);
2513 Py_INCREF(elem);
2514 oldelem = PyTuple_GET_ITEM(result, k);
2515 PyTuple_SET_ITEM(result, k, elem);
2516 Py_DECREF(oldelem);
2517 }
2518 break;
2519 }
2520 }
2521 /* If i is negative, then the cycles have all
2522 rolled-over and we're done. */
2523 if (i < 0)
2524 goto empty;
2525 }
2526 Py_INCREF(result);
2527 return result;
2528
2529empty:
2530 po->stopped = 1;
2531 return NULL;
2532}
2533
2534PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00002535"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002536\n\
2537Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00002538permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002539
2540static PyTypeObject permutations_type = {
2541 PyVarObject_HEAD_INIT(NULL, 0)
2542 "itertools.permutations", /* tp_name */
2543 sizeof(permutationsobject), /* tp_basicsize */
2544 0, /* tp_itemsize */
2545 /* methods */
2546 (destructor)permutations_dealloc, /* tp_dealloc */
2547 0, /* tp_print */
2548 0, /* tp_getattr */
2549 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00002550 0, /* tp_reserved */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002551 0, /* tp_repr */
2552 0, /* tp_as_number */
2553 0, /* tp_as_sequence */
2554 0, /* tp_as_mapping */
2555 0, /* tp_hash */
2556 0, /* tp_call */
2557 0, /* tp_str */
2558 PyObject_GenericGetAttr, /* tp_getattro */
2559 0, /* tp_setattro */
2560 0, /* tp_as_buffer */
2561 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2562 Py_TPFLAGS_BASETYPE, /* tp_flags */
2563 permutations_doc, /* tp_doc */
2564 (traverseproc)permutations_traverse, /* tp_traverse */
2565 0, /* tp_clear */
2566 0, /* tp_richcompare */
2567 0, /* tp_weaklistoffset */
2568 PyObject_SelfIter, /* tp_iter */
2569 (iternextfunc)permutations_next, /* tp_iternext */
2570 0, /* tp_methods */
2571 0, /* tp_members */
2572 0, /* tp_getset */
2573 0, /* tp_base */
2574 0, /* tp_dict */
2575 0, /* tp_descr_get */
2576 0, /* tp_descr_set */
2577 0, /* tp_dictoffset */
2578 0, /* tp_init */
2579 0, /* tp_alloc */
2580 permutations_new, /* tp_new */
2581 PyObject_GC_Del, /* tp_free */
2582};
2583
2584
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002585/* compress object ************************************************************/
2586
2587/* Equivalent to:
2588
2589 def compress(data, selectors):
2590 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2591 return (d for d, s in zip(data, selectors) if s)
2592*/
2593
2594typedef struct {
2595 PyObject_HEAD
2596 PyObject *data;
2597 PyObject *selectors;
2598} compressobject;
2599
2600static PyTypeObject compress_type;
2601
2602static PyObject *
2603compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2604{
2605 PyObject *seq1, *seq2;
2606 PyObject *data=NULL, *selectors=NULL;
2607 compressobject *lz;
2608
2609 if (type == &compress_type && !_PyArg_NoKeywords("compress()", kwds))
2610 return NULL;
2611
2612 if (!PyArg_UnpackTuple(args, "compress", 2, 2, &seq1, &seq2))
2613 return NULL;
2614
2615 data = PyObject_GetIter(seq1);
2616 if (data == NULL)
2617 goto fail;
2618 selectors = PyObject_GetIter(seq2);
2619 if (selectors == NULL)
2620 goto fail;
2621
2622 /* create compressobject structure */
2623 lz = (compressobject *)type->tp_alloc(type, 0);
2624 if (lz == NULL)
2625 goto fail;
2626 lz->data = data;
2627 lz->selectors = selectors;
2628 return (PyObject *)lz;
2629
2630fail:
2631 Py_XDECREF(data);
2632 Py_XDECREF(selectors);
2633 return NULL;
2634}
2635
2636static void
2637compress_dealloc(compressobject *lz)
2638{
2639 PyObject_GC_UnTrack(lz);
2640 Py_XDECREF(lz->data);
2641 Py_XDECREF(lz->selectors);
2642 Py_TYPE(lz)->tp_free(lz);
2643}
2644
2645static int
2646compress_traverse(compressobject *lz, visitproc visit, void *arg)
2647{
2648 Py_VISIT(lz->data);
2649 Py_VISIT(lz->selectors);
2650 return 0;
2651}
2652
2653static PyObject *
2654compress_next(compressobject *lz)
2655{
2656 PyObject *data = lz->data, *selectors = lz->selectors;
2657 PyObject *datum, *selector;
2658 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2659 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2660 int ok;
2661
2662 while (1) {
2663 /* Steps: get datum, get selector, evaluate selector.
2664 Order is important (to match the pure python version
2665 in terms of which input gets a chance to raise an
2666 exception first).
2667 */
2668
2669 datum = datanext(data);
2670 if (datum == NULL)
2671 return NULL;
2672
2673 selector = selectornext(selectors);
2674 if (selector == NULL) {
2675 Py_DECREF(datum);
2676 return NULL;
2677 }
2678
2679 ok = PyObject_IsTrue(selector);
2680 Py_DECREF(selector);
2681 if (ok == 1)
2682 return datum;
2683 Py_DECREF(datum);
2684 if (ok == -1)
2685 return NULL;
2686 }
2687}
2688
2689PyDoc_STRVAR(compress_doc,
2690"compress(data sequence, selector sequence) --> iterator over selected data\n\
2691\n\
2692Return data elements corresponding to true selector elements.\n\
2693Forms a shorter iterator from selected data elements using the\n\
2694selectors to choose the data elements.");
2695
2696static PyTypeObject compress_type = {
2697 PyVarObject_HEAD_INIT(NULL, 0)
2698 "itertools.compress", /* tp_name */
2699 sizeof(compressobject), /* tp_basicsize */
2700 0, /* tp_itemsize */
2701 /* methods */
2702 (destructor)compress_dealloc, /* tp_dealloc */
2703 0, /* tp_print */
2704 0, /* tp_getattr */
2705 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00002706 0, /* tp_reserved */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002707 0, /* tp_repr */
2708 0, /* tp_as_number */
2709 0, /* tp_as_sequence */
2710 0, /* tp_as_mapping */
2711 0, /* tp_hash */
2712 0, /* tp_call */
2713 0, /* tp_str */
2714 PyObject_GenericGetAttr, /* tp_getattro */
2715 0, /* tp_setattro */
2716 0, /* tp_as_buffer */
2717 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2718 Py_TPFLAGS_BASETYPE, /* tp_flags */
2719 compress_doc, /* tp_doc */
2720 (traverseproc)compress_traverse, /* tp_traverse */
2721 0, /* tp_clear */
2722 0, /* tp_richcompare */
2723 0, /* tp_weaklistoffset */
2724 PyObject_SelfIter, /* tp_iter */
2725 (iternextfunc)compress_next, /* tp_iternext */
2726 0, /* tp_methods */
2727 0, /* tp_members */
2728 0, /* tp_getset */
2729 0, /* tp_base */
2730 0, /* tp_dict */
2731 0, /* tp_descr_get */
2732 0, /* tp_descr_set */
2733 0, /* tp_dictoffset */
2734 0, /* tp_init */
2735 0, /* tp_alloc */
2736 compress_new, /* tp_new */
2737 PyObject_GC_Del, /* tp_free */
2738};
2739
2740
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002741/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00002742
2743typedef struct {
2744 PyObject_HEAD
2745 PyObject *func;
2746 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002747} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002748
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002749static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002750
2751static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002752filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002753{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00002754 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002755 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002756 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002757
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002758 if (type == &filterfalse_type &&
2759 !_PyArg_NoKeywords("filterfalse()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002760 return NULL;
2761
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002762 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
Raymond Hettinger60eca932003-02-09 06:40:58 +00002763 return NULL;
2764
2765 /* Get iterator. */
2766 it = PyObject_GetIter(seq);
2767 if (it == NULL)
2768 return NULL;
2769
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002770 /* create filterfalseobject structure */
2771 lz = (filterfalseobject *)type->tp_alloc(type, 0);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002772 if (lz == NULL) {
2773 Py_DECREF(it);
2774 return NULL;
2775 }
2776 Py_INCREF(func);
2777 lz->func = func;
2778 lz->it = it;
2779
2780 return (PyObject *)lz;
2781}
2782
2783static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002784filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002785{
2786 PyObject_GC_UnTrack(lz);
2787 Py_XDECREF(lz->func);
2788 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00002789 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002790}
2791
2792static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002793filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002794{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002795 Py_VISIT(lz->it);
2796 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002797 return 0;
2798}
2799
2800static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002801filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002802{
2803 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002804 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002805 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002806 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002807
Christian Heimes90aa7642007-12-19 02:45:37 +00002808 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002809 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002810 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002811 if (item == NULL)
2812 return NULL;
2813
Christian Heimes836baa52008-02-26 08:18:30 +00002814 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
Raymond Hettinger60eca932003-02-09 06:40:58 +00002815 ok = PyObject_IsTrue(item);
2816 } else {
2817 PyObject *good;
2818 good = PyObject_CallFunctionObjArgs(lz->func,
2819 item, NULL);
2820 if (good == NULL) {
2821 Py_DECREF(item);
2822 return NULL;
2823 }
2824 ok = PyObject_IsTrue(good);
2825 Py_DECREF(good);
2826 }
2827 if (!ok)
2828 return item;
2829 Py_DECREF(item);
2830 }
2831}
2832
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002833PyDoc_STRVAR(filterfalse_doc,
2834"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002835\n\
2836Return those items of sequence for which function(item) is false.\n\
2837If function is None, return the items that are false.");
2838
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002839static PyTypeObject filterfalse_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002840 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002841 "itertools.filterfalse", /* tp_name */
2842 sizeof(filterfalseobject), /* tp_basicsize */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002843 0, /* tp_itemsize */
2844 /* methods */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002845 (destructor)filterfalse_dealloc, /* tp_dealloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002846 0, /* tp_print */
2847 0, /* tp_getattr */
2848 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00002849 0, /* tp_reserved */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002850 0, /* tp_repr */
2851 0, /* tp_as_number */
2852 0, /* tp_as_sequence */
2853 0, /* tp_as_mapping */
2854 0, /* tp_hash */
2855 0, /* tp_call */
2856 0, /* tp_str */
2857 PyObject_GenericGetAttr, /* tp_getattro */
2858 0, /* tp_setattro */
2859 0, /* tp_as_buffer */
2860 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2861 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002862 filterfalse_doc, /* tp_doc */
2863 (traverseproc)filterfalse_traverse, /* tp_traverse */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002864 0, /* tp_clear */
2865 0, /* tp_richcompare */
2866 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002867 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002868 (iternextfunc)filterfalse_next, /* tp_iternext */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002869 0, /* tp_methods */
2870 0, /* tp_members */
2871 0, /* tp_getset */
2872 0, /* tp_base */
2873 0, /* tp_dict */
2874 0, /* tp_descr_get */
2875 0, /* tp_descr_set */
2876 0, /* tp_dictoffset */
2877 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002878 0, /* tp_alloc */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002879 filterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002880 PyObject_GC_Del, /* tp_free */
2881};
2882
2883
2884/* count object ************************************************************/
2885
2886typedef struct {
2887 PyObject_HEAD
Thomas Wouters477c8d52006-05-27 19:21:47 +00002888 Py_ssize_t cnt;
Raymond Hettinger30729212009-02-12 06:28:27 +00002889 PyObject *long_cnt;
2890 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002891} countobject;
2892
Raymond Hettinger30729212009-02-12 06:28:27 +00002893/* Counting logic and invariants:
2894
2895C_add_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
2896
2897 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
2898 Advances with: cnt += 1
2899 When count hits Y_SSIZE_T_MAX, switch to Py_add_mode.
2900
2901Py_add_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
2902
2903 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
2904 All counting is done with python objects (no overflows or underflows).
2905 Advances with: long_cnt += long_step
2906 Step may be zero -- effectively a slow version of repeat(cnt).
2907 Either long_cnt or long_step may be a float.
2908*/
2909
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002910static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002911
2912static PyObject *
2913count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2914{
2915 countobject *lz;
Thomas Wouters477c8d52006-05-27 19:21:47 +00002916 Py_ssize_t cnt = 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002917 PyObject *long_cnt = NULL;
Raymond Hettinger30729212009-02-12 06:28:27 +00002918 PyObject *long_step = NULL;
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +00002919 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002920
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +00002921 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
2922 kwlist, &long_cnt, &long_step))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002923 return NULL;
2924
Raymond Hettinger30729212009-02-12 06:28:27 +00002925 if (long_cnt != NULL && !PyNumber_Check(long_cnt) ||
2926 long_step != NULL && !PyNumber_Check(long_step)) {
2927 PyErr_SetString(PyExc_TypeError, "a number is required");
2928 return NULL;
2929 }
2930
2931 if (long_step == NULL) {
2932 /* If not specified, step defaults to 1 */
2933 long_step = PyLong_FromLong(1);
2934 if (long_step == NULL)
2935 return NULL;
2936 } else
2937 Py_INCREF(long_step);
2938 assert(long_step != NULL);
2939
2940 if (long_cnt != NULL) {
2941 cnt = PyLong_AsSsize_t(long_cnt);
2942 if ((cnt == -1 && PyErr_Occurred()) ||
2943 !PyIndex_Check(long_cnt) ||
2944 !PyLong_Check(long_step) ||
2945 PyLong_AS_LONG(long_step) != 1) {
2946 /* Switch to Py_add_mode */
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002947 PyErr_Clear();
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002948 Py_INCREF(long_cnt);
2949 cnt = PY_SSIZE_T_MAX;
Raymond Hettinger30729212009-02-12 06:28:27 +00002950 } else
2951 long_cnt = NULL;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002952 }
Raymond Hettinger30729212009-02-12 06:28:27 +00002953 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL ||
2954 cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002955
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002956 /* create countobject structure */
Raymond Hettingerd6280f42009-02-16 20:50:56 +00002957 lz = (countobject *)type->tp_alloc(type, 0);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002958 if (lz == NULL) {
2959 Py_XDECREF(long_cnt);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002960 return NULL;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002961 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002962 lz->cnt = cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002963 lz->long_cnt = long_cnt;
Raymond Hettinger30729212009-02-12 06:28:27 +00002964 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002965
2966 return (PyObject *)lz;
2967}
2968
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002969static void
2970count_dealloc(countobject *lz)
2971{
Raymond Hettingerd6280f42009-02-16 20:50:56 +00002972 PyObject_GC_UnTrack(lz);
Raymond Hettinger30729212009-02-12 06:28:27 +00002973 Py_XDECREF(lz->long_cnt);
2974 Py_XDECREF(lz->long_step);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00002975 Py_TYPE(lz)->tp_free(lz);
2976}
2977
2978count_traverse(countobject *lz, visitproc visit, void *arg)
2979{
2980 Py_VISIT(lz->long_cnt);
2981 Py_VISIT(lz->long_step);
2982 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002983}
2984
2985static PyObject *
2986count_nextlong(countobject *lz)
2987{
2988 static PyObject *one = NULL;
Raymond Hettinger30729212009-02-12 06:28:27 +00002989 PyObject *long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002990 PyObject *stepped_up;
2991
Raymond Hettinger30729212009-02-12 06:28:27 +00002992 long_cnt = lz->long_cnt;
2993 if (long_cnt == NULL) {
2994 /* Switch to Py_add_mode */
2995 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
2996 if (long_cnt == NULL)
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002997 return NULL;
2998 }
Raymond Hettinger30729212009-02-12 06:28:27 +00002999 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
3000
3001 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003002 if (stepped_up == NULL)
3003 return NULL;
3004 lz->long_cnt = stepped_up;
Raymond Hettinger30729212009-02-12 06:28:27 +00003005 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003006}
3007
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003008static PyObject *
3009count_next(countobject *lz)
3010{
Raymond Hettinger30729212009-02-12 06:28:27 +00003011 if (lz->cnt == PY_SSIZE_T_MAX)
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003012 return count_nextlong(lz);
Christian Heimes217cfd12007-12-02 14:31:20 +00003013 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003014}
3015
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003016static PyObject *
3017count_repr(countobject *lz)
3018{
Raymond Hettinger30729212009-02-12 06:28:27 +00003019 if (lz->cnt != PY_SSIZE_T_MAX)
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003020 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
3021
Raymond Hettinger30729212009-02-12 06:28:27 +00003022 if (PyLong_Check(lz->long_step)) {
3023 long step = PyLong_AsLong(lz->long_step);
3024 if (step == -1 && PyErr_Occurred()) {
3025 PyErr_Clear();
3026 }
3027 if (step == 1) {
3028 /* Don't display step when it is an integer equal to 1 */
3029 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
3030 }
3031 }
3032 return PyUnicode_FromFormat("count(%R, %R)",
3033 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003034}
3035
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003036PyDoc_STRVAR(count_doc,
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +00003037 "count([start[, step]]) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003038\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00003039Return a count object whose .__next__() method returns consecutive\n\
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +00003040integers starting from zero or, if specified, from start.\n\
Raymond Hettinger6dfff192009-02-12 10:19:59 +00003041If step is specified, counts by that interval. Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00003042 def count(firstval=0, step=1):\n\
3043 x = firstval\n\
Raymond Hettinger6dfff192009-02-12 10:19:59 +00003044 while 1:\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00003045 yield x\n\
Raymond Hettinger6dfff192009-02-12 10:19:59 +00003046 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003047
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003048static PyTypeObject count_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00003049 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003050 "itertools.count", /* tp_name */
3051 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003052 0, /* tp_itemsize */
3053 /* methods */
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003054 (destructor)count_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003055 0, /* tp_print */
3056 0, /* tp_getattr */
3057 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00003058 0, /* tp_reserved */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003059 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003060 0, /* tp_as_number */
3061 0, /* tp_as_sequence */
3062 0, /* tp_as_mapping */
3063 0, /* tp_hash */
3064 0, /* tp_call */
3065 0, /* tp_str */
3066 PyObject_GenericGetAttr, /* tp_getattro */
3067 0, /* tp_setattro */
3068 0, /* tp_as_buffer */
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003069 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3070 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003071 count_doc, /* tp_doc */
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003072 (traverseproc)count_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003073 0, /* tp_clear */
3074 0, /* tp_richcompare */
3075 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003076 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003077 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003078 0, /* tp_methods */
3079 0, /* tp_members */
3080 0, /* tp_getset */
3081 0, /* tp_base */
3082 0, /* tp_dict */
3083 0, /* tp_descr_get */
3084 0, /* tp_descr_set */
3085 0, /* tp_dictoffset */
3086 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003087 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003088 count_new, /* tp_new */
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003089 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003090};
3091
3092
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003093/* repeat object ************************************************************/
3094
3095typedef struct {
3096 PyObject_HEAD
3097 PyObject *element;
Thomas Wouters477c8d52006-05-27 19:21:47 +00003098 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003099} repeatobject;
3100
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003101static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003102
3103static PyObject *
3104repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3105{
3106 repeatobject *ro;
3107 PyObject *element;
Thomas Wouters477c8d52006-05-27 19:21:47 +00003108 Py_ssize_t cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003109
Thomas Woutersb2137042007-02-01 18:02:27 +00003110 if (type == &repeat_type && !_PyArg_NoKeywords("repeat()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00003111 return NULL;
3112
Thomas Wouters477c8d52006-05-27 19:21:47 +00003113 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003114 return NULL;
3115
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003116 if (PyTuple_Size(args) == 2 && cnt < 0)
3117 cnt = 0;
3118
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003119 ro = (repeatobject *)type->tp_alloc(type, 0);
3120 if (ro == NULL)
3121 return NULL;
3122 Py_INCREF(element);
3123 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003124 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003125 return (PyObject *)ro;
3126}
3127
3128static void
3129repeat_dealloc(repeatobject *ro)
3130{
3131 PyObject_GC_UnTrack(ro);
3132 Py_XDECREF(ro->element);
Christian Heimes90aa7642007-12-19 02:45:37 +00003133 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003134}
3135
3136static int
3137repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3138{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00003139 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003140 return 0;
3141}
3142
3143static PyObject *
3144repeat_next(repeatobject *ro)
3145{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003146 if (ro->cnt == 0)
3147 return NULL;
3148 if (ro->cnt > 0)
3149 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003150 Py_INCREF(ro->element);
3151 return ro->element;
3152}
3153
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003154static PyObject *
3155repeat_repr(repeatobject *ro)
3156{
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003157 if (ro->cnt == -1)
Walter Dörwald7569dfe2007-05-19 21:49:49 +00003158 return PyUnicode_FromFormat("repeat(%R)", ro->element);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003159 else
Walter Dörwald7569dfe2007-05-19 21:49:49 +00003160 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003161}
3162
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003163static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003164repeat_len(repeatobject *ro)
3165{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003166 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003167 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003168 return NULL;
3169 }
Christian Heimes217cfd12007-12-02 14:31:20 +00003170 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003171}
3172
Armin Rigof5b3e362006-02-11 21:32:43 +00003173PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003174
3175static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00003176 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003177 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003178};
3179
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003180PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003181"repeat(element [,times]) -> create an iterator which returns the element\n\
3182for the specified number of times. If not specified, returns the element\n\
3183endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003184
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003185static PyTypeObject repeat_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00003186 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003187 "itertools.repeat", /* tp_name */
3188 sizeof(repeatobject), /* tp_basicsize */
3189 0, /* tp_itemsize */
3190 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003191 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003192 0, /* tp_print */
3193 0, /* tp_getattr */
3194 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00003195 0, /* tp_reserved */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003196 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003197 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003198 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003199 0, /* tp_as_mapping */
3200 0, /* tp_hash */
3201 0, /* tp_call */
3202 0, /* tp_str */
3203 PyObject_GenericGetAttr, /* tp_getattro */
3204 0, /* tp_setattro */
3205 0, /* tp_as_buffer */
3206 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3207 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003208 repeat_doc, /* tp_doc */
3209 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003210 0, /* tp_clear */
3211 0, /* tp_richcompare */
3212 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003213 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003214 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003215 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003216 0, /* tp_members */
3217 0, /* tp_getset */
3218 0, /* tp_base */
3219 0, /* tp_dict */
3220 0, /* tp_descr_get */
3221 0, /* tp_descr_set */
3222 0, /* tp_dictoffset */
3223 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003224 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003225 repeat_new, /* tp_new */
3226 PyObject_GC_Del, /* tp_free */
3227};
3228
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003229/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00003230
3231#include "Python.h"
3232
3233typedef struct {
3234 PyObject_HEAD
3235 Py_ssize_t tuplesize;
3236 Py_ssize_t numactive;
3237 PyObject *ittuple; /* tuple of iterators */
3238 PyObject *result;
3239 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003240} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003241
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003242static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003243
3244static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003245zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003246{
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003247 ziplongestobject *lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003248 Py_ssize_t i;
3249 PyObject *ittuple; /* tuple of iterators */
3250 PyObject *result;
3251 PyObject *fillvalue = Py_None;
3252 Py_ssize_t tuplesize = PySequence_Length(args);
3253
3254 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3255 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3256 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3257 PyErr_SetString(PyExc_TypeError,
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003258 "zip_longest() got an unexpected keyword argument");
Thomas Wouterscf297e42007-02-23 15:07:44 +00003259 return NULL;
3260 }
3261 }
3262
3263 /* args must be a tuple */
3264 assert(PyTuple_Check(args));
3265
3266 /* obtain iterators */
3267 ittuple = PyTuple_New(tuplesize);
3268 if (ittuple == NULL)
3269 return NULL;
3270 for (i=0; i < tuplesize; ++i) {
3271 PyObject *item = PyTuple_GET_ITEM(args, i);
3272 PyObject *it = PyObject_GetIter(item);
3273 if (it == NULL) {
3274 if (PyErr_ExceptionMatches(PyExc_TypeError))
3275 PyErr_Format(PyExc_TypeError,
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003276 "zip_longest argument #%zd must support iteration",
Thomas Wouterscf297e42007-02-23 15:07:44 +00003277 i+1);
3278 Py_DECREF(ittuple);
3279 return NULL;
3280 }
3281 PyTuple_SET_ITEM(ittuple, i, it);
3282 }
3283
3284 /* create a result holder */
3285 result = PyTuple_New(tuplesize);
3286 if (result == NULL) {
3287 Py_DECREF(ittuple);
3288 return NULL;
3289 }
3290 for (i=0 ; i < tuplesize ; i++) {
3291 Py_INCREF(Py_None);
3292 PyTuple_SET_ITEM(result, i, Py_None);
3293 }
3294
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003295 /* create ziplongestobject structure */
3296 lz = (ziplongestobject *)type->tp_alloc(type, 0);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003297 if (lz == NULL) {
3298 Py_DECREF(ittuple);
3299 Py_DECREF(result);
3300 return NULL;
3301 }
3302 lz->ittuple = ittuple;
3303 lz->tuplesize = tuplesize;
3304 lz->numactive = tuplesize;
3305 lz->result = result;
3306 Py_INCREF(fillvalue);
3307 lz->fillvalue = fillvalue;
3308 return (PyObject *)lz;
3309}
3310
3311static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003312zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003313{
3314 PyObject_GC_UnTrack(lz);
3315 Py_XDECREF(lz->ittuple);
3316 Py_XDECREF(lz->result);
3317 Py_XDECREF(lz->fillvalue);
Christian Heimes90aa7642007-12-19 02:45:37 +00003318 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003319}
3320
3321static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003322zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003323{
3324 Py_VISIT(lz->ittuple);
3325 Py_VISIT(lz->result);
3326 Py_VISIT(lz->fillvalue);
3327 return 0;
3328}
3329
3330static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003331zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003332{
3333 Py_ssize_t i;
3334 Py_ssize_t tuplesize = lz->tuplesize;
3335 PyObject *result = lz->result;
3336 PyObject *it;
3337 PyObject *item;
3338 PyObject *olditem;
3339
3340 if (tuplesize == 0)
3341 return NULL;
3342 if (lz->numactive == 0)
3343 return NULL;
Christian Heimes90aa7642007-12-19 02:45:37 +00003344 if (Py_REFCNT(result) == 1) {
Thomas Wouterscf297e42007-02-23 15:07:44 +00003345 Py_INCREF(result);
3346 for (i=0 ; i < tuplesize ; i++) {
3347 it = PyTuple_GET_ITEM(lz->ittuple, i);
3348 if (it == NULL) {
3349 Py_INCREF(lz->fillvalue);
3350 item = lz->fillvalue;
3351 } else {
Christian Heimes90aa7642007-12-19 02:45:37 +00003352 item = (*Py_TYPE(it)->tp_iternext)(it);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003353 if (item == NULL) {
3354 lz->numactive -= 1;
3355 if (lz->numactive == 0) {
3356 Py_DECREF(result);
3357 return NULL;
3358 } else {
3359 Py_INCREF(lz->fillvalue);
3360 item = lz->fillvalue;
3361 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3362 Py_DECREF(it);
3363 }
3364 }
3365 }
3366 olditem = PyTuple_GET_ITEM(result, i);
3367 PyTuple_SET_ITEM(result, i, item);
3368 Py_DECREF(olditem);
3369 }
3370 } else {
3371 result = PyTuple_New(tuplesize);
3372 if (result == NULL)
3373 return NULL;
3374 for (i=0 ; i < tuplesize ; i++) {
3375 it = PyTuple_GET_ITEM(lz->ittuple, i);
3376 if (it == NULL) {
3377 Py_INCREF(lz->fillvalue);
3378 item = lz->fillvalue;
3379 } else {
Christian Heimes90aa7642007-12-19 02:45:37 +00003380 item = (*Py_TYPE(it)->tp_iternext)(it);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003381 if (item == NULL) {
3382 lz->numactive -= 1;
3383 if (lz->numactive == 0) {
3384 Py_DECREF(result);
3385 return NULL;
3386 } else {
3387 Py_INCREF(lz->fillvalue);
3388 item = lz->fillvalue;
3389 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3390 Py_DECREF(it);
3391 }
3392 }
3393 }
3394 PyTuple_SET_ITEM(result, i, item);
3395 }
3396 }
3397 return result;
3398}
3399
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003400PyDoc_STRVAR(zip_longest_doc,
3401"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003402\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003403Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00003404the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003405method continues until the longest iterable in the argument sequence\n\
3406is exhausted and then it raises StopIteration. When the shorter iterables\n\
3407are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3408defaults to None or can be specified by a keyword argument.\n\
3409");
3410
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003411static PyTypeObject ziplongest_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00003412 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003413 "itertools.zip_longest", /* tp_name */
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003414 sizeof(ziplongestobject), /* tp_basicsize */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003415 0, /* tp_itemsize */
3416 /* methods */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003417 (destructor)zip_longest_dealloc, /* tp_dealloc */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003418 0, /* tp_print */
3419 0, /* tp_getattr */
3420 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00003421 0, /* tp_reserved */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003422 0, /* tp_repr */
3423 0, /* tp_as_number */
3424 0, /* tp_as_sequence */
3425 0, /* tp_as_mapping */
3426 0, /* tp_hash */
3427 0, /* tp_call */
3428 0, /* tp_str */
3429 PyObject_GenericGetAttr, /* tp_getattro */
3430 0, /* tp_setattro */
3431 0, /* tp_as_buffer */
3432 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3433 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003434 zip_longest_doc, /* tp_doc */
3435 (traverseproc)zip_longest_traverse, /* tp_traverse */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003436 0, /* tp_clear */
3437 0, /* tp_richcompare */
3438 0, /* tp_weaklistoffset */
3439 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003440 (iternextfunc)zip_longest_next, /* tp_iternext */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003441 0, /* tp_methods */
3442 0, /* tp_members */
3443 0, /* tp_getset */
3444 0, /* tp_base */
3445 0, /* tp_dict */
3446 0, /* tp_descr_get */
3447 0, /* tp_descr_set */
3448 0, /* tp_dictoffset */
3449 0, /* tp_init */
3450 0, /* tp_alloc */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003451 zip_longest_new, /* tp_new */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003452 PyObject_GC_Del, /* tp_free */
3453};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003454
3455/* module level code ********************************************************/
3456
3457PyDoc_STRVAR(module_doc,
3458"Functional tools for creating and using iterators.\n\
3459\n\
3460Infinite iterators:\n\
3461count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003462cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003463repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003464\n\
3465Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003466chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
3467compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
3468dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
3469groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003470filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003471islice(seq, [start,] stop [, step]) --> elements from\n\
3472 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003473starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003474tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003475takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003476zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00003477\n\
3478Combinatoric generators:\n\
3479product(p, q, ... [repeat=1]) --> cartesian product\n\
3480permutations(p[, r])\n\
3481combinations(p[, r])\n\
3482combinations_with_replacement(p[, r])\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003483");
3484
3485
Raymond Hettingerad983e72003-11-12 14:32:26 +00003486static PyMethodDef module_methods[] = {
3487 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3488 {NULL, NULL} /* sentinel */
3489};
3490
Martin v. Löwis1a214512008-06-11 05:26:20 +00003491
3492static struct PyModuleDef itertoolsmodule = {
3493 PyModuleDef_HEAD_INIT,
3494 "itertools",
3495 module_doc,
3496 -1,
3497 module_methods,
3498 NULL,
3499 NULL,
3500 NULL,
3501 NULL
3502};
3503
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003504PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00003505PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003506{
Raymond Hettinger60eca932003-02-09 06:40:58 +00003507 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003508 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003509 char *name;
3510 PyTypeObject *typelist[] = {
Christian Heimes380f7f22008-02-28 11:19:05 +00003511 &combinations_type,
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003512 &cwr_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003513 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003514 &dropwhile_type,
3515 &takewhile_type,
3516 &islice_type,
3517 &starmap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003518 &chain_type,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003519 &compress_type,
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003520 &filterfalse_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003521 &count_type,
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003522 &ziplongest_type,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003523 &permutations_type,
Christian Heimes380f7f22008-02-28 11:19:05 +00003524 &product_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003525 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003526 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003527 NULL
3528 };
3529
Christian Heimes90aa7642007-12-19 02:45:37 +00003530 Py_TYPE(&teedataobject_type) = &PyType_Type;
Martin v. Löwis1a214512008-06-11 05:26:20 +00003531 m = PyModule_Create(&itertoolsmodule);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00003532 if (m == NULL)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003533 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003534
Raymond Hettinger60eca932003-02-09 06:40:58 +00003535 for (i=0 ; typelist[i] != NULL ; i++) {
3536 if (PyType_Ready(typelist[i]) < 0)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003537 return NULL;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003538 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003539 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003540 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003541 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003542 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003543
3544 if (PyType_Ready(&teedataobject_type) < 0)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003545 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +00003546 if (PyType_Ready(&tee_type) < 0)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003547 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003548 if (PyType_Ready(&_grouper_type) < 0)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003549 return NULL;
3550 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003551}