blob: 74396b68b866ce0340a0d63f9794cdc673a497e8 [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) {
Benjamin Peterson25e26a02009-02-16 21:28:29 +0000748 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
749 Py_DECREF(item);
750 return NULL;
751 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000752 return item;
753 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000754 if (PyErr_Occurred()) {
755 if (PyErr_ExceptionMatches(PyExc_StopIteration))
756 PyErr_Clear();
757 else
758 return NULL;
759 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000760 if (PyList_Size(lz->saved) == 0)
761 return NULL;
762 it = PyObject_GetIter(lz->saved);
763 if (it == NULL)
764 return NULL;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000765 tmp = lz->it;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000766 lz->it = it;
767 lz->firstpass = 1;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000768 Py_DECREF(tmp);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000769 }
770}
771
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000772PyDoc_STRVAR(cycle_doc,
773"cycle(iterable) --> cycle object\n\
774\n\
775Return elements from the iterable until it is exhausted.\n\
776Then repeat the sequence indefinitely.");
777
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000778static PyTypeObject cycle_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +0000779 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000780 "itertools.cycle", /* tp_name */
781 sizeof(cycleobject), /* tp_basicsize */
782 0, /* tp_itemsize */
783 /* methods */
784 (destructor)cycle_dealloc, /* tp_dealloc */
785 0, /* tp_print */
786 0, /* tp_getattr */
787 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +0000788 0, /* tp_reserved */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000789 0, /* tp_repr */
790 0, /* tp_as_number */
791 0, /* tp_as_sequence */
792 0, /* tp_as_mapping */
793 0, /* tp_hash */
794 0, /* tp_call */
795 0, /* tp_str */
796 PyObject_GenericGetAttr, /* tp_getattro */
797 0, /* tp_setattro */
798 0, /* tp_as_buffer */
799 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
800 Py_TPFLAGS_BASETYPE, /* tp_flags */
801 cycle_doc, /* tp_doc */
802 (traverseproc)cycle_traverse, /* tp_traverse */
803 0, /* tp_clear */
804 0, /* tp_richcompare */
805 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000806 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000807 (iternextfunc)cycle_next, /* tp_iternext */
808 0, /* tp_methods */
809 0, /* tp_members */
810 0, /* tp_getset */
811 0, /* tp_base */
812 0, /* tp_dict */
813 0, /* tp_descr_get */
814 0, /* tp_descr_set */
815 0, /* tp_dictoffset */
816 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000817 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000818 cycle_new, /* tp_new */
819 PyObject_GC_Del, /* tp_free */
820};
821
822
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000823/* dropwhile object **********************************************************/
824
825typedef struct {
826 PyObject_HEAD
827 PyObject *func;
828 PyObject *it;
829 long start;
830} dropwhileobject;
831
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000832static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000833
834static PyObject *
835dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
836{
837 PyObject *func, *seq;
838 PyObject *it;
839 dropwhileobject *lz;
840
Thomas Woutersb2137042007-02-01 18:02:27 +0000841 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000842 return NULL;
843
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000844 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
845 return NULL;
846
847 /* Get iterator. */
848 it = PyObject_GetIter(seq);
849 if (it == NULL)
850 return NULL;
851
852 /* create dropwhileobject structure */
853 lz = (dropwhileobject *)type->tp_alloc(type, 0);
854 if (lz == NULL) {
855 Py_DECREF(it);
856 return NULL;
857 }
858 Py_INCREF(func);
859 lz->func = func;
860 lz->it = it;
861 lz->start = 0;
862
863 return (PyObject *)lz;
864}
865
866static void
867dropwhile_dealloc(dropwhileobject *lz)
868{
869 PyObject_GC_UnTrack(lz);
870 Py_XDECREF(lz->func);
871 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +0000872 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000873}
874
875static int
876dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
877{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000878 Py_VISIT(lz->it);
879 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000880 return 0;
881}
882
883static PyObject *
884dropwhile_next(dropwhileobject *lz)
885{
886 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000887 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000888 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000889 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000890
Christian Heimes90aa7642007-12-19 02:45:37 +0000891 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000892 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000893 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000894 if (item == NULL)
895 return NULL;
896 if (lz->start == 1)
897 return item;
898
899 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
900 if (good == NULL) {
901 Py_DECREF(item);
902 return NULL;
903 }
904 ok = PyObject_IsTrue(good);
905 Py_DECREF(good);
906 if (!ok) {
907 lz->start = 1;
908 return item;
909 }
910 Py_DECREF(item);
911 }
912}
913
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000914PyDoc_STRVAR(dropwhile_doc,
915"dropwhile(predicate, iterable) --> dropwhile object\n\
916\n\
917Drop items from the iterable while predicate(item) is true.\n\
918Afterwards, return every element until the iterable is exhausted.");
919
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000920static PyTypeObject dropwhile_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +0000921 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000922 "itertools.dropwhile", /* tp_name */
923 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000924 0, /* tp_itemsize */
925 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000926 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000927 0, /* tp_print */
928 0, /* tp_getattr */
929 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +0000930 0, /* tp_reserved */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000931 0, /* tp_repr */
932 0, /* tp_as_number */
933 0, /* tp_as_sequence */
934 0, /* tp_as_mapping */
935 0, /* tp_hash */
936 0, /* tp_call */
937 0, /* tp_str */
938 PyObject_GenericGetAttr, /* tp_getattro */
939 0, /* tp_setattro */
940 0, /* tp_as_buffer */
941 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
942 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000943 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000944 (traverseproc)dropwhile_traverse, /* tp_traverse */
945 0, /* tp_clear */
946 0, /* tp_richcompare */
947 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000948 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000949 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000950 0, /* tp_methods */
951 0, /* tp_members */
952 0, /* tp_getset */
953 0, /* tp_base */
954 0, /* tp_dict */
955 0, /* tp_descr_get */
956 0, /* tp_descr_set */
957 0, /* tp_dictoffset */
958 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000959 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000960 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000961 PyObject_GC_Del, /* tp_free */
962};
963
964
965/* takewhile object **********************************************************/
966
967typedef struct {
968 PyObject_HEAD
969 PyObject *func;
970 PyObject *it;
971 long stop;
972} takewhileobject;
973
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000974static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000975
976static PyObject *
977takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
978{
979 PyObject *func, *seq;
980 PyObject *it;
981 takewhileobject *lz;
982
Thomas Woutersb2137042007-02-01 18:02:27 +0000983 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000984 return NULL;
985
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000986 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
987 return NULL;
988
989 /* Get iterator. */
990 it = PyObject_GetIter(seq);
991 if (it == NULL)
992 return NULL;
993
994 /* create takewhileobject structure */
995 lz = (takewhileobject *)type->tp_alloc(type, 0);
996 if (lz == NULL) {
997 Py_DECREF(it);
998 return NULL;
999 }
1000 Py_INCREF(func);
1001 lz->func = func;
1002 lz->it = it;
1003 lz->stop = 0;
1004
1005 return (PyObject *)lz;
1006}
1007
1008static void
1009takewhile_dealloc(takewhileobject *lz)
1010{
1011 PyObject_GC_UnTrack(lz);
1012 Py_XDECREF(lz->func);
1013 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00001014 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001015}
1016
1017static int
1018takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1019{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001020 Py_VISIT(lz->it);
1021 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001022 return 0;
1023}
1024
1025static PyObject *
1026takewhile_next(takewhileobject *lz)
1027{
1028 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001029 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001030 long ok;
1031
1032 if (lz->stop == 1)
1033 return NULL;
1034
Christian Heimes90aa7642007-12-19 02:45:37 +00001035 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001036 if (item == NULL)
1037 return NULL;
1038
1039 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1040 if (good == NULL) {
1041 Py_DECREF(item);
1042 return NULL;
1043 }
1044 ok = PyObject_IsTrue(good);
1045 Py_DECREF(good);
1046 if (ok)
1047 return item;
1048 Py_DECREF(item);
1049 lz->stop = 1;
1050 return NULL;
1051}
1052
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053PyDoc_STRVAR(takewhile_doc,
1054"takewhile(predicate, iterable) --> takewhile object\n\
1055\n\
1056Return successive entries from an iterable as long as the \n\
1057predicate evaluates to true for each entry.");
1058
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001059static PyTypeObject takewhile_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001060 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001061 "itertools.takewhile", /* tp_name */
1062 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001063 0, /* tp_itemsize */
1064 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001065 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001066 0, /* tp_print */
1067 0, /* tp_getattr */
1068 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00001069 0, /* tp_reserved */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001070 0, /* tp_repr */
1071 0, /* tp_as_number */
1072 0, /* tp_as_sequence */
1073 0, /* tp_as_mapping */
1074 0, /* tp_hash */
1075 0, /* tp_call */
1076 0, /* tp_str */
1077 PyObject_GenericGetAttr, /* tp_getattro */
1078 0, /* tp_setattro */
1079 0, /* tp_as_buffer */
1080 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1081 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001082 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001083 (traverseproc)takewhile_traverse, /* tp_traverse */
1084 0, /* tp_clear */
1085 0, /* tp_richcompare */
1086 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001087 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001088 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001089 0, /* tp_methods */
1090 0, /* tp_members */
1091 0, /* tp_getset */
1092 0, /* tp_base */
1093 0, /* tp_dict */
1094 0, /* tp_descr_get */
1095 0, /* tp_descr_set */
1096 0, /* tp_dictoffset */
1097 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001098 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001099 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001100 PyObject_GC_Del, /* tp_free */
1101};
1102
1103
1104/* islice object ************************************************************/
1105
1106typedef struct {
1107 PyObject_HEAD
1108 PyObject *it;
Thomas Wouters477c8d52006-05-27 19:21:47 +00001109 Py_ssize_t next;
1110 Py_ssize_t stop;
1111 Py_ssize_t step;
1112 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001113} isliceobject;
1114
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001115static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116
1117static PyObject *
1118islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1119{
1120 PyObject *seq;
Thomas Wouters477c8d52006-05-27 19:21:47 +00001121 Py_ssize_t start=0, stop=-1, step=1;
Raymond Hettingerb2594052004-12-05 09:25:51 +00001122 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001123 Py_ssize_t numargs;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001124 isliceobject *lz;
1125
Thomas Woutersb2137042007-02-01 18:02:27 +00001126 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001127 return NULL;
1128
Raymond Hettingerb2594052004-12-05 09:25:51 +00001129 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001130 return NULL;
1131
Raymond Hettingerb2594052004-12-05 09:25:51 +00001132 numargs = PyTuple_Size(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001133 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001134 if (a1 != Py_None) {
Christian Heimes217cfd12007-12-02 14:31:20 +00001135 stop = PyLong_AsSsize_t(a1);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001136 if (stop == -1) {
1137 if (PyErr_Occurred())
1138 PyErr_Clear();
1139 PyErr_SetString(PyExc_ValueError,
Benjamin Petersonbbcd1eb2009-06-24 02:29:57 +00001140 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001141 return NULL;
1142 }
1143 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001144 } else {
Raymond Hettingerb2594052004-12-05 09:25:51 +00001145 if (a1 != Py_None)
Christian Heimes217cfd12007-12-02 14:31:20 +00001146 start = PyLong_AsSsize_t(a1);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001147 if (start == -1 && PyErr_Occurred())
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001148 PyErr_Clear();
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001149 if (a2 != Py_None) {
Christian Heimes217cfd12007-12-02 14:31:20 +00001150 stop = PyLong_AsSsize_t(a2);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001151 if (stop == -1) {
1152 if (PyErr_Occurred())
1153 PyErr_Clear();
1154 PyErr_SetString(PyExc_ValueError,
Benjamin Petersonbbcd1eb2009-06-24 02:29:57 +00001155 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001156 return NULL;
1157 }
1158 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001159 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001160 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001161 PyErr_SetString(PyExc_ValueError,
Benjamin Petersonbbcd1eb2009-06-24 02:29:57 +00001162 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001163 return NULL;
1164 }
1165
Raymond Hettingerb2594052004-12-05 09:25:51 +00001166 if (a3 != NULL) {
1167 if (a3 != Py_None)
Christian Heimes217cfd12007-12-02 14:31:20 +00001168 step = PyLong_AsSsize_t(a3);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001169 if (step == -1 && PyErr_Occurred())
1170 PyErr_Clear();
1171 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001172 if (step<1) {
1173 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001174 "Step for islice() must be a positive integer or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001175 return NULL;
1176 }
1177
1178 /* Get iterator. */
1179 it = PyObject_GetIter(seq);
1180 if (it == NULL)
1181 return NULL;
1182
1183 /* create isliceobject structure */
1184 lz = (isliceobject *)type->tp_alloc(type, 0);
1185 if (lz == NULL) {
1186 Py_DECREF(it);
1187 return NULL;
1188 }
1189 lz->it = it;
1190 lz->next = start;
1191 lz->stop = stop;
1192 lz->step = step;
1193 lz->cnt = 0L;
1194
1195 return (PyObject *)lz;
1196}
1197
1198static void
1199islice_dealloc(isliceobject *lz)
1200{
1201 PyObject_GC_UnTrack(lz);
1202 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00001203 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204}
1205
1206static int
1207islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1208{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001209 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001210 return 0;
1211}
1212
1213static PyObject *
1214islice_next(isliceobject *lz)
1215{
1216 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001217 PyObject *it = lz->it;
Thomas Wouters477c8d52006-05-27 19:21:47 +00001218 Py_ssize_t oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001219 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001220
Christian Heimes90aa7642007-12-19 02:45:37 +00001221 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001222 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001223 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001224 if (item == NULL)
1225 return NULL;
1226 Py_DECREF(item);
1227 lz->cnt++;
1228 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001229 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001230 return NULL;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001231 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001232 if (item == NULL)
1233 return NULL;
1234 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001235 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001236 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001237 if (lz->next < oldnext) /* Check for overflow */
1238 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001239 return item;
1240}
1241
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001242PyDoc_STRVAR(islice_doc,
1243"islice(iterable, [start,] stop [, step]) --> islice object\n\
1244\n\
1245Return an iterator whose next() method returns selected values from an\n\
1246iterable. If start is specified, will skip all preceding elements;\n\
1247otherwise, start defaults to zero. Step defaults to one. If\n\
1248specified as another value, step determines how many values are \n\
1249skipped between successive calls. Works like a slice() on a list\n\
1250but returns an iterator.");
1251
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001252static PyTypeObject islice_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001253 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001254 "itertools.islice", /* tp_name */
1255 sizeof(isliceobject), /* tp_basicsize */
1256 0, /* tp_itemsize */
1257 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001258 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001259 0, /* tp_print */
1260 0, /* tp_getattr */
1261 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00001262 0, /* tp_reserved */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001263 0, /* tp_repr */
1264 0, /* tp_as_number */
1265 0, /* tp_as_sequence */
1266 0, /* tp_as_mapping */
1267 0, /* tp_hash */
1268 0, /* tp_call */
1269 0, /* tp_str */
1270 PyObject_GenericGetAttr, /* tp_getattro */
1271 0, /* tp_setattro */
1272 0, /* tp_as_buffer */
1273 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1274 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001275 islice_doc, /* tp_doc */
1276 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001277 0, /* tp_clear */
1278 0, /* tp_richcompare */
1279 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001280 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001281 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001282 0, /* tp_methods */
1283 0, /* tp_members */
1284 0, /* tp_getset */
1285 0, /* tp_base */
1286 0, /* tp_dict */
1287 0, /* tp_descr_get */
1288 0, /* tp_descr_set */
1289 0, /* tp_dictoffset */
1290 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001291 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001292 islice_new, /* tp_new */
1293 PyObject_GC_Del, /* tp_free */
1294};
1295
1296
1297/* starmap object ************************************************************/
1298
1299typedef struct {
1300 PyObject_HEAD
1301 PyObject *func;
1302 PyObject *it;
1303} starmapobject;
1304
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001305static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001306
1307static PyObject *
1308starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1309{
1310 PyObject *func, *seq;
1311 PyObject *it;
1312 starmapobject *lz;
1313
Thomas Woutersb2137042007-02-01 18:02:27 +00001314 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001315 return NULL;
1316
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001317 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1318 return NULL;
1319
1320 /* Get iterator. */
1321 it = PyObject_GetIter(seq);
1322 if (it == NULL)
1323 return NULL;
1324
1325 /* create starmapobject structure */
1326 lz = (starmapobject *)type->tp_alloc(type, 0);
1327 if (lz == NULL) {
1328 Py_DECREF(it);
1329 return NULL;
1330 }
1331 Py_INCREF(func);
1332 lz->func = func;
1333 lz->it = it;
1334
1335 return (PyObject *)lz;
1336}
1337
1338static void
1339starmap_dealloc(starmapobject *lz)
1340{
1341 PyObject_GC_UnTrack(lz);
1342 Py_XDECREF(lz->func);
1343 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00001344 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001345}
1346
1347static int
1348starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1349{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001350 Py_VISIT(lz->it);
1351 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001352 return 0;
1353}
1354
1355static PyObject *
1356starmap_next(starmapobject *lz)
1357{
1358 PyObject *args;
1359 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001360 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001361
Christian Heimes90aa7642007-12-19 02:45:37 +00001362 args = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001363 if (args == NULL)
1364 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001365 if (!PyTuple_CheckExact(args)) {
Christian Heimes679db4a2008-01-18 09:56:22 +00001366 PyObject *newargs = PySequence_Tuple(args);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001367 Py_DECREF(args);
Christian Heimes679db4a2008-01-18 09:56:22 +00001368 if (newargs == NULL)
1369 return NULL;
1370 args = newargs;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001371 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001372 result = PyObject_Call(lz->func, args, NULL);
1373 Py_DECREF(args);
1374 return result;
1375}
1376
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001377PyDoc_STRVAR(starmap_doc,
1378"starmap(function, sequence) --> starmap object\n\
1379\n\
1380Return an iterator whose values are returned from the function evaluated\n\
1381with a argument tuple taken from the given sequence.");
1382
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001383static PyTypeObject starmap_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001384 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001385 "itertools.starmap", /* tp_name */
1386 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001387 0, /* tp_itemsize */
1388 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001389 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001390 0, /* tp_print */
1391 0, /* tp_getattr */
1392 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00001393 0, /* tp_reserved */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001394 0, /* tp_repr */
1395 0, /* tp_as_number */
1396 0, /* tp_as_sequence */
1397 0, /* tp_as_mapping */
1398 0, /* tp_hash */
1399 0, /* tp_call */
1400 0, /* tp_str */
1401 PyObject_GenericGetAttr, /* tp_getattro */
1402 0, /* tp_setattro */
1403 0, /* tp_as_buffer */
1404 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1405 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001406 starmap_doc, /* tp_doc */
1407 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001408 0, /* tp_clear */
1409 0, /* tp_richcompare */
1410 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001411 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001412 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001413 0, /* tp_methods */
1414 0, /* tp_members */
1415 0, /* tp_getset */
1416 0, /* tp_base */
1417 0, /* tp_dict */
1418 0, /* tp_descr_get */
1419 0, /* tp_descr_set */
1420 0, /* tp_dictoffset */
1421 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001422 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001423 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001424 PyObject_GC_Del, /* tp_free */
1425};
1426
1427
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001428/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001429
1430typedef struct {
1431 PyObject_HEAD
Christian Heimesf16baeb2008-02-29 14:57:44 +00001432 PyObject *source; /* Iterator over input iterables */
1433 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001434} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001435
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001436static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001437
Christian Heimesf16baeb2008-02-29 14:57:44 +00001438static PyObject *
1439chain_new_internal(PyTypeObject *type, PyObject *source)
1440{
1441 chainobject *lz;
1442
1443 lz = (chainobject *)type->tp_alloc(type, 0);
1444 if (lz == NULL) {
1445 Py_DECREF(source);
1446 return NULL;
1447 }
1448
1449 lz->source = source;
1450 lz->active = NULL;
1451 return (PyObject *)lz;
1452}
1453
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001454static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001455chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001456{
Christian Heimesf16baeb2008-02-29 14:57:44 +00001457 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001458
Thomas Woutersb2137042007-02-01 18:02:27 +00001459 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001460 return NULL;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001461
1462 source = PyObject_GetIter(args);
1463 if (source == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001464 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001465
Christian Heimesf16baeb2008-02-29 14:57:44 +00001466 return chain_new_internal(type, source);
1467}
1468
1469static PyObject *
1470chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1471{
1472 PyObject *source;
1473
1474 source = PyObject_GetIter(arg);
1475 if (source == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001476 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001477
Christian Heimesf16baeb2008-02-29 14:57:44 +00001478 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001479}
1480
1481static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001482chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001483{
1484 PyObject_GC_UnTrack(lz);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001485 Py_XDECREF(lz->active);
1486 Py_XDECREF(lz->source);
Christian Heimes90aa7642007-12-19 02:45:37 +00001487 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001488}
1489
Raymond Hettinger2012f172003-02-07 05:32:58 +00001490static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001491chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001492{
Christian Heimesf16baeb2008-02-29 14:57:44 +00001493 Py_VISIT(lz->source);
1494 Py_VISIT(lz->active);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001495 return 0;
1496}
1497
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001498static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001499chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001500{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001501 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001502
Christian Heimesf16baeb2008-02-29 14:57:44 +00001503 if (lz->source == NULL)
1504 return NULL; /* already stopped */
1505
1506 if (lz->active == NULL) {
1507 PyObject *iterable = PyIter_Next(lz->source);
1508 if (iterable == NULL) {
1509 Py_CLEAR(lz->source);
1510 return NULL; /* no more input sources */
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001511 }
Christian Heimesf16baeb2008-02-29 14:57:44 +00001512 lz->active = PyObject_GetIter(iterable);
Christian Heimes78644762008-03-04 23:39:23 +00001513 Py_DECREF(iterable);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001514 if (lz->active == NULL) {
Christian Heimesf16baeb2008-02-29 14:57:44 +00001515 Py_CLEAR(lz->source);
1516 return NULL; /* input not iterable */
1517 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001518 }
Christian Heimesf16baeb2008-02-29 14:57:44 +00001519 item = PyIter_Next(lz->active);
1520 if (item != NULL)
1521 return item;
1522 if (PyErr_Occurred()) {
1523 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1524 PyErr_Clear();
1525 else
1526 return NULL; /* input raised an exception */
1527 }
1528 Py_CLEAR(lz->active);
1529 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001530}
1531
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001532PyDoc_STRVAR(chain_doc,
1533"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001534\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001535Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001536first iterable until it is exhausted, then elements from the next\n\
1537iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001538
Christian Heimesf16baeb2008-02-29 14:57:44 +00001539PyDoc_STRVAR(chain_from_iterable_doc,
1540"chain.from_iterable(iterable) --> chain object\n\
1541\n\
1542Alternate chain() contructor taking a single iterable argument\n\
1543that evaluates lazily.");
1544
1545static PyMethodDef chain_methods[] = {
1546 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1547 chain_from_iterable_doc},
1548 {NULL, NULL} /* sentinel */
1549};
1550
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001551static PyTypeObject chain_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001552 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001553 "itertools.chain", /* tp_name */
1554 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001555 0, /* tp_itemsize */
1556 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001557 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001558 0, /* tp_print */
1559 0, /* tp_getattr */
1560 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00001561 0, /* tp_reserved */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001562 0, /* tp_repr */
1563 0, /* tp_as_number */
1564 0, /* tp_as_sequence */
1565 0, /* tp_as_mapping */
1566 0, /* tp_hash */
1567 0, /* tp_call */
1568 0, /* tp_str */
1569 PyObject_GenericGetAttr, /* tp_getattro */
1570 0, /* tp_setattro */
1571 0, /* tp_as_buffer */
1572 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1573 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001574 chain_doc, /* tp_doc */
1575 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001576 0, /* tp_clear */
1577 0, /* tp_richcompare */
1578 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001579 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001580 (iternextfunc)chain_next, /* tp_iternext */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001581 chain_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001582 0, /* tp_members */
1583 0, /* tp_getset */
1584 0, /* tp_base */
1585 0, /* tp_dict */
1586 0, /* tp_descr_get */
1587 0, /* tp_descr_set */
1588 0, /* tp_dictoffset */
1589 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001590 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001591 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001592 PyObject_GC_Del, /* tp_free */
1593};
1594
1595
Christian Heimesc3f30c42008-02-22 16:37:40 +00001596/* product object ************************************************************/
1597
1598typedef struct {
1599 PyObject_HEAD
1600 PyObject *pools; /* tuple of pool tuples */
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001601 Py_ssize_t *indices; /* one index per pool */
1602 PyObject *result; /* most recently returned result tuple */
1603 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001604} productobject;
1605
1606static PyTypeObject product_type;
1607
1608static PyObject *
1609product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1610{
1611 productobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001612 Py_ssize_t nargs, npools, repeat=1;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001613 PyObject *pools = NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001614 Py_ssize_t *indices = NULL;
1615 Py_ssize_t i;
1616
Christian Heimesf16baeb2008-02-29 14:57:44 +00001617 if (kwds != NULL) {
1618 char *kwlist[] = {"repeat", 0};
1619 PyObject *tmpargs = PyTuple_New(0);
1620 if (tmpargs == NULL)
1621 return NULL;
1622 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1623 Py_DECREF(tmpargs);
1624 return NULL;
1625 }
1626 Py_DECREF(tmpargs);
1627 if (repeat < 0) {
1628 PyErr_SetString(PyExc_ValueError,
1629 "repeat argument cannot be negative");
1630 return NULL;
1631 }
1632 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001633
1634 assert(PyTuple_Check(args));
Christian Heimesf16baeb2008-02-29 14:57:44 +00001635 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1636 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001637
Christian Heimesc3f30c42008-02-22 16:37:40 +00001638 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
Christian Heimesb558a2e2008-03-02 22:46:37 +00001639 if (indices == NULL) {
Christian Heimesc3f30c42008-02-22 16:37:40 +00001640 PyErr_NoMemory();
1641 goto error;
1642 }
1643
1644 pools = PyTuple_New(npools);
1645 if (pools == NULL)
1646 goto error;
1647
Christian Heimesf16baeb2008-02-29 14:57:44 +00001648 for (i=0; i < nargs ; ++i) {
Christian Heimesc3f30c42008-02-22 16:37:40 +00001649 PyObject *item = PyTuple_GET_ITEM(args, i);
1650 PyObject *pool = PySequence_Tuple(item);
1651 if (pool == NULL)
1652 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001653 PyTuple_SET_ITEM(pools, i, pool);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001654 indices[i] = 0;
1655 }
Christian Heimesf16baeb2008-02-29 14:57:44 +00001656 for ( ; i < npools; ++i) {
1657 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1658 Py_INCREF(pool);
1659 PyTuple_SET_ITEM(pools, i, pool);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001660 indices[i] = 0;
1661 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001662
1663 /* create productobject structure */
1664 lz = (productobject *)type->tp_alloc(type, 0);
Christian Heimes380f7f22008-02-28 11:19:05 +00001665 if (lz == NULL)
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001666 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001667
1668 lz->pools = pools;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001669 lz->indices = indices;
1670 lz->result = NULL;
1671 lz->stopped = 0;
1672
1673 return (PyObject *)lz;
1674
1675error:
Christian Heimesc3f30c42008-02-22 16:37:40 +00001676 if (indices != NULL)
1677 PyMem_Free(indices);
1678 Py_XDECREF(pools);
1679 return NULL;
1680}
1681
1682static void
1683product_dealloc(productobject *lz)
1684{
1685 PyObject_GC_UnTrack(lz);
1686 Py_XDECREF(lz->pools);
1687 Py_XDECREF(lz->result);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001688 if (lz->indices != NULL)
1689 PyMem_Free(lz->indices);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001690 Py_TYPE(lz)->tp_free(lz);
1691}
1692
1693static int
1694product_traverse(productobject *lz, visitproc visit, void *arg)
1695{
1696 Py_VISIT(lz->pools);
1697 Py_VISIT(lz->result);
1698 return 0;
1699}
1700
1701static PyObject *
1702product_next(productobject *lz)
1703{
1704 PyObject *pool;
1705 PyObject *elem;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001706 PyObject *oldelem;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001707 PyObject *pools = lz->pools;
1708 PyObject *result = lz->result;
1709 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1710 Py_ssize_t i;
1711
1712 if (lz->stopped)
1713 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001714
Christian Heimesc3f30c42008-02-22 16:37:40 +00001715 if (result == NULL) {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001716 /* On the first pass, return an initial tuple filled with the
Christian Heimes78644762008-03-04 23:39:23 +00001717 first element from each pool. */
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001718 result = PyTuple_New(npools);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001719 if (result == NULL)
1720 goto empty;
1721 lz->result = result;
1722 for (i=0; i < npools; i++) {
1723 pool = PyTuple_GET_ITEM(pools, i);
1724 if (PyTuple_GET_SIZE(pool) == 0)
1725 goto empty;
1726 elem = PyTuple_GET_ITEM(pool, 0);
1727 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001728 PyTuple_SET_ITEM(result, i, elem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001729 }
1730 } else {
1731 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001732
1733 /* Copy the previous result tuple or re-use it if available */
1734 if (Py_REFCNT(result) > 1) {
1735 PyObject *old_result = result;
1736 result = PyTuple_New(npools);
1737 if (result == NULL)
1738 goto empty;
1739 lz->result = result;
1740 for (i=0; i < npools; i++) {
1741 elem = PyTuple_GET_ITEM(old_result, i);
1742 Py_INCREF(elem);
1743 PyTuple_SET_ITEM(result, i, elem);
1744 }
1745 Py_DECREF(old_result);
1746 }
1747 /* Now, we've got the only copy so we can update it in-place */
Christian Heimesb558a2e2008-03-02 22:46:37 +00001748 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001749
1750 /* Update the pool indices right-to-left. Only advance to the
1751 next pool when the previous one rolls-over */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001752 for (i=npools-1 ; i >= 0 ; i--) {
1753 pool = PyTuple_GET_ITEM(pools, i);
1754 indices[i]++;
Christian Heimesb558a2e2008-03-02 22:46:37 +00001755 if (indices[i] == PyTuple_GET_SIZE(pool)) {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001756 /* Roll-over and advance to next pool */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001757 indices[i] = 0;
1758 elem = PyTuple_GET_ITEM(pool, 0);
1759 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001760 oldelem = PyTuple_GET_ITEM(result, i);
1761 PyTuple_SET_ITEM(result, i, elem);
1762 Py_DECREF(oldelem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001763 } else {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001764 /* No rollover. Just increment and stop here. */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001765 elem = PyTuple_GET_ITEM(pool, indices[i]);
1766 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001767 oldelem = PyTuple_GET_ITEM(result, i);
1768 PyTuple_SET_ITEM(result, i, elem);
1769 Py_DECREF(oldelem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001770 break;
1771 }
1772 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001773
1774 /* If i is negative, then the indices have all rolled-over
1775 and we're done. */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001776 if (i < 0)
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001777 goto empty;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001778 }
1779
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001780 Py_INCREF(result);
1781 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001782
1783empty:
1784 lz->stopped = 1;
1785 return NULL;
1786}
1787
1788PyDoc_STRVAR(product_doc,
1789"product(*iterables) --> product object\n\
1790\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001791Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001792For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1793The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1794cycle in a manner similar to an odometer (with the rightmost element changing\n\
1795on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00001796To compute the product of an iterable with itself, specify the number\n\
1797of repetitions with the optional repeat keyword argument. For example,\n\
1798product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001799product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1800product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1801
1802static PyTypeObject product_type = {
1803 PyVarObject_HEAD_INIT(NULL, 0)
1804 "itertools.product", /* tp_name */
1805 sizeof(productobject), /* tp_basicsize */
1806 0, /* tp_itemsize */
1807 /* methods */
1808 (destructor)product_dealloc, /* tp_dealloc */
1809 0, /* tp_print */
1810 0, /* tp_getattr */
1811 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00001812 0, /* tp_reserved */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001813 0, /* tp_repr */
1814 0, /* tp_as_number */
1815 0, /* tp_as_sequence */
1816 0, /* tp_as_mapping */
1817 0, /* tp_hash */
1818 0, /* tp_call */
1819 0, /* tp_str */
1820 PyObject_GenericGetAttr, /* tp_getattro */
1821 0, /* tp_setattro */
1822 0, /* tp_as_buffer */
1823 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1824 Py_TPFLAGS_BASETYPE, /* tp_flags */
1825 product_doc, /* tp_doc */
1826 (traverseproc)product_traverse, /* tp_traverse */
1827 0, /* tp_clear */
1828 0, /* tp_richcompare */
1829 0, /* tp_weaklistoffset */
1830 PyObject_SelfIter, /* tp_iter */
1831 (iternextfunc)product_next, /* tp_iternext */
1832 0, /* tp_methods */
1833 0, /* tp_members */
1834 0, /* tp_getset */
1835 0, /* tp_base */
1836 0, /* tp_dict */
1837 0, /* tp_descr_get */
1838 0, /* tp_descr_set */
1839 0, /* tp_dictoffset */
1840 0, /* tp_init */
1841 0, /* tp_alloc */
1842 product_new, /* tp_new */
1843 PyObject_GC_Del, /* tp_free */
1844};
1845
1846
Christian Heimes380f7f22008-02-28 11:19:05 +00001847/* combinations object ************************************************************/
1848
1849typedef struct {
1850 PyObject_HEAD
1851 PyObject *pool; /* input converted to a tuple */
1852 Py_ssize_t *indices; /* one index per result element */
1853 PyObject *result; /* most recently returned result tuple */
1854 Py_ssize_t r; /* size of result tuple */
1855 int stopped; /* set to 1 when the combinations iterator is exhausted */
1856} combinationsobject;
1857
1858static PyTypeObject combinations_type;
1859
1860static PyObject *
1861combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1862{
1863 combinationsobject *co;
1864 Py_ssize_t n;
1865 Py_ssize_t r;
1866 PyObject *pool = NULL;
1867 PyObject *iterable = NULL;
1868 Py_ssize_t *indices = NULL;
1869 Py_ssize_t i;
1870 static char *kwargs[] = {"iterable", "r", NULL};
1871
1872 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
1873 &iterable, &r))
1874 return NULL;
1875
1876 pool = PySequence_Tuple(iterable);
1877 if (pool == NULL)
1878 goto error;
1879 n = PyTuple_GET_SIZE(pool);
1880 if (r < 0) {
1881 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
1882 goto error;
1883 }
Christian Heimes380f7f22008-02-28 11:19:05 +00001884
1885 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
1886 if (indices == NULL) {
1887 PyErr_NoMemory();
1888 goto error;
1889 }
1890
1891 for (i=0 ; i<r ; i++)
1892 indices[i] = i;
1893
1894 /* create combinationsobject structure */
1895 co = (combinationsobject *)type->tp_alloc(type, 0);
1896 if (co == NULL)
1897 goto error;
1898
1899 co->pool = pool;
1900 co->indices = indices;
1901 co->result = NULL;
1902 co->r = r;
Raymond Hettinger5bad41e2009-01-08 21:01:54 +00001903 co->stopped = r > n ? 1 : 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00001904
1905 return (PyObject *)co;
1906
1907error:
1908 if (indices != NULL)
1909 PyMem_Free(indices);
1910 Py_XDECREF(pool);
1911 return NULL;
1912}
1913
1914static void
1915combinations_dealloc(combinationsobject *co)
1916{
1917 PyObject_GC_UnTrack(co);
1918 Py_XDECREF(co->pool);
1919 Py_XDECREF(co->result);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001920 if (co->indices != NULL)
1921 PyMem_Free(co->indices);
Christian Heimes380f7f22008-02-28 11:19:05 +00001922 Py_TYPE(co)->tp_free(co);
1923}
1924
1925static int
1926combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
1927{
1928 Py_VISIT(co->pool);
1929 Py_VISIT(co->result);
1930 return 0;
1931}
1932
1933static PyObject *
1934combinations_next(combinationsobject *co)
1935{
1936 PyObject *elem;
1937 PyObject *oldelem;
1938 PyObject *pool = co->pool;
1939 Py_ssize_t *indices = co->indices;
1940 PyObject *result = co->result;
1941 Py_ssize_t n = PyTuple_GET_SIZE(pool);
1942 Py_ssize_t r = co->r;
1943 Py_ssize_t i, j, index;
1944
1945 if (co->stopped)
1946 return NULL;
1947
1948 if (result == NULL) {
1949 /* On the first pass, initialize result tuple using the indices */
1950 result = PyTuple_New(r);
1951 if (result == NULL)
1952 goto empty;
1953 co->result = result;
1954 for (i=0; i<r ; i++) {
1955 index = indices[i];
1956 elem = PyTuple_GET_ITEM(pool, index);
1957 Py_INCREF(elem);
1958 PyTuple_SET_ITEM(result, i, elem);
1959 }
1960 } else {
1961 /* Copy the previous result tuple or re-use it if available */
1962 if (Py_REFCNT(result) > 1) {
1963 PyObject *old_result = result;
1964 result = PyTuple_New(r);
1965 if (result == NULL)
1966 goto empty;
1967 co->result = result;
1968 for (i=0; i<r ; i++) {
1969 elem = PyTuple_GET_ITEM(old_result, i);
1970 Py_INCREF(elem);
1971 PyTuple_SET_ITEM(result, i, elem);
1972 }
1973 Py_DECREF(old_result);
1974 }
1975 /* Now, we've got the only copy so we can update it in-place
1976 * CPython's empty tuple is a singleton and cached in
1977 * PyTuple's freelist.
1978 */
1979 assert(r == 0 || Py_REFCNT(result) == 1);
1980
1981 /* Scan indices right-to-left until finding one that is not
1982 at its maximum (i + n - r). */
1983 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
1984 ;
1985
1986 /* If i is negative, then the indices are all at
1987 their maximum value and we're done. */
1988 if (i < 0)
1989 goto empty;
1990
1991 /* Increment the current index which we know is not at its
1992 maximum. Then move back to the right setting each index
1993 to its lowest possible value (one higher than the index
1994 to its left -- this maintains the sort order invariant). */
1995 indices[i]++;
1996 for (j=i+1 ; j<r ; j++)
1997 indices[j] = indices[j-1] + 1;
1998
1999 /* Update the result tuple for the new indices
2000 starting with i, the leftmost index that changed */
2001 for ( ; i<r ; i++) {
2002 index = indices[i];
2003 elem = PyTuple_GET_ITEM(pool, index);
2004 Py_INCREF(elem);
2005 oldelem = PyTuple_GET_ITEM(result, i);
2006 PyTuple_SET_ITEM(result, i, elem);
2007 Py_DECREF(oldelem);
2008 }
2009 }
2010
2011 Py_INCREF(result);
2012 return result;
2013
2014empty:
2015 co->stopped = 1;
2016 return NULL;
2017}
2018
2019PyDoc_STRVAR(combinations_doc,
Raymond Hettingerd3a77c02009-11-19 01:23:41 +00002020"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002021\n\
2022Return successive r-length combinations of elements in the iterable.\n\n\
2023combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2024
2025static PyTypeObject combinations_type = {
2026 PyVarObject_HEAD_INIT(NULL, 0)
2027 "itertools.combinations", /* tp_name */
2028 sizeof(combinationsobject), /* tp_basicsize */
2029 0, /* tp_itemsize */
2030 /* methods */
2031 (destructor)combinations_dealloc, /* tp_dealloc */
2032 0, /* tp_print */
2033 0, /* tp_getattr */
2034 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00002035 0, /* tp_reserved */
Christian Heimes380f7f22008-02-28 11:19:05 +00002036 0, /* tp_repr */
2037 0, /* tp_as_number */
2038 0, /* tp_as_sequence */
2039 0, /* tp_as_mapping */
2040 0, /* tp_hash */
2041 0, /* tp_call */
2042 0, /* tp_str */
2043 PyObject_GenericGetAttr, /* tp_getattro */
2044 0, /* tp_setattro */
2045 0, /* tp_as_buffer */
2046 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2047 Py_TPFLAGS_BASETYPE, /* tp_flags */
2048 combinations_doc, /* tp_doc */
2049 (traverseproc)combinations_traverse, /* tp_traverse */
2050 0, /* tp_clear */
2051 0, /* tp_richcompare */
2052 0, /* tp_weaklistoffset */
2053 PyObject_SelfIter, /* tp_iter */
2054 (iternextfunc)combinations_next, /* tp_iternext */
2055 0, /* tp_methods */
2056 0, /* tp_members */
2057 0, /* tp_getset */
2058 0, /* tp_base */
2059 0, /* tp_dict */
2060 0, /* tp_descr_get */
2061 0, /* tp_descr_set */
2062 0, /* tp_dictoffset */
2063 0, /* tp_init */
2064 0, /* tp_alloc */
2065 combinations_new, /* tp_new */
2066 PyObject_GC_Del, /* tp_free */
2067};
2068
2069
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002070/* combinations with replacement object *******************************************/
2071
2072/* Equivalent to:
2073
2074 def combinations_with_replacement(iterable, r):
2075 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2076 # number items returned: (n+r-1)! / r! / (n-1)!
2077 pool = tuple(iterable)
2078 n = len(pool)
2079 indices = [0] * r
2080 yield tuple(pool[i] for i in indices)
2081 while 1:
2082 for i in reversed(range(r)):
2083 if indices[i] != n - 1:
2084 break
2085 else:
2086 return
2087 indices[i:] = [indices[i] + 1] * (r - i)
2088 yield tuple(pool[i] for i in indices)
2089
2090 def combinations_with_replacement2(iterable, r):
2091 'Alternate version that filters from product()'
2092 pool = tuple(iterable)
2093 n = len(pool)
2094 for indices in product(range(n), repeat=r):
2095 if sorted(indices) == list(indices):
2096 yield tuple(pool[i] for i in indices)
2097*/
2098typedef struct {
2099 PyObject_HEAD
2100 PyObject *pool; /* input converted to a tuple */
2101 Py_ssize_t *indices; /* one index per result element */
2102 PyObject *result; /* most recently returned result tuple */
2103 Py_ssize_t r; /* size of result tuple */
2104 int stopped; /* set to 1 when the cwr iterator is exhausted */
2105} cwrobject;
2106
2107static PyTypeObject cwr_type;
2108
2109static PyObject *
2110cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2111{
2112 cwrobject *co;
2113 Py_ssize_t n;
2114 Py_ssize_t r;
2115 PyObject *pool = NULL;
2116 PyObject *iterable = NULL;
2117 Py_ssize_t *indices = NULL;
2118 Py_ssize_t i;
2119 static char *kwargs[] = {"iterable", "r", NULL};
2120
2121 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2122 &iterable, &r))
2123 return NULL;
2124
2125 pool = PySequence_Tuple(iterable);
2126 if (pool == NULL)
2127 goto error;
2128 n = PyTuple_GET_SIZE(pool);
2129 if (r < 0) {
2130 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2131 goto error;
2132 }
2133
2134 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2135 if (indices == NULL) {
2136 PyErr_NoMemory();
2137 goto error;
2138 }
2139
2140 for (i=0 ; i<r ; i++)
2141 indices[i] = 0;
2142
2143 /* create cwrobject structure */
2144 co = (cwrobject *)type->tp_alloc(type, 0);
2145 if (co == NULL)
2146 goto error;
2147
2148 co->pool = pool;
2149 co->indices = indices;
2150 co->result = NULL;
2151 co->r = r;
2152 co->stopped = !n && r;
2153
2154 return (PyObject *)co;
2155
2156error:
2157 if (indices != NULL)
2158 PyMem_Free(indices);
2159 Py_XDECREF(pool);
2160 return NULL;
2161}
2162
2163static void
2164cwr_dealloc(cwrobject *co)
2165{
2166 PyObject_GC_UnTrack(co);
2167 Py_XDECREF(co->pool);
2168 Py_XDECREF(co->result);
2169 if (co->indices != NULL)
2170 PyMem_Free(co->indices);
2171 Py_TYPE(co)->tp_free(co);
2172}
2173
2174static int
2175cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2176{
2177 Py_VISIT(co->pool);
2178 Py_VISIT(co->result);
2179 return 0;
2180}
2181
2182static PyObject *
2183cwr_next(cwrobject *co)
2184{
2185 PyObject *elem;
2186 PyObject *oldelem;
2187 PyObject *pool = co->pool;
2188 Py_ssize_t *indices = co->indices;
2189 PyObject *result = co->result;
2190 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2191 Py_ssize_t r = co->r;
2192 Py_ssize_t i, j, index;
2193
2194 if (co->stopped)
2195 return NULL;
2196
2197 if (result == NULL) {
2198 /* On the first pass, initialize result tuple using the indices */
2199 result = PyTuple_New(r);
2200 if (result == NULL)
2201 goto empty;
2202 co->result = result;
2203 for (i=0; i<r ; i++) {
2204 index = indices[i];
2205 elem = PyTuple_GET_ITEM(pool, index);
2206 Py_INCREF(elem);
2207 PyTuple_SET_ITEM(result, i, elem);
2208 }
2209 } else {
2210 /* Copy the previous result tuple or re-use it if available */
2211 if (Py_REFCNT(result) > 1) {
2212 PyObject *old_result = result;
2213 result = PyTuple_New(r);
2214 if (result == NULL)
2215 goto empty;
2216 co->result = result;
2217 for (i=0; i<r ; i++) {
2218 elem = PyTuple_GET_ITEM(old_result, i);
2219 Py_INCREF(elem);
2220 PyTuple_SET_ITEM(result, i, elem);
2221 }
2222 Py_DECREF(old_result);
2223 }
2224 /* Now, we've got the only copy so we can update it in-place CPython's
2225 empty tuple is a singleton and cached in PyTuple's freelist. */
2226 assert(r == 0 || Py_REFCNT(result) == 1);
2227
2228 /* Scan indices right-to-left until finding one that is not
2229 * at its maximum (n-1). */
2230 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2231 ;
2232
2233 /* If i is negative, then the indices are all at
2234 their maximum value and we're done. */
2235 if (i < 0)
2236 goto empty;
2237
2238 /* Increment the current index which we know is not at its
2239 maximum. Then set all to the right to the same value. */
2240 indices[i]++;
2241 for (j=i+1 ; j<r ; j++)
2242 indices[j] = indices[j-1];
2243
2244 /* Update the result tuple for the new indices
2245 starting with i, the leftmost index that changed */
2246 for ( ; i<r ; i++) {
2247 index = indices[i];
2248 elem = PyTuple_GET_ITEM(pool, index);
2249 Py_INCREF(elem);
2250 oldelem = PyTuple_GET_ITEM(result, i);
2251 PyTuple_SET_ITEM(result, i, elem);
2252 Py_DECREF(oldelem);
2253 }
2254 }
2255
2256 Py_INCREF(result);
2257 return result;
2258
2259empty:
2260 co->stopped = 1;
2261 return NULL;
2262}
2263
2264PyDoc_STRVAR(cwr_doc,
Raymond Hettingerd3a77c02009-11-19 01:23:41 +00002265"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002266\n\
2267Return successive r-length combinations of elements in the iterable\n\
2268allowing individual elements to have successive repeats.\n\
2269combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2270
2271static PyTypeObject cwr_type = {
2272 PyVarObject_HEAD_INIT(NULL, 0)
2273 "itertools.combinations_with_replacement", /* tp_name */
2274 sizeof(cwrobject), /* tp_basicsize */
2275 0, /* tp_itemsize */
2276 /* methods */
2277 (destructor)cwr_dealloc, /* tp_dealloc */
2278 0, /* tp_print */
2279 0, /* tp_getattr */
2280 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00002281 0, /* tp_reserved */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002282 0, /* tp_repr */
2283 0, /* tp_as_number */
2284 0, /* tp_as_sequence */
2285 0, /* tp_as_mapping */
2286 0, /* tp_hash */
2287 0, /* tp_call */
2288 0, /* tp_str */
2289 PyObject_GenericGetAttr, /* tp_getattro */
2290 0, /* tp_setattro */
2291 0, /* tp_as_buffer */
2292 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2293 Py_TPFLAGS_BASETYPE, /* tp_flags */
2294 cwr_doc, /* tp_doc */
2295 (traverseproc)cwr_traverse, /* tp_traverse */
2296 0, /* tp_clear */
2297 0, /* tp_richcompare */
2298 0, /* tp_weaklistoffset */
2299 PyObject_SelfIter, /* tp_iter */
2300 (iternextfunc)cwr_next, /* tp_iternext */
2301 0, /* tp_methods */
2302 0, /* tp_members */
2303 0, /* tp_getset */
2304 0, /* tp_base */
2305 0, /* tp_dict */
2306 0, /* tp_descr_get */
2307 0, /* tp_descr_set */
2308 0, /* tp_dictoffset */
2309 0, /* tp_init */
2310 0, /* tp_alloc */
2311 cwr_new, /* tp_new */
2312 PyObject_GC_Del, /* tp_free */
2313};
2314
2315
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002316/* permutations object ************************************************************
2317
2318def permutations(iterable, r=None):
2319 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2320 pool = tuple(iterable)
2321 n = len(pool)
2322 r = n if r is None else r
2323 indices = range(n)
2324 cycles = range(n-r+1, n+1)[::-1]
2325 yield tuple(pool[i] for i in indices[:r])
2326 while n:
2327 for i in reversed(range(r)):
2328 cycles[i] -= 1
2329 if cycles[i] == 0:
2330 indices[i:] = indices[i+1:] + indices[i:i+1]
2331 cycles[i] = n - i
2332 else:
2333 j = cycles[i]
2334 indices[i], indices[-j] = indices[-j], indices[i]
2335 yield tuple(pool[i] for i in indices[:r])
2336 break
2337 else:
2338 return
2339*/
2340
2341typedef struct {
2342 PyObject_HEAD
2343 PyObject *pool; /* input converted to a tuple */
2344 Py_ssize_t *indices; /* one index per element in the pool */
2345 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2346 PyObject *result; /* most recently returned result tuple */
2347 Py_ssize_t r; /* size of result tuple */
2348 int stopped; /* set to 1 when the permutations iterator is exhausted */
2349} permutationsobject;
2350
2351static PyTypeObject permutations_type;
2352
2353static PyObject *
2354permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2355{
2356 permutationsobject *po;
2357 Py_ssize_t n;
2358 Py_ssize_t r;
2359 PyObject *robj = Py_None;
2360 PyObject *pool = NULL;
2361 PyObject *iterable = NULL;
2362 Py_ssize_t *indices = NULL;
2363 Py_ssize_t *cycles = NULL;
2364 Py_ssize_t i;
2365 static char *kwargs[] = {"iterable", "r", NULL};
2366
2367 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2368 &iterable, &robj))
2369 return NULL;
2370
2371 pool = PySequence_Tuple(iterable);
2372 if (pool == NULL)
2373 goto error;
2374 n = PyTuple_GET_SIZE(pool);
2375
2376 r = n;
2377 if (robj != Py_None) {
2378 if (!PyLong_Check(robj)) {
2379 PyErr_SetString(PyExc_TypeError, "Expected int as r");
Neal Norwitzd2bee322008-04-01 07:37:58 +00002380 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002381 }
2382 r = PyLong_AsSsize_t(robj);
2383 if (r == -1 && PyErr_Occurred())
2384 goto error;
2385 }
2386 if (r < 0) {
2387 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2388 goto error;
2389 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002390
2391 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2392 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2393 if (indices == NULL || cycles == NULL) {
2394 PyErr_NoMemory();
2395 goto error;
2396 }
2397
2398 for (i=0 ; i<n ; i++)
2399 indices[i] = i;
2400 for (i=0 ; i<r ; i++)
2401 cycles[i] = n - i;
2402
2403 /* create permutationsobject structure */
2404 po = (permutationsobject *)type->tp_alloc(type, 0);
2405 if (po == NULL)
2406 goto error;
2407
2408 po->pool = pool;
2409 po->indices = indices;
2410 po->cycles = cycles;
2411 po->result = NULL;
2412 po->r = r;
Raymond Hettinger5bad41e2009-01-08 21:01:54 +00002413 po->stopped = r > n ? 1 : 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002414
2415 return (PyObject *)po;
2416
2417error:
2418 if (indices != NULL)
2419 PyMem_Free(indices);
2420 if (cycles != NULL)
2421 PyMem_Free(cycles);
2422 Py_XDECREF(pool);
2423 return NULL;
2424}
2425
2426static void
2427permutations_dealloc(permutationsobject *po)
2428{
2429 PyObject_GC_UnTrack(po);
2430 Py_XDECREF(po->pool);
2431 Py_XDECREF(po->result);
2432 PyMem_Free(po->indices);
2433 PyMem_Free(po->cycles);
2434 Py_TYPE(po)->tp_free(po);
2435}
2436
2437static int
2438permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2439{
2440 Py_VISIT(po->pool);
2441 Py_VISIT(po->result);
2442 return 0;
2443}
2444
2445static PyObject *
2446permutations_next(permutationsobject *po)
2447{
2448 PyObject *elem;
2449 PyObject *oldelem;
2450 PyObject *pool = po->pool;
2451 Py_ssize_t *indices = po->indices;
2452 Py_ssize_t *cycles = po->cycles;
2453 PyObject *result = po->result;
2454 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2455 Py_ssize_t r = po->r;
2456 Py_ssize_t i, j, k, index;
2457
2458 if (po->stopped)
2459 return NULL;
2460
2461 if (result == NULL) {
2462 /* On the first pass, initialize result tuple using the indices */
2463 result = PyTuple_New(r);
2464 if (result == NULL)
2465 goto empty;
2466 po->result = result;
2467 for (i=0; i<r ; i++) {
2468 index = indices[i];
2469 elem = PyTuple_GET_ITEM(pool, index);
2470 Py_INCREF(elem);
2471 PyTuple_SET_ITEM(result, i, elem);
2472 }
2473 } else {
2474 if (n == 0)
2475 goto empty;
2476
2477 /* Copy the previous result tuple or re-use it if available */
2478 if (Py_REFCNT(result) > 1) {
2479 PyObject *old_result = result;
2480 result = PyTuple_New(r);
2481 if (result == NULL)
2482 goto empty;
2483 po->result = result;
2484 for (i=0; i<r ; i++) {
2485 elem = PyTuple_GET_ITEM(old_result, i);
2486 Py_INCREF(elem);
2487 PyTuple_SET_ITEM(result, i, elem);
2488 }
2489 Py_DECREF(old_result);
2490 }
2491 /* Now, we've got the only copy so we can update it in-place */
2492 assert(r == 0 || Py_REFCNT(result) == 1);
2493
2494 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2495 for (i=r-1 ; i>=0 ; i--) {
2496 cycles[i] -= 1;
2497 if (cycles[i] == 0) {
2498 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2499 index = indices[i];
2500 for (j=i ; j<n-1 ; j++)
2501 indices[j] = indices[j+1];
2502 indices[n-1] = index;
2503 cycles[i] = n - i;
2504 } else {
2505 j = cycles[i];
2506 index = indices[i];
2507 indices[i] = indices[n-j];
2508 indices[n-j] = index;
2509
2510 for (k=i; k<r ; k++) {
2511 /* start with i, the leftmost element that changed */
2512 /* yield tuple(pool[k] for k in indices[:r]) */
2513 index = indices[k];
2514 elem = PyTuple_GET_ITEM(pool, index);
2515 Py_INCREF(elem);
2516 oldelem = PyTuple_GET_ITEM(result, k);
2517 PyTuple_SET_ITEM(result, k, elem);
2518 Py_DECREF(oldelem);
2519 }
2520 break;
2521 }
2522 }
2523 /* If i is negative, then the cycles have all
2524 rolled-over and we're done. */
2525 if (i < 0)
2526 goto empty;
2527 }
2528 Py_INCREF(result);
2529 return result;
2530
2531empty:
2532 po->stopped = 1;
2533 return NULL;
2534}
2535
2536PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00002537"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002538\n\
2539Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00002540permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002541
2542static PyTypeObject permutations_type = {
2543 PyVarObject_HEAD_INIT(NULL, 0)
2544 "itertools.permutations", /* tp_name */
2545 sizeof(permutationsobject), /* tp_basicsize */
2546 0, /* tp_itemsize */
2547 /* methods */
2548 (destructor)permutations_dealloc, /* tp_dealloc */
2549 0, /* tp_print */
2550 0, /* tp_getattr */
2551 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00002552 0, /* tp_reserved */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002553 0, /* tp_repr */
2554 0, /* tp_as_number */
2555 0, /* tp_as_sequence */
2556 0, /* tp_as_mapping */
2557 0, /* tp_hash */
2558 0, /* tp_call */
2559 0, /* tp_str */
2560 PyObject_GenericGetAttr, /* tp_getattro */
2561 0, /* tp_setattro */
2562 0, /* tp_as_buffer */
2563 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2564 Py_TPFLAGS_BASETYPE, /* tp_flags */
2565 permutations_doc, /* tp_doc */
2566 (traverseproc)permutations_traverse, /* tp_traverse */
2567 0, /* tp_clear */
2568 0, /* tp_richcompare */
2569 0, /* tp_weaklistoffset */
2570 PyObject_SelfIter, /* tp_iter */
2571 (iternextfunc)permutations_next, /* tp_iternext */
2572 0, /* tp_methods */
2573 0, /* tp_members */
2574 0, /* tp_getset */
2575 0, /* tp_base */
2576 0, /* tp_dict */
2577 0, /* tp_descr_get */
2578 0, /* tp_descr_set */
2579 0, /* tp_dictoffset */
2580 0, /* tp_init */
2581 0, /* tp_alloc */
2582 permutations_new, /* tp_new */
2583 PyObject_GC_Del, /* tp_free */
2584};
2585
2586
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002587/* compress object ************************************************************/
2588
2589/* Equivalent to:
2590
2591 def compress(data, selectors):
2592 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2593 return (d for d, s in zip(data, selectors) if s)
2594*/
2595
2596typedef struct {
2597 PyObject_HEAD
2598 PyObject *data;
2599 PyObject *selectors;
2600} compressobject;
2601
2602static PyTypeObject compress_type;
2603
2604static PyObject *
2605compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2606{
2607 PyObject *seq1, *seq2;
2608 PyObject *data=NULL, *selectors=NULL;
2609 compressobject *lz;
Raymond Hettinger15a49502009-02-19 02:17:09 +00002610 static char *kwargs[] = {"data", "selectors", NULL};
2611
2612 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002613 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,
Raymond Hettinger15a49502009-02-19 02:17:09 +00002690"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002691\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
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002895fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00002896
2897 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
2898 Advances with: cnt += 1
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002899 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00002900
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002901slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00002902
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).
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002907 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00002908*/
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;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002916 int slow_mode = 0;
Thomas Wouters477c8d52006-05-27 19:21:47 +00002917 Py_ssize_t cnt = 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002918 PyObject *long_cnt = NULL;
Raymond Hettinger30729212009-02-12 06:28:27 +00002919 PyObject *long_step = NULL;
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +00002920 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002921
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +00002922 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
2923 kwlist, &long_cnt, &long_step))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002924 return NULL;
2925
Benjamin Peterson912fbca2009-02-21 23:14:55 +00002926 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
2927 (long_step != NULL && !PyNumber_Check(long_step))) {
Raymond Hettinger30729212009-02-12 06:28:27 +00002928 PyErr_SetString(PyExc_TypeError, "a number is required");
2929 return NULL;
2930 }
2931
Raymond Hettinger30729212009-02-12 06:28:27 +00002932 if (long_cnt != NULL) {
2933 cnt = PyLong_AsSsize_t(long_cnt);
Benjamin Peterson912fbca2009-02-21 23:14:55 +00002934 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002935 PyErr_Clear();
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002936 slow_mode = 1;
2937 }
2938 Py_INCREF(long_cnt);
2939 } else {
2940 cnt = 0;
2941 long_cnt = PyLong_FromLong(0);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002942 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002943
2944 /* If not specified, step defaults to 1 */
2945 if (long_step == NULL) {
2946 long_step = PyLong_FromLong(1);
2947 if (long_step == NULL) {
2948 Py_DECREF(long_cnt);
2949 return NULL;
2950 }
2951 } else
2952 Py_INCREF(long_step);
2953
2954 assert(long_cnt != NULL && long_step != NULL);
2955
2956 /* Fast mode only works when the step is 1 */
2957 if (!PyLong_Check(long_step) ||
2958 PyLong_AS_LONG(long_step) != 1) {
2959 slow_mode = 1;
2960 }
2961
2962 if (slow_mode)
2963 cnt = PY_SSIZE_T_MAX;
2964 else
2965 Py_CLEAR(long_cnt);
2966
Benjamin Peterson912fbca2009-02-21 23:14:55 +00002967 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
2968 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002969 assert(slow_mode ||
Benjamin Peterson912fbca2009-02-21 23:14:55 +00002970 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002971
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002972 /* create countobject structure */
Raymond Hettingerd6280f42009-02-16 20:50:56 +00002973 lz = (countobject *)type->tp_alloc(type, 0);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002974 if (lz == NULL) {
2975 Py_XDECREF(long_cnt);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002976 return NULL;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002977 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002978 lz->cnt = cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002979 lz->long_cnt = long_cnt;
Raymond Hettinger30729212009-02-12 06:28:27 +00002980 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002981
2982 return (PyObject *)lz;
2983}
2984
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002985static void
2986count_dealloc(countobject *lz)
2987{
Raymond Hettingerd6280f42009-02-16 20:50:56 +00002988 PyObject_GC_UnTrack(lz);
Raymond Hettinger30729212009-02-12 06:28:27 +00002989 Py_XDECREF(lz->long_cnt);
2990 Py_XDECREF(lz->long_step);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00002991 Py_TYPE(lz)->tp_free(lz);
2992}
2993
Benjamin Peterson25e26a02009-02-16 21:28:29 +00002994static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00002995count_traverse(countobject *lz, visitproc visit, void *arg)
2996{
2997 Py_VISIT(lz->long_cnt);
2998 Py_VISIT(lz->long_step);
2999 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003000}
3001
3002static PyObject *
3003count_nextlong(countobject *lz)
3004{
Raymond Hettinger30729212009-02-12 06:28:27 +00003005 PyObject *long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003006 PyObject *stepped_up;
3007
Raymond Hettinger30729212009-02-12 06:28:27 +00003008 long_cnt = lz->long_cnt;
3009 if (long_cnt == NULL) {
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003010 /* Switch to slow_mode */
Raymond Hettinger30729212009-02-12 06:28:27 +00003011 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3012 if (long_cnt == NULL)
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003013 return NULL;
3014 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003015 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
3016
3017 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003018 if (stepped_up == NULL)
3019 return NULL;
3020 lz->long_cnt = stepped_up;
Raymond Hettinger30729212009-02-12 06:28:27 +00003021 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003022}
3023
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003024static PyObject *
3025count_next(countobject *lz)
3026{
Raymond Hettinger30729212009-02-12 06:28:27 +00003027 if (lz->cnt == PY_SSIZE_T_MAX)
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003028 return count_nextlong(lz);
Christian Heimes217cfd12007-12-02 14:31:20 +00003029 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003030}
3031
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003032static PyObject *
3033count_repr(countobject *lz)
3034{
Raymond Hettinger30729212009-02-12 06:28:27 +00003035 if (lz->cnt != PY_SSIZE_T_MAX)
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003036 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
3037
Raymond Hettinger30729212009-02-12 06:28:27 +00003038 if (PyLong_Check(lz->long_step)) {
3039 long step = PyLong_AsLong(lz->long_step);
3040 if (step == -1 && PyErr_Occurred()) {
3041 PyErr_Clear();
3042 }
3043 if (step == 1) {
3044 /* Don't display step when it is an integer equal to 1 */
3045 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
3046 }
3047 }
3048 return PyUnicode_FromFormat("count(%R, %R)",
3049 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003050}
3051
Raymond Hettinger6c8ee7a2009-11-30 21:55:17 +00003052static PyObject *
3053count_reduce(countobject *lz)
3054{
3055 if (lz->cnt == PY_SSIZE_T_MAX)
3056 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3057 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
3058}
3059
3060PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3061
3062static PyMethodDef count_methods[] = {
3063 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3064 count_reduce_doc},
3065 {NULL, NULL} /* sentinel */
3066};
3067
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003068PyDoc_STRVAR(count_doc,
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003069 "count(start=0, step=1]) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003070\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003071Return a count object whose .__next__() method returns consecutive values.\n\
3072Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00003073 def count(firstval=0, step=1):\n\
3074 x = firstval\n\
Raymond Hettinger6dfff192009-02-12 10:19:59 +00003075 while 1:\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00003076 yield x\n\
Raymond Hettinger6dfff192009-02-12 10:19:59 +00003077 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003078
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003079static PyTypeObject count_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00003080 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003081 "itertools.count", /* tp_name */
3082 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003083 0, /* tp_itemsize */
3084 /* methods */
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003085 (destructor)count_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003086 0, /* tp_print */
3087 0, /* tp_getattr */
3088 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00003089 0, /* tp_reserved */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003090 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003091 0, /* tp_as_number */
3092 0, /* tp_as_sequence */
3093 0, /* tp_as_mapping */
3094 0, /* tp_hash */
3095 0, /* tp_call */
3096 0, /* tp_str */
3097 PyObject_GenericGetAttr, /* tp_getattro */
3098 0, /* tp_setattro */
3099 0, /* tp_as_buffer */
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003100 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3101 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003102 count_doc, /* tp_doc */
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003103 (traverseproc)count_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003104 0, /* tp_clear */
3105 0, /* tp_richcompare */
3106 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003107 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003108 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger6c8ee7a2009-11-30 21:55:17 +00003109 count_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003110 0, /* tp_members */
3111 0, /* tp_getset */
3112 0, /* tp_base */
3113 0, /* tp_dict */
3114 0, /* tp_descr_get */
3115 0, /* tp_descr_set */
3116 0, /* tp_dictoffset */
3117 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003118 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003119 count_new, /* tp_new */
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003120 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003121};
3122
3123
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003124/* repeat object ************************************************************/
3125
3126typedef struct {
3127 PyObject_HEAD
3128 PyObject *element;
Thomas Wouters477c8d52006-05-27 19:21:47 +00003129 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003130} repeatobject;
3131
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003132static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003133
3134static PyObject *
3135repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3136{
3137 repeatobject *ro;
3138 PyObject *element;
Thomas Wouters477c8d52006-05-27 19:21:47 +00003139 Py_ssize_t cnt = -1;
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00003140 static char *kwargs[] = {"object", "times", NULL};
3141
3142 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3143 &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003144 return NULL;
3145
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003146 if (PyTuple_Size(args) == 2 && cnt < 0)
3147 cnt = 0;
3148
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003149 ro = (repeatobject *)type->tp_alloc(type, 0);
3150 if (ro == NULL)
3151 return NULL;
3152 Py_INCREF(element);
3153 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003154 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003155 return (PyObject *)ro;
3156}
3157
3158static void
3159repeat_dealloc(repeatobject *ro)
3160{
3161 PyObject_GC_UnTrack(ro);
3162 Py_XDECREF(ro->element);
Christian Heimes90aa7642007-12-19 02:45:37 +00003163 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003164}
3165
3166static int
3167repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3168{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00003169 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003170 return 0;
3171}
3172
3173static PyObject *
3174repeat_next(repeatobject *ro)
3175{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003176 if (ro->cnt == 0)
3177 return NULL;
3178 if (ro->cnt > 0)
3179 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003180 Py_INCREF(ro->element);
3181 return ro->element;
3182}
3183
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003184static PyObject *
3185repeat_repr(repeatobject *ro)
3186{
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003187 if (ro->cnt == -1)
Walter Dörwald7569dfe2007-05-19 21:49:49 +00003188 return PyUnicode_FromFormat("repeat(%R)", ro->element);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003189 else
Walter Dörwald7569dfe2007-05-19 21:49:49 +00003190 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003191}
3192
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003193static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003194repeat_len(repeatobject *ro)
3195{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003196 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003197 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003198 return NULL;
3199 }
Christian Heimes217cfd12007-12-02 14:31:20 +00003200 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003201}
3202
Armin Rigof5b3e362006-02-11 21:32:43 +00003203PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003204
3205static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00003206 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003207 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003208};
3209
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003210PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00003211"repeat(object [,times]) -> create an iterator which returns the object\n\
3212for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003213endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003214
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003215static PyTypeObject repeat_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00003216 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003217 "itertools.repeat", /* tp_name */
3218 sizeof(repeatobject), /* tp_basicsize */
3219 0, /* tp_itemsize */
3220 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003221 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003222 0, /* tp_print */
3223 0, /* tp_getattr */
3224 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00003225 0, /* tp_reserved */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003226 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003227 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003228 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003229 0, /* tp_as_mapping */
3230 0, /* tp_hash */
3231 0, /* tp_call */
3232 0, /* tp_str */
3233 PyObject_GenericGetAttr, /* tp_getattro */
3234 0, /* tp_setattro */
3235 0, /* tp_as_buffer */
3236 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3237 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003238 repeat_doc, /* tp_doc */
3239 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003240 0, /* tp_clear */
3241 0, /* tp_richcompare */
3242 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003243 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003244 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003245 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003246 0, /* tp_members */
3247 0, /* tp_getset */
3248 0, /* tp_base */
3249 0, /* tp_dict */
3250 0, /* tp_descr_get */
3251 0, /* tp_descr_set */
3252 0, /* tp_dictoffset */
3253 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003254 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003255 repeat_new, /* tp_new */
3256 PyObject_GC_Del, /* tp_free */
3257};
3258
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003259/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00003260
3261#include "Python.h"
3262
3263typedef struct {
3264 PyObject_HEAD
3265 Py_ssize_t tuplesize;
3266 Py_ssize_t numactive;
3267 PyObject *ittuple; /* tuple of iterators */
3268 PyObject *result;
3269 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003270} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003271
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003272static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003273
3274static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003275zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003276{
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003277 ziplongestobject *lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003278 Py_ssize_t i;
3279 PyObject *ittuple; /* tuple of iterators */
3280 PyObject *result;
3281 PyObject *fillvalue = Py_None;
3282 Py_ssize_t tuplesize = PySequence_Length(args);
3283
3284 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3285 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3286 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3287 PyErr_SetString(PyExc_TypeError,
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003288 "zip_longest() got an unexpected keyword argument");
Thomas Wouterscf297e42007-02-23 15:07:44 +00003289 return NULL;
3290 }
3291 }
3292
3293 /* args must be a tuple */
3294 assert(PyTuple_Check(args));
3295
3296 /* obtain iterators */
3297 ittuple = PyTuple_New(tuplesize);
3298 if (ittuple == NULL)
3299 return NULL;
3300 for (i=0; i < tuplesize; ++i) {
3301 PyObject *item = PyTuple_GET_ITEM(args, i);
3302 PyObject *it = PyObject_GetIter(item);
3303 if (it == NULL) {
3304 if (PyErr_ExceptionMatches(PyExc_TypeError))
3305 PyErr_Format(PyExc_TypeError,
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003306 "zip_longest argument #%zd must support iteration",
Thomas Wouterscf297e42007-02-23 15:07:44 +00003307 i+1);
3308 Py_DECREF(ittuple);
3309 return NULL;
3310 }
3311 PyTuple_SET_ITEM(ittuple, i, it);
3312 }
3313
3314 /* create a result holder */
3315 result = PyTuple_New(tuplesize);
3316 if (result == NULL) {
3317 Py_DECREF(ittuple);
3318 return NULL;
3319 }
3320 for (i=0 ; i < tuplesize ; i++) {
3321 Py_INCREF(Py_None);
3322 PyTuple_SET_ITEM(result, i, Py_None);
3323 }
3324
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003325 /* create ziplongestobject structure */
3326 lz = (ziplongestobject *)type->tp_alloc(type, 0);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003327 if (lz == NULL) {
3328 Py_DECREF(ittuple);
3329 Py_DECREF(result);
3330 return NULL;
3331 }
3332 lz->ittuple = ittuple;
3333 lz->tuplesize = tuplesize;
3334 lz->numactive = tuplesize;
3335 lz->result = result;
3336 Py_INCREF(fillvalue);
3337 lz->fillvalue = fillvalue;
3338 return (PyObject *)lz;
3339}
3340
3341static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003342zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003343{
3344 PyObject_GC_UnTrack(lz);
3345 Py_XDECREF(lz->ittuple);
3346 Py_XDECREF(lz->result);
3347 Py_XDECREF(lz->fillvalue);
Christian Heimes90aa7642007-12-19 02:45:37 +00003348 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003349}
3350
3351static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003352zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003353{
3354 Py_VISIT(lz->ittuple);
3355 Py_VISIT(lz->result);
3356 Py_VISIT(lz->fillvalue);
3357 return 0;
3358}
3359
3360static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003361zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003362{
Raymond Hettingera9311a32009-11-01 21:02:38 +00003363 Py_ssize_t i;
3364 Py_ssize_t tuplesize = lz->tuplesize;
3365 PyObject *result = lz->result;
3366 PyObject *it;
3367 PyObject *item;
3368 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003369
Raymond Hettingera9311a32009-11-01 21:02:38 +00003370 if (tuplesize == 0)
3371 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003372 if (lz->numactive == 0)
3373 return NULL;
Raymond Hettingera9311a32009-11-01 21:02:38 +00003374 if (Py_REFCNT(result) == 1) {
3375 Py_INCREF(result);
3376 for (i=0 ; i < tuplesize ; i++) {
3377 it = PyTuple_GET_ITEM(lz->ittuple, i);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003378 if (it == NULL) {
3379 Py_INCREF(lz->fillvalue);
3380 item = lz->fillvalue;
3381 } else {
Raymond Hettingera9311a32009-11-01 21:02:38 +00003382 item = PyIter_Next(it);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003383 if (item == NULL) {
Raymond Hettingera9311a32009-11-01 21:02:38 +00003384 lz->numactive -= 1;
3385 if (lz->numactive == 0 || PyErr_Occurred()) {
3386 lz->numactive = 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003387 Py_DECREF(result);
3388 return NULL;
3389 } else {
3390 Py_INCREF(lz->fillvalue);
Raymond Hettingera9311a32009-11-01 21:02:38 +00003391 item = lz->fillvalue;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003392 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3393 Py_DECREF(it);
3394 }
3395 }
3396 }
Raymond Hettingera9311a32009-11-01 21:02:38 +00003397 olditem = PyTuple_GET_ITEM(result, i);
3398 PyTuple_SET_ITEM(result, i, item);
3399 Py_DECREF(olditem);
3400 }
3401 } else {
3402 result = PyTuple_New(tuplesize);
3403 if (result == NULL)
3404 return NULL;
3405 for (i=0 ; i < tuplesize ; i++) {
3406 it = PyTuple_GET_ITEM(lz->ittuple, i);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003407 if (it == NULL) {
3408 Py_INCREF(lz->fillvalue);
3409 item = lz->fillvalue;
3410 } else {
Raymond Hettingera9311a32009-11-01 21:02:38 +00003411 item = PyIter_Next(it);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003412 if (item == NULL) {
Raymond Hettingera9311a32009-11-01 21:02:38 +00003413 lz->numactive -= 1;
3414 if (lz->numactive == 0 || PyErr_Occurred()) {
3415 lz->numactive = 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003416 Py_DECREF(result);
3417 return NULL;
3418 } else {
3419 Py_INCREF(lz->fillvalue);
Raymond Hettingera9311a32009-11-01 21:02:38 +00003420 item = lz->fillvalue;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003421 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3422 Py_DECREF(it);
3423 }
3424 }
3425 }
Raymond Hettingera9311a32009-11-01 21:02:38 +00003426 PyTuple_SET_ITEM(result, i, item);
3427 }
3428 }
3429 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003430}
3431
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003432PyDoc_STRVAR(zip_longest_doc,
3433"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003434\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003435Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00003436the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003437method continues until the longest iterable in the argument sequence\n\
3438is exhausted and then it raises StopIteration. When the shorter iterables\n\
3439are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3440defaults to None or can be specified by a keyword argument.\n\
3441");
3442
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003443static PyTypeObject ziplongest_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00003444 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003445 "itertools.zip_longest", /* tp_name */
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003446 sizeof(ziplongestobject), /* tp_basicsize */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003447 0, /* tp_itemsize */
3448 /* methods */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003449 (destructor)zip_longest_dealloc, /* tp_dealloc */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003450 0, /* tp_print */
3451 0, /* tp_getattr */
3452 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00003453 0, /* tp_reserved */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003454 0, /* tp_repr */
3455 0, /* tp_as_number */
3456 0, /* tp_as_sequence */
3457 0, /* tp_as_mapping */
3458 0, /* tp_hash */
3459 0, /* tp_call */
3460 0, /* tp_str */
3461 PyObject_GenericGetAttr, /* tp_getattro */
3462 0, /* tp_setattro */
3463 0, /* tp_as_buffer */
3464 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3465 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003466 zip_longest_doc, /* tp_doc */
3467 (traverseproc)zip_longest_traverse, /* tp_traverse */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003468 0, /* tp_clear */
3469 0, /* tp_richcompare */
3470 0, /* tp_weaklistoffset */
3471 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003472 (iternextfunc)zip_longest_next, /* tp_iternext */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003473 0, /* tp_methods */
3474 0, /* tp_members */
3475 0, /* tp_getset */
3476 0, /* tp_base */
3477 0, /* tp_dict */
3478 0, /* tp_descr_get */
3479 0, /* tp_descr_set */
3480 0, /* tp_dictoffset */
3481 0, /* tp_init */
3482 0, /* tp_alloc */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003483 zip_longest_new, /* tp_new */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003484 PyObject_GC_Del, /* tp_free */
3485};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003486
3487/* module level code ********************************************************/
3488
3489PyDoc_STRVAR(module_doc,
3490"Functional tools for creating and using iterators.\n\
3491\n\
3492Infinite iterators:\n\
3493count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003494cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003495repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003496\n\
3497Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003498chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
3499compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
3500dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
3501groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003502filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003503islice(seq, [start,] stop [, step]) --> elements from\n\
3504 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003505starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003506tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003507takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003508zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00003509\n\
3510Combinatoric generators:\n\
3511product(p, q, ... [repeat=1]) --> cartesian product\n\
3512permutations(p[, r])\n\
Raymond Hettingerd3a77c02009-11-19 01:23:41 +00003513combinations(p, r)\n\
3514combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003515");
3516
3517
Raymond Hettingerad983e72003-11-12 14:32:26 +00003518static PyMethodDef module_methods[] = {
3519 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3520 {NULL, NULL} /* sentinel */
3521};
3522
Martin v. Löwis1a214512008-06-11 05:26:20 +00003523
3524static struct PyModuleDef itertoolsmodule = {
3525 PyModuleDef_HEAD_INIT,
3526 "itertools",
3527 module_doc,
3528 -1,
3529 module_methods,
3530 NULL,
3531 NULL,
3532 NULL,
3533 NULL
3534};
3535
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003536PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00003537PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003538{
Raymond Hettinger60eca932003-02-09 06:40:58 +00003539 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003540 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003541 char *name;
3542 PyTypeObject *typelist[] = {
Christian Heimes380f7f22008-02-28 11:19:05 +00003543 &combinations_type,
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003544 &cwr_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003545 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003546 &dropwhile_type,
3547 &takewhile_type,
3548 &islice_type,
3549 &starmap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003550 &chain_type,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003551 &compress_type,
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003552 &filterfalse_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003553 &count_type,
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003554 &ziplongest_type,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003555 &permutations_type,
Christian Heimes380f7f22008-02-28 11:19:05 +00003556 &product_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003557 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003558 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003559 NULL
3560 };
3561
Christian Heimes90aa7642007-12-19 02:45:37 +00003562 Py_TYPE(&teedataobject_type) = &PyType_Type;
Martin v. Löwis1a214512008-06-11 05:26:20 +00003563 m = PyModule_Create(&itertoolsmodule);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00003564 if (m == NULL)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003565 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003566
Raymond Hettinger60eca932003-02-09 06:40:58 +00003567 for (i=0 ; typelist[i] != NULL ; i++) {
3568 if (PyType_Ready(typelist[i]) < 0)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003569 return NULL;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003570 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003571 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003572 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003573 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003574 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003575
3576 if (PyType_Ready(&teedataobject_type) < 0)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003577 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +00003578 if (PyType_Ready(&tee_type) < 0)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003579 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003580 if (PyType_Ready(&_grouper_type) < 0)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003581 return NULL;
3582 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003583}