blob: f6d0e2335d955210047db767dbbb8991965cf0b5 [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
Antoine Pitrouc83ea132010-05-09 14:46:46 +00005/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00006 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 {
Antoine Pitrouc83ea132010-05-09 14:46:46 +000015 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000021} groupbyobject;
22
23static PyTypeObject groupby_type;
24static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26static PyObject *
27groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28{
Antoine Pitrouc83ea132010-05-09 14:46:46 +000029 static char *kwargs[] = {"iterable", "key", NULL};
30 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000032
Antoine Pitrouc83ea132010-05-09 14:46:46 +000033 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;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000051}
52
53static void
54groupby_dealloc(groupbyobject *gbo)
55{
Antoine Pitrouc83ea132010-05-09 14:46:46 +000056 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);
62 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{
Antoine Pitrouc83ea132010-05-09 14:46:46 +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);
73 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000074}
75
76static PyObject *
77groupby_next(groupbyobject *gbo)
78{
Antoine Pitrouc83ea132010-05-09 14:46:46 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
Antoine Pitrouc83ea132010-05-09 14:46:46 +000081 /* 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;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000089
Antoine Pitrouc83ea132010-05-09 14:46:46 +000090 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 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000097
Antoine Pitrouc83ea132010-05-09 14:46:46 +000098 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000101
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000102 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 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000113
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000114 tmp = gbo->currkey;
115 gbo->currkey = newkey;
116 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000117
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000118 tmp = gbo->currvalue;
119 gbo->currvalue = newvalue;
120 Py_XDECREF(tmp);
121 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000123 Py_INCREF(gbo->currkey);
124 tmp = gbo->tgtkey;
125 gbo->tgtkey = gbo->currkey;
126 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000127
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000131
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000135}
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 = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000142 PyVarObject_HEAD_INIT(NULL, 0)
143 "itertools.groupby", /* tp_name */
144 sizeof(groupbyobject), /* tp_basicsize */
145 0, /* tp_itemsize */
146 /* methods */
147 (destructor)groupby_dealloc, /* tp_dealloc */
148 0, /* tp_print */
149 0, /* tp_getattr */
150 0, /* tp_setattr */
151 0, /* tp_compare */
152 0, /* tp_repr */
153 0, /* tp_as_number */
154 0, /* tp_as_sequence */
155 0, /* tp_as_mapping */
156 0, /* tp_hash */
157 0, /* tp_call */
158 0, /* tp_str */
159 PyObject_GenericGetAttr, /* tp_getattro */
160 0, /* tp_setattro */
161 0, /* tp_as_buffer */
162 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
163 Py_TPFLAGS_BASETYPE, /* tp_flags */
164 groupby_doc, /* tp_doc */
165 (traverseproc)groupby_traverse, /* tp_traverse */
166 0, /* tp_clear */
167 0, /* tp_richcompare */
168 0, /* tp_weaklistoffset */
169 PyObject_SelfIter, /* tp_iter */
170 (iternextfunc)groupby_next, /* tp_iternext */
171 0, /* tp_methods */
172 0, /* tp_members */
173 0, /* tp_getset */
174 0, /* tp_base */
175 0, /* tp_dict */
176 0, /* tp_descr_get */
177 0, /* tp_descr_set */
178 0, /* tp_dictoffset */
179 0, /* tp_init */
180 0, /* tp_alloc */
181 groupby_new, /* tp_new */
182 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000183};
184
185
186/* _grouper object (internal) ************************************************/
187
188typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000189 PyObject_HEAD
190 PyObject *parent;
191 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000192} _grouperobject;
193
194static PyTypeObject _grouper_type;
195
196static PyObject *
197_grouper_create(groupbyobject *parent, PyObject *tgtkey)
198{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000199 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000200
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000201 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
202 if (igo == NULL)
203 return NULL;
204 igo->parent = (PyObject *)parent;
205 Py_INCREF(parent);
206 igo->tgtkey = tgtkey;
207 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000208
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000209 PyObject_GC_Track(igo);
210 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000211}
212
213static void
214_grouper_dealloc(_grouperobject *igo)
215{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000216 PyObject_GC_UnTrack(igo);
217 Py_DECREF(igo->parent);
218 Py_DECREF(igo->tgtkey);
219 PyObject_GC_Del(igo);
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000220}
221
222static int
223_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
224{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000225 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{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000233 groupbyobject *gbo = (groupbyobject *)igo->parent;
234 PyObject *newvalue, *newkey, *r;
235 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000236
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000237 if (gbo->currvalue == NULL) {
238 newvalue = PyIter_Next(gbo->it);
239 if (newvalue == NULL)
240 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000241
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000242 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 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000253
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000254 assert(gbo->currkey == NULL);
255 gbo->currkey = newkey;
256 gbo->currvalue = newvalue;
257 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000259 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;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000264
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000265 r = gbo->currvalue;
266 gbo->currvalue = NULL;
267 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000268
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000269 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000270}
271
272static PyTypeObject _grouper_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000273 PyVarObject_HEAD_INIT(NULL, 0)
274 "itertools._grouper", /* tp_name */
275 sizeof(_grouperobject), /* tp_basicsize */
276 0, /* tp_itemsize */
277 /* methods */
278 (destructor)_grouper_dealloc, /* tp_dealloc */
279 0, /* tp_print */
280 0, /* tp_getattr */
281 0, /* tp_setattr */
282 0, /* tp_compare */
283 0, /* tp_repr */
284 0, /* tp_as_number */
285 0, /* tp_as_sequence */
286 0, /* tp_as_mapping */
287 0, /* tp_hash */
288 0, /* tp_call */
289 0, /* tp_str */
290 PyObject_GenericGetAttr, /* tp_getattro */
291 0, /* tp_setattro */
292 0, /* tp_as_buffer */
293 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
294 0, /* tp_doc */
295 (traverseproc)_grouper_traverse,/* tp_traverse */
296 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 */
312 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000313};
314
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000315
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000316
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
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000321 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000322 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 {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000330 PyObject_HEAD
331 PyObject *it;
332 int numread;
333 PyObject *nextlink;
334 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000335} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000336
337typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000338 PyObject_HEAD
339 teedataobject *dataobj;
340 int index;
341 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{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000349 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000350
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000351 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
352 if (tdo == NULL)
353 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000354
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000355 tdo->numread = 0;
356 tdo->nextlink = NULL;
357 Py_INCREF(it);
358 tdo->it = it;
359 PyObject_GC_Track(tdo);
360 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000361}
362
363static PyObject *
364teedataobject_jumplink(teedataobject *tdo)
365{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000366 if (tdo->nextlink == NULL)
367 tdo->nextlink = teedataobject_new(tdo->it);
368 Py_XINCREF(tdo->nextlink);
369 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000370}
371
372static PyObject *
373teedataobject_getitem(teedataobject *tdo, int i)
374{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000375 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000376
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000377 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;
388 }
389 Py_INCREF(value);
390 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000391}
392
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000393static int
394teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
395{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000396 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;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000402}
403
404static int
405teedataobject_clear(teedataobject *tdo)
406{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000407 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;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000413}
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{
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000426 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
427 "itertools.tee_dataobject", /* tp_name */
428 sizeof(teedataobject), /* tp_basicsize */
429 0, /* tp_itemsize */
430 /* methods */
431 (destructor)teedataobject_dealloc, /* tp_dealloc */
432 0, /* tp_print */
433 0, /* tp_getattr */
434 0, /* tp_setattr */
435 0, /* tp_compare */
436 0, /* tp_repr */
437 0, /* tp_as_number */
438 0, /* tp_as_sequence */
439 0, /* tp_as_mapping */
440 0, /* tp_hash */
441 0, /* tp_call */
442 0, /* tp_str */
443 PyObject_GenericGetAttr, /* tp_getattro */
444 0, /* tp_setattro */
445 0, /* tp_as_buffer */
446 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
447 teedataobject_doc, /* tp_doc */
448 (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{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000474 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000475
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000476 if (to->index >= LINKCELLS) {
477 link = teedataobject_jumplink(to->dataobj);
478 Py_DECREF(to->dataobj);
479 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;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000487}
488
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000489static int
490tee_traverse(teeobject *to, visitproc visit, void *arg)
491{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000492 Py_VISIT((PyObject *)to->dataobj);
493 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000494}
495
Raymond Hettingerad983e72003-11-12 14:32:26 +0000496static PyObject *
497tee_copy(teeobject *to)
498{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000499 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000500
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000501 newto = PyObject_GC_New(teeobject, &tee_type);
502 if (newto == NULL)
503 return NULL;
504 Py_INCREF(to->dataobj);
505 newto->dataobj = to->dataobj;
506 newto->index = to->index;
507 newto->weakreflist = NULL;
508 PyObject_GC_Track(newto);
509 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000510}
511
512PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
513
514static PyObject *
515tee_fromiterable(PyObject *iterable)
516{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000517 teeobject *to;
518 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000519
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000520 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 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000527
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000528 to = PyObject_GC_New(teeobject, &tee_type);
529 if (to == NULL)
530 goto done;
531 to->dataobj = (teedataobject *)teedataobject_new(it);
532 if (!to->dataobj) {
533 PyObject_GC_Del(to);
534 to = NULL;
535 goto done;
536 }
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000537
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000538 to->index = 0;
539 to->weakreflist = NULL;
540 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000541done:
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000542 Py_XDECREF(it);
543 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000544}
545
546static PyObject *
547tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
548{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000549 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000550
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000551 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
552 return NULL;
553 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000554}
555
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000556static int
557tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000558{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000559 if (to->weakreflist != NULL)
560 PyObject_ClearWeakRefs((PyObject *) to);
561 Py_CLEAR(to->dataobj);
562 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000563}
564
565static void
566tee_dealloc(teeobject *to)
567{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000568 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[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000577 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
578 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000579};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000580
581static PyTypeObject tee_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000582 PyVarObject_HEAD_INIT(NULL, 0)
583 "itertools.tee", /* tp_name */
584 sizeof(teeobject), /* tp_basicsize */
585 0, /* tp_itemsize */
586 /* methods */
587 (destructor)tee_dealloc, /* tp_dealloc */
588 0, /* tp_print */
589 0, /* tp_getattr */
590 0, /* tp_setattr */
591 0, /* tp_compare */
592 0, /* tp_repr */
593 0, /* tp_as_number */
594 0, /* tp_as_sequence */
595 0, /* tp_as_mapping */
596 0, /* tp_hash */
597 0, /* tp_call */
598 0, /* tp_str */
599 0, /* tp_getattro */
600 0, /* tp_setattro */
601 0, /* tp_as_buffer */
602 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
603 teeobject_doc, /* tp_doc */
604 (traverseproc)tee_traverse, /* tp_traverse */
605 (inquiry)tee_clear, /* tp_clear */
606 0, /* tp_richcompare */
607 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
608 PyObject_SelfIter, /* tp_iter */
609 (iternextfunc)tee_next, /* tp_iternext */
610 tee_methods, /* tp_methods */
611 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 */
621 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{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000627 Py_ssize_t i, n=2;
628 PyObject *it, *iterable, *copyable, *result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000629
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000630 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
631 return NULL;
632 if (n < 0) {
633 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
634 return NULL;
635 }
636 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;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000665}
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 {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000674 PyObject_HEAD
675 PyObject *it;
676 PyObject *saved;
677 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000678} 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{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000685 PyObject *it;
686 PyObject *iterable;
687 PyObject *saved;
688 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000689
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000690 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
691 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000692
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000693 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
694 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000695
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000696 /* Get iterator. */
697 it = PyObject_GetIter(iterable);
698 if (it == NULL)
699 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000700
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000701 saved = PyList_New(0);
702 if (saved == NULL) {
703 Py_DECREF(it);
704 return NULL;
705 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000706
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000707 /* 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;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000717
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000718 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000719}
720
721static void
722cycle_dealloc(cycleobject *lz)
723{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000724 PyObject_GC_UnTrack(lz);
725 Py_XDECREF(lz->saved);
726 Py_XDECREF(lz->it);
727 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{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000733 Py_VISIT(lz->it);
734 Py_VISIT(lz->saved);
735 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000736}
737
738static PyObject *
739cycle_next(cycleobject *lz)
740{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000741 PyObject *item;
742 PyObject *it;
743 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000744
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000745 while (1) {
746 item = PyIter_Next(lz->it);
747 if (item != NULL) {
748 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
749 Py_DECREF(item);
750 return NULL;
751 }
752 return item;
753 }
754 if (PyErr_Occurred()) {
755 if (PyErr_ExceptionMatches(PyExc_StopIteration))
756 PyErr_Clear();
757 else
758 return NULL;
759 }
760 if (PyList_Size(lz->saved) == 0)
761 return NULL;
762 it = PyObject_GetIter(lz->saved);
763 if (it == NULL)
764 return NULL;
765 tmp = lz->it;
766 lz->it = it;
767 lz->firstpass = 1;
768 Py_DECREF(tmp);
769 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000770}
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 = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000779 PyVarObject_HEAD_INIT(NULL, 0)
780 "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 */
788 0, /* tp_compare */
789 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 */
806 PyObject_SelfIter, /* tp_iter */
807 (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 */
817 0, /* tp_alloc */
818 cycle_new, /* tp_new */
819 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000820};
821
822
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000823/* dropwhile object **********************************************************/
824
825typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000826 PyObject_HEAD
827 PyObject *func;
828 PyObject *it;
829 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000830} 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{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000837 PyObject *func, *seq;
838 PyObject *it;
839 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000840
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000841 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
842 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000843
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000844 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
845 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000846
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000847 /* Get iterator. */
848 it = PyObject_GetIter(seq);
849 if (it == NULL)
850 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000851
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000852 /* 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;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000862
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000863 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000864}
865
866static void
867dropwhile_dealloc(dropwhileobject *lz)
868{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000869 PyObject_GC_UnTrack(lz);
870 Py_XDECREF(lz->func);
871 Py_XDECREF(lz->it);
872 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{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000878 Py_VISIT(lz->it);
879 Py_VISIT(lz->func);
880 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000881}
882
883static PyObject *
884dropwhile_next(dropwhileobject *lz)
885{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000886 PyObject *item, *good;
887 PyObject *it = lz->it;
888 long ok;
889 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000890
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000891 iternext = *Py_TYPE(it)->tp_iternext;
892 for (;;) {
893 item = iternext(it);
894 if (item == NULL)
895 return NULL;
896 if (lz->start == 1)
897 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000898
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000899 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);
Antoine Pitrouc5bef752012-08-15 23:16:51 +0200906 if (ok == 0) {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000907 lz->start = 1;
908 return item;
909 }
910 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +0200911 if (ok < 0)
912 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000913 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000914}
915
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000916PyDoc_STRVAR(dropwhile_doc,
917"dropwhile(predicate, iterable) --> dropwhile object\n\
918\n\
919Drop items from the iterable while predicate(item) is true.\n\
920Afterwards, return every element until the iterable is exhausted.");
921
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000922static PyTypeObject dropwhile_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000923 PyVarObject_HEAD_INIT(NULL, 0)
924 "itertools.dropwhile", /* tp_name */
925 sizeof(dropwhileobject), /* tp_basicsize */
926 0, /* tp_itemsize */
927 /* methods */
928 (destructor)dropwhile_dealloc, /* tp_dealloc */
929 0, /* tp_print */
930 0, /* tp_getattr */
931 0, /* tp_setattr */
932 0, /* tp_compare */
933 0, /* tp_repr */
934 0, /* tp_as_number */
935 0, /* tp_as_sequence */
936 0, /* tp_as_mapping */
937 0, /* tp_hash */
938 0, /* tp_call */
939 0, /* tp_str */
940 PyObject_GenericGetAttr, /* tp_getattro */
941 0, /* tp_setattro */
942 0, /* tp_as_buffer */
943 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
944 Py_TPFLAGS_BASETYPE, /* tp_flags */
945 dropwhile_doc, /* tp_doc */
946 (traverseproc)dropwhile_traverse, /* tp_traverse */
947 0, /* tp_clear */
948 0, /* tp_richcompare */
949 0, /* tp_weaklistoffset */
950 PyObject_SelfIter, /* tp_iter */
951 (iternextfunc)dropwhile_next, /* tp_iternext */
952 0, /* tp_methods */
953 0, /* tp_members */
954 0, /* tp_getset */
955 0, /* tp_base */
956 0, /* tp_dict */
957 0, /* tp_descr_get */
958 0, /* tp_descr_set */
959 0, /* tp_dictoffset */
960 0, /* tp_init */
961 0, /* tp_alloc */
962 dropwhile_new, /* tp_new */
963 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000964};
965
966
967/* takewhile object **********************************************************/
968
969typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000970 PyObject_HEAD
971 PyObject *func;
972 PyObject *it;
973 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000974} takewhileobject;
975
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000976static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000977
978static PyObject *
979takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
980{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000981 PyObject *func, *seq;
982 PyObject *it;
983 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000984
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000985 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
986 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000987
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000988 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
989 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000990
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000991 /* Get iterator. */
992 it = PyObject_GetIter(seq);
993 if (it == NULL)
994 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000995
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000996 /* create takewhileobject structure */
997 lz = (takewhileobject *)type->tp_alloc(type, 0);
998 if (lz == NULL) {
999 Py_DECREF(it);
1000 return NULL;
1001 }
1002 Py_INCREF(func);
1003 lz->func = func;
1004 lz->it = it;
1005 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001006
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001007 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001008}
1009
1010static void
1011takewhile_dealloc(takewhileobject *lz)
1012{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001013 PyObject_GC_UnTrack(lz);
1014 Py_XDECREF(lz->func);
1015 Py_XDECREF(lz->it);
1016 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001017}
1018
1019static int
1020takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1021{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001022 Py_VISIT(lz->it);
1023 Py_VISIT(lz->func);
1024 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001025}
1026
1027static PyObject *
1028takewhile_next(takewhileobject *lz)
1029{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001030 PyObject *item, *good;
1031 PyObject *it = lz->it;
1032 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001033
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001034 if (lz->stop == 1)
1035 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001036
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001037 item = (*Py_TYPE(it)->tp_iternext)(it);
1038 if (item == NULL)
1039 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001040
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001041 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1042 if (good == NULL) {
1043 Py_DECREF(item);
1044 return NULL;
1045 }
1046 ok = PyObject_IsTrue(good);
1047 Py_DECREF(good);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02001048 if (ok > 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001049 return item;
1050 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02001051 if (ok == 0)
1052 lz->stop = 1;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001053 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001054}
1055
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001056PyDoc_STRVAR(takewhile_doc,
1057"takewhile(predicate, iterable) --> takewhile object\n\
1058\n\
1059Return successive entries from an iterable as long as the \n\
1060predicate evaluates to true for each entry.");
1061
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001062static PyTypeObject takewhile_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001063 PyVarObject_HEAD_INIT(NULL, 0)
1064 "itertools.takewhile", /* tp_name */
1065 sizeof(takewhileobject), /* tp_basicsize */
1066 0, /* tp_itemsize */
1067 /* methods */
1068 (destructor)takewhile_dealloc, /* tp_dealloc */
1069 0, /* tp_print */
1070 0, /* tp_getattr */
1071 0, /* tp_setattr */
1072 0, /* tp_compare */
1073 0, /* tp_repr */
1074 0, /* tp_as_number */
1075 0, /* tp_as_sequence */
1076 0, /* tp_as_mapping */
1077 0, /* tp_hash */
1078 0, /* tp_call */
1079 0, /* tp_str */
1080 PyObject_GenericGetAttr, /* tp_getattro */
1081 0, /* tp_setattro */
1082 0, /* tp_as_buffer */
1083 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1084 Py_TPFLAGS_BASETYPE, /* tp_flags */
1085 takewhile_doc, /* tp_doc */
1086 (traverseproc)takewhile_traverse, /* tp_traverse */
1087 0, /* tp_clear */
1088 0, /* tp_richcompare */
1089 0, /* tp_weaklistoffset */
1090 PyObject_SelfIter, /* tp_iter */
1091 (iternextfunc)takewhile_next, /* tp_iternext */
1092 0, /* tp_methods */
1093 0, /* tp_members */
1094 0, /* tp_getset */
1095 0, /* tp_base */
1096 0, /* tp_dict */
1097 0, /* tp_descr_get */
1098 0, /* tp_descr_set */
1099 0, /* tp_dictoffset */
1100 0, /* tp_init */
1101 0, /* tp_alloc */
1102 takewhile_new, /* tp_new */
1103 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001104};
1105
1106
1107/* islice object ************************************************************/
1108
1109typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001110 PyObject_HEAD
1111 PyObject *it;
1112 Py_ssize_t next;
1113 Py_ssize_t stop;
1114 Py_ssize_t step;
1115 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116} isliceobject;
1117
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001118static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001119
1120static PyObject *
1121islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1122{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001123 PyObject *seq;
1124 Py_ssize_t start=0, stop=-1, step=1;
1125 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1126 Py_ssize_t numargs;
1127 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001128
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001129 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1130 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001131
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001132 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1133 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001134
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001135 numargs = PyTuple_Size(args);
1136 if (numargs == 2) {
1137 if (a1 != Py_None) {
1138 stop = PyInt_AsSsize_t(a1);
1139 if (stop == -1) {
1140 if (PyErr_Occurred())
1141 PyErr_Clear();
1142 PyErr_SetString(PyExc_ValueError,
1143 "Stop argument for islice() must be None or an integer: 0 <= x <= maxint.");
1144 return NULL;
1145 }
1146 }
1147 } else {
1148 if (a1 != Py_None)
1149 start = PyInt_AsSsize_t(a1);
1150 if (start == -1 && PyErr_Occurred())
1151 PyErr_Clear();
1152 if (a2 != Py_None) {
1153 stop = PyInt_AsSsize_t(a2);
1154 if (stop == -1) {
1155 if (PyErr_Occurred())
1156 PyErr_Clear();
1157 PyErr_SetString(PyExc_ValueError,
1158 "Stop argument for islice() must be None or an integer: 0 <= x <= maxint.");
1159 return NULL;
1160 }
1161 }
1162 }
1163 if (start<0 || stop<-1) {
1164 PyErr_SetString(PyExc_ValueError,
1165 "Indices for islice() must be None or an integer: 0 <= x <= maxint.");
1166 return NULL;
1167 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001168
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001169 if (a3 != NULL) {
1170 if (a3 != Py_None)
1171 step = PyInt_AsSsize_t(a3);
1172 if (step == -1 && PyErr_Occurred())
1173 PyErr_Clear();
1174 }
1175 if (step<1) {
1176 PyErr_SetString(PyExc_ValueError,
1177 "Step for islice() must be a positive integer or None.");
1178 return NULL;
1179 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001180
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001181 /* Get iterator. */
1182 it = PyObject_GetIter(seq);
1183 if (it == NULL)
1184 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001185
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001186 /* create isliceobject structure */
1187 lz = (isliceobject *)type->tp_alloc(type, 0);
1188 if (lz == NULL) {
1189 Py_DECREF(it);
1190 return NULL;
1191 }
1192 lz->it = it;
1193 lz->next = start;
1194 lz->stop = stop;
1195 lz->step = step;
1196 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001197
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001198 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001199}
1200
1201static void
1202islice_dealloc(isliceobject *lz)
1203{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001204 PyObject_GC_UnTrack(lz);
1205 Py_XDECREF(lz->it);
1206 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001207}
1208
1209static int
1210islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1211{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001212 Py_VISIT(lz->it);
1213 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001214}
1215
1216static PyObject *
1217islice_next(isliceobject *lz)
1218{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001219 PyObject *item;
1220 PyObject *it = lz->it;
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001221 Py_ssize_t stop = lz->stop;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001222 Py_ssize_t oldnext;
1223 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001224
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001225 iternext = *Py_TYPE(it)->tp_iternext;
1226 while (lz->cnt < lz->next) {
1227 item = iternext(it);
1228 if (item == NULL)
1229 return NULL;
1230 Py_DECREF(item);
1231 lz->cnt++;
1232 }
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001233 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001234 return NULL;
1235 item = iternext(it);
1236 if (item == NULL)
1237 return NULL;
1238 lz->cnt++;
1239 oldnext = lz->next;
Mark Dickinsona96b0d12011-09-24 09:01:16 +01001240 /* The (size_t) cast below avoids the danger of undefined
1241 behaviour from signed integer overflow. */
1242 lz->next += (size_t)lz->step;
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001243 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1244 lz->next = stop;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001245 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001246}
1247
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001248PyDoc_STRVAR(islice_doc,
1249"islice(iterable, [start,] stop [, step]) --> islice object\n\
1250\n\
1251Return an iterator whose next() method returns selected values from an\n\
1252iterable. If start is specified, will skip all preceding elements;\n\
1253otherwise, start defaults to zero. Step defaults to one. If\n\
1254specified as another value, step determines how many values are \n\
1255skipped between successive calls. Works like a slice() on a list\n\
1256but returns an iterator.");
1257
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001258static PyTypeObject islice_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001259 PyVarObject_HEAD_INIT(NULL, 0)
1260 "itertools.islice", /* tp_name */
1261 sizeof(isliceobject), /* tp_basicsize */
1262 0, /* tp_itemsize */
1263 /* methods */
1264 (destructor)islice_dealloc, /* tp_dealloc */
1265 0, /* tp_print */
1266 0, /* tp_getattr */
1267 0, /* tp_setattr */
1268 0, /* tp_compare */
1269 0, /* tp_repr */
1270 0, /* tp_as_number */
1271 0, /* tp_as_sequence */
1272 0, /* tp_as_mapping */
1273 0, /* tp_hash */
1274 0, /* tp_call */
1275 0, /* tp_str */
1276 PyObject_GenericGetAttr, /* tp_getattro */
1277 0, /* tp_setattro */
1278 0, /* tp_as_buffer */
1279 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1280 Py_TPFLAGS_BASETYPE, /* tp_flags */
1281 islice_doc, /* tp_doc */
1282 (traverseproc)islice_traverse, /* tp_traverse */
1283 0, /* tp_clear */
1284 0, /* tp_richcompare */
1285 0, /* tp_weaklistoffset */
1286 PyObject_SelfIter, /* tp_iter */
1287 (iternextfunc)islice_next, /* tp_iternext */
1288 0, /* tp_methods */
1289 0, /* tp_members */
1290 0, /* tp_getset */
1291 0, /* tp_base */
1292 0, /* tp_dict */
1293 0, /* tp_descr_get */
1294 0, /* tp_descr_set */
1295 0, /* tp_dictoffset */
1296 0, /* tp_init */
1297 0, /* tp_alloc */
1298 islice_new, /* tp_new */
1299 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001300};
1301
1302
1303/* starmap object ************************************************************/
1304
1305typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001306 PyObject_HEAD
1307 PyObject *func;
1308 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001309} starmapobject;
1310
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001311static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001312
1313static PyObject *
1314starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1315{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001316 PyObject *func, *seq;
1317 PyObject *it;
1318 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001319
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001320 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1321 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001322
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001323 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1324 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001325
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001326 /* Get iterator. */
1327 it = PyObject_GetIter(seq);
1328 if (it == NULL)
1329 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001330
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001331 /* create starmapobject structure */
1332 lz = (starmapobject *)type->tp_alloc(type, 0);
1333 if (lz == NULL) {
1334 Py_DECREF(it);
1335 return NULL;
1336 }
1337 Py_INCREF(func);
1338 lz->func = func;
1339 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001340
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001341 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001342}
1343
1344static void
1345starmap_dealloc(starmapobject *lz)
1346{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001347 PyObject_GC_UnTrack(lz);
1348 Py_XDECREF(lz->func);
1349 Py_XDECREF(lz->it);
1350 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001351}
1352
1353static int
1354starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1355{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001356 Py_VISIT(lz->it);
1357 Py_VISIT(lz->func);
1358 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001359}
1360
1361static PyObject *
1362starmap_next(starmapobject *lz)
1363{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001364 PyObject *args;
1365 PyObject *result;
1366 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001367
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001368 args = (*Py_TYPE(it)->tp_iternext)(it);
1369 if (args == NULL)
1370 return NULL;
1371 if (!PyTuple_CheckExact(args)) {
1372 PyObject *newargs = PySequence_Tuple(args);
1373 Py_DECREF(args);
1374 if (newargs == NULL)
1375 return NULL;
1376 args = newargs;
1377 }
1378 result = PyObject_Call(lz->func, args, NULL);
1379 Py_DECREF(args);
1380 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001381}
1382
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001383PyDoc_STRVAR(starmap_doc,
1384"starmap(function, sequence) --> starmap object\n\
1385\n\
1386Return an iterator whose values are returned from the function evaluated\n\
1387with a argument tuple taken from the given sequence.");
1388
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001389static PyTypeObject starmap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001390 PyVarObject_HEAD_INIT(NULL, 0)
1391 "itertools.starmap", /* tp_name */
1392 sizeof(starmapobject), /* tp_basicsize */
1393 0, /* tp_itemsize */
1394 /* methods */
1395 (destructor)starmap_dealloc, /* tp_dealloc */
1396 0, /* tp_print */
1397 0, /* tp_getattr */
1398 0, /* tp_setattr */
1399 0, /* tp_compare */
1400 0, /* tp_repr */
1401 0, /* tp_as_number */
1402 0, /* tp_as_sequence */
1403 0, /* tp_as_mapping */
1404 0, /* tp_hash */
1405 0, /* tp_call */
1406 0, /* tp_str */
1407 PyObject_GenericGetAttr, /* tp_getattro */
1408 0, /* tp_setattro */
1409 0, /* tp_as_buffer */
1410 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1411 Py_TPFLAGS_BASETYPE, /* tp_flags */
1412 starmap_doc, /* tp_doc */
1413 (traverseproc)starmap_traverse, /* tp_traverse */
1414 0, /* tp_clear */
1415 0, /* tp_richcompare */
1416 0, /* tp_weaklistoffset */
1417 PyObject_SelfIter, /* tp_iter */
1418 (iternextfunc)starmap_next, /* tp_iternext */
1419 0, /* tp_methods */
1420 0, /* tp_members */
1421 0, /* tp_getset */
1422 0, /* tp_base */
1423 0, /* tp_dict */
1424 0, /* tp_descr_get */
1425 0, /* tp_descr_set */
1426 0, /* tp_dictoffset */
1427 0, /* tp_init */
1428 0, /* tp_alloc */
1429 starmap_new, /* tp_new */
1430 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001431};
1432
1433
1434/* imap object ************************************************************/
1435
1436typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001437 PyObject_HEAD
1438 PyObject *iters;
1439 PyObject *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001440} imapobject;
1441
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001442static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001443
1444static PyObject *
1445imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1446{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001447 PyObject *it, *iters, *func;
1448 imapobject *lz;
1449 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001450
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001451 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
1452 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001453
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001454 numargs = PyTuple_Size(args);
1455 if (numargs < 2) {
1456 PyErr_SetString(PyExc_TypeError,
1457 "imap() must have at least two arguments.");
1458 return NULL;
1459 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001460
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001461 iters = PyTuple_New(numargs-1);
1462 if (iters == NULL)
1463 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001464
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001465 for (i=1 ; i<numargs ; i++) {
1466 /* Get iterator. */
1467 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1468 if (it == NULL) {
1469 Py_DECREF(iters);
1470 return NULL;
1471 }
1472 PyTuple_SET_ITEM(iters, i-1, it);
1473 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001474
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001475 /* create imapobject structure */
1476 lz = (imapobject *)type->tp_alloc(type, 0);
1477 if (lz == NULL) {
1478 Py_DECREF(iters);
1479 return NULL;
1480 }
1481 lz->iters = iters;
1482 func = PyTuple_GET_ITEM(args, 0);
1483 Py_INCREF(func);
1484 lz->func = func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001485
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001486 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001487}
1488
1489static void
1490imap_dealloc(imapobject *lz)
1491{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001492 PyObject_GC_UnTrack(lz);
1493 Py_XDECREF(lz->iters);
1494 Py_XDECREF(lz->func);
1495 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001496}
1497
1498static int
1499imap_traverse(imapobject *lz, visitproc visit, void *arg)
1500{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001501 Py_VISIT(lz->iters);
1502 Py_VISIT(lz->func);
1503 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001504}
1505
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001506/*
Raymond Hettinger2012f172003-02-07 05:32:58 +00001507imap() is an iterator version of __builtins__.map() except that it does
1508not have the None fill-in feature. That was intentionally left out for
1509the following reasons:
1510
1511 1) Itertools are designed to be easily combined and chained together.
1512 Having all tools stop with the shortest input is a unifying principle
1513 that makes it easier to combine finite iterators (supplying data) with
1514 infinite iterators like count() and repeat() (for supplying sequential
1515 or constant arguments to a function).
1516
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001517 2) In typical use cases for combining itertools, having one finite data
Raymond Hettinger2012f172003-02-07 05:32:58 +00001518 supplier run out before another is likely to be an error condition which
1519 should not pass silently by automatically supplying None.
1520
1521 3) The use cases for automatic None fill-in are rare -- not many functions
1522 do something useful when a parameter suddenly switches type and becomes
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001523 None.
Raymond Hettinger2012f172003-02-07 05:32:58 +00001524
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001525 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001526 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001527
1528 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1529*/
1530
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001531static PyObject *
1532imap_next(imapobject *lz)
1533{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001534 PyObject *val;
1535 PyObject *argtuple;
1536 PyObject *result;
1537 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001538
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001539 numargs = PyTuple_Size(lz->iters);
1540 argtuple = PyTuple_New(numargs);
1541 if (argtuple == NULL)
1542 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001543
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001544 for (i=0 ; i<numargs ; i++) {
1545 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1546 if (val == NULL) {
1547 Py_DECREF(argtuple);
1548 return NULL;
1549 }
1550 PyTuple_SET_ITEM(argtuple, i, val);
1551 }
1552 if (lz->func == Py_None)
1553 return argtuple;
1554 result = PyObject_Call(lz->func, argtuple, NULL);
1555 Py_DECREF(argtuple);
1556 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001557}
1558
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001559PyDoc_STRVAR(imap_doc,
1560"imap(func, *iterables) --> imap object\n\
1561\n\
1562Make an iterator that computes the function using arguments from\n\
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001563each of the iterables. Like map() except that it returns\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001564an iterator instead of a list and that it stops when the shortest\n\
1565iterable is exhausted instead of filling in None for shorter\n\
1566iterables.");
1567
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001568static PyTypeObject imap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001569 PyVarObject_HEAD_INIT(NULL, 0)
1570 "itertools.imap", /* tp_name */
1571 sizeof(imapobject), /* tp_basicsize */
1572 0, /* tp_itemsize */
1573 /* methods */
1574 (destructor)imap_dealloc, /* tp_dealloc */
1575 0, /* tp_print */
1576 0, /* tp_getattr */
1577 0, /* tp_setattr */
1578 0, /* tp_compare */
1579 0, /* tp_repr */
1580 0, /* tp_as_number */
1581 0, /* tp_as_sequence */
1582 0, /* tp_as_mapping */
1583 0, /* tp_hash */
1584 0, /* tp_call */
1585 0, /* tp_str */
1586 PyObject_GenericGetAttr, /* tp_getattro */
1587 0, /* tp_setattro */
1588 0, /* tp_as_buffer */
1589 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1590 Py_TPFLAGS_BASETYPE, /* tp_flags */
1591 imap_doc, /* tp_doc */
1592 (traverseproc)imap_traverse, /* tp_traverse */
1593 0, /* tp_clear */
1594 0, /* tp_richcompare */
1595 0, /* tp_weaklistoffset */
1596 PyObject_SelfIter, /* tp_iter */
1597 (iternextfunc)imap_next, /* tp_iternext */
1598 0, /* tp_methods */
1599 0, /* tp_members */
1600 0, /* tp_getset */
1601 0, /* tp_base */
1602 0, /* tp_dict */
1603 0, /* tp_descr_get */
1604 0, /* tp_descr_set */
1605 0, /* tp_dictoffset */
1606 0, /* tp_init */
1607 0, /* tp_alloc */
1608 imap_new, /* tp_new */
1609 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001610};
1611
1612
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001613/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001614
1615typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001616 PyObject_HEAD
1617 PyObject *source; /* Iterator over input iterables */
1618 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001619} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001620
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001621static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001622
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001623static PyObject *
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001624chain_new_internal(PyTypeObject *type, PyObject *source)
1625{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001626 chainobject *lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001627
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001628 lz = (chainobject *)type->tp_alloc(type, 0);
1629 if (lz == NULL) {
1630 Py_DECREF(source);
1631 return NULL;
1632 }
1633
1634 lz->source = source;
1635 lz->active = NULL;
1636 return (PyObject *)lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001637}
1638
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001639static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001640chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001641{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001642 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001643
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001644 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1645 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001646
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001647 source = PyObject_GetIter(args);
1648 if (source == NULL)
1649 return NULL;
1650
1651 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001652}
1653
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001654static PyObject *
1655chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1656{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001657 PyObject *source;
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001658
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001659 source = PyObject_GetIter(arg);
1660 if (source == NULL)
1661 return NULL;
1662
1663 return chain_new_internal(type, source);
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001664}
1665
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001666static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001667chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001668{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001669 PyObject_GC_UnTrack(lz);
1670 Py_XDECREF(lz->active);
1671 Py_XDECREF(lz->source);
1672 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001673}
1674
Raymond Hettinger2012f172003-02-07 05:32:58 +00001675static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001676chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001677{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001678 Py_VISIT(lz->source);
1679 Py_VISIT(lz->active);
1680 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001681}
1682
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001683static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001684chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001685{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001686 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001687
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001688 if (lz->source == NULL)
1689 return NULL; /* already stopped */
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001690
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001691 if (lz->active == NULL) {
1692 PyObject *iterable = PyIter_Next(lz->source);
1693 if (iterable == NULL) {
1694 Py_CLEAR(lz->source);
1695 return NULL; /* no more input sources */
1696 }
1697 lz->active = PyObject_GetIter(iterable);
1698 Py_DECREF(iterable);
1699 if (lz->active == NULL) {
1700 Py_CLEAR(lz->source);
1701 return NULL; /* input not iterable */
1702 }
1703 }
1704 item = PyIter_Next(lz->active);
1705 if (item != NULL)
1706 return item;
1707 if (PyErr_Occurred()) {
1708 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1709 PyErr_Clear();
1710 else
1711 return NULL; /* input raised an exception */
1712 }
1713 Py_CLEAR(lz->active);
1714 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001715}
1716
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001717PyDoc_STRVAR(chain_doc,
1718"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001719\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001720Return a chain object whose .next() method returns elements from the\n\
1721first iterable until it is exhausted, then elements from the next\n\
1722iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001723
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001724PyDoc_STRVAR(chain_from_iterable_doc,
1725"chain.from_iterable(iterable) --> chain object\n\
1726\n\
1727Alternate chain() contructor taking a single iterable argument\n\
1728that evaluates lazily.");
1729
1730static PyMethodDef chain_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001731 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1732 chain_from_iterable_doc},
1733 {NULL, NULL} /* sentinel */
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001734};
1735
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001736static PyTypeObject chain_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001737 PyVarObject_HEAD_INIT(NULL, 0)
1738 "itertools.chain", /* tp_name */
1739 sizeof(chainobject), /* tp_basicsize */
1740 0, /* tp_itemsize */
1741 /* methods */
1742 (destructor)chain_dealloc, /* tp_dealloc */
1743 0, /* tp_print */
1744 0, /* tp_getattr */
1745 0, /* tp_setattr */
1746 0, /* tp_compare */
1747 0, /* tp_repr */
1748 0, /* tp_as_number */
1749 0, /* tp_as_sequence */
1750 0, /* tp_as_mapping */
1751 0, /* tp_hash */
1752 0, /* tp_call */
1753 0, /* tp_str */
1754 PyObject_GenericGetAttr, /* tp_getattro */
1755 0, /* tp_setattro */
1756 0, /* tp_as_buffer */
1757 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1758 Py_TPFLAGS_BASETYPE, /* tp_flags */
1759 chain_doc, /* tp_doc */
1760 (traverseproc)chain_traverse, /* tp_traverse */
1761 0, /* tp_clear */
1762 0, /* tp_richcompare */
1763 0, /* tp_weaklistoffset */
1764 PyObject_SelfIter, /* tp_iter */
1765 (iternextfunc)chain_next, /* tp_iternext */
1766 chain_methods, /* tp_methods */
1767 0, /* tp_members */
1768 0, /* tp_getset */
1769 0, /* tp_base */
1770 0, /* tp_dict */
1771 0, /* tp_descr_get */
1772 0, /* tp_descr_set */
1773 0, /* tp_dictoffset */
1774 0, /* tp_init */
1775 0, /* tp_alloc */
1776 chain_new, /* tp_new */
1777 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001778};
1779
1780
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001781/* product object ************************************************************/
1782
1783typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001784 PyObject_HEAD
1785 PyObject *pools; /* tuple of pool tuples */
1786 Py_ssize_t *indices; /* one index per pool */
1787 PyObject *result; /* most recently returned result tuple */
1788 int stopped; /* set to 1 when the product iterator is exhausted */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001789} productobject;
1790
1791static PyTypeObject product_type;
1792
1793static PyObject *
1794product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1795{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001796 productobject *lz;
1797 Py_ssize_t nargs, npools, repeat=1;
1798 PyObject *pools = NULL;
1799 Py_ssize_t *indices = NULL;
1800 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001801
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001802 if (kwds != NULL) {
1803 char *kwlist[] = {"repeat", 0};
1804 PyObject *tmpargs = PyTuple_New(0);
1805 if (tmpargs == NULL)
1806 return NULL;
1807 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1808 Py_DECREF(tmpargs);
1809 return NULL;
1810 }
1811 Py_DECREF(tmpargs);
1812 if (repeat < 0) {
1813 PyErr_SetString(PyExc_ValueError,
1814 "repeat argument cannot be negative");
1815 return NULL;
1816 }
1817 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001818
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001819 assert(PyTuple_Check(args));
1820 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1821 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001822
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001823 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1824 if (indices == NULL) {
1825 PyErr_NoMemory();
1826 goto error;
1827 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001828
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001829 pools = PyTuple_New(npools);
1830 if (pools == NULL)
1831 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001832
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001833 for (i=0; i < nargs ; ++i) {
1834 PyObject *item = PyTuple_GET_ITEM(args, i);
1835 PyObject *pool = PySequence_Tuple(item);
1836 if (pool == NULL)
1837 goto error;
1838 PyTuple_SET_ITEM(pools, i, pool);
1839 indices[i] = 0;
1840 }
1841 for ( ; i < npools; ++i) {
1842 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1843 Py_INCREF(pool);
1844 PyTuple_SET_ITEM(pools, i, pool);
1845 indices[i] = 0;
1846 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001847
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001848 /* create productobject structure */
1849 lz = (productobject *)type->tp_alloc(type, 0);
1850 if (lz == NULL)
1851 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001852
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001853 lz->pools = pools;
1854 lz->indices = indices;
1855 lz->result = NULL;
1856 lz->stopped = 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001857
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001858 return (PyObject *)lz;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001859
1860error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001861 if (indices != NULL)
1862 PyMem_Free(indices);
1863 Py_XDECREF(pools);
1864 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001865}
1866
1867static void
1868product_dealloc(productobject *lz)
1869{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001870 PyObject_GC_UnTrack(lz);
1871 Py_XDECREF(lz->pools);
1872 Py_XDECREF(lz->result);
1873 if (lz->indices != NULL)
1874 PyMem_Free(lz->indices);
1875 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001876}
1877
1878static int
1879product_traverse(productobject *lz, visitproc visit, void *arg)
1880{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001881 Py_VISIT(lz->pools);
1882 Py_VISIT(lz->result);
1883 return 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001884}
1885
1886static PyObject *
1887product_next(productobject *lz)
1888{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001889 PyObject *pool;
1890 PyObject *elem;
1891 PyObject *oldelem;
1892 PyObject *pools = lz->pools;
1893 PyObject *result = lz->result;
1894 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1895 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001896
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001897 if (lz->stopped)
1898 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001899
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001900 if (result == NULL) {
1901 /* On the first pass, return an initial tuple filled with the
1902 first element from each pool. */
1903 result = PyTuple_New(npools);
1904 if (result == NULL)
1905 goto empty;
1906 lz->result = result;
1907 for (i=0; i < npools; i++) {
1908 pool = PyTuple_GET_ITEM(pools, i);
1909 if (PyTuple_GET_SIZE(pool) == 0)
1910 goto empty;
1911 elem = PyTuple_GET_ITEM(pool, 0);
1912 Py_INCREF(elem);
1913 PyTuple_SET_ITEM(result, i, elem);
1914 }
1915 } else {
1916 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001917
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001918 /* Copy the previous result tuple or re-use it if available */
1919 if (Py_REFCNT(result) > 1) {
1920 PyObject *old_result = result;
1921 result = PyTuple_New(npools);
1922 if (result == NULL)
1923 goto empty;
1924 lz->result = result;
1925 for (i=0; i < npools; i++) {
1926 elem = PyTuple_GET_ITEM(old_result, i);
1927 Py_INCREF(elem);
1928 PyTuple_SET_ITEM(result, i, elem);
1929 }
1930 Py_DECREF(old_result);
1931 }
1932 /* Now, we've got the only copy so we can update it in-place */
1933 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001934
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001935 /* Update the pool indices right-to-left. Only advance to the
1936 next pool when the previous one rolls-over */
1937 for (i=npools-1 ; i >= 0 ; i--) {
1938 pool = PyTuple_GET_ITEM(pools, i);
1939 indices[i]++;
1940 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1941 /* Roll-over and advance to next pool */
1942 indices[i] = 0;
1943 elem = PyTuple_GET_ITEM(pool, 0);
1944 Py_INCREF(elem);
1945 oldelem = PyTuple_GET_ITEM(result, i);
1946 PyTuple_SET_ITEM(result, i, elem);
1947 Py_DECREF(oldelem);
1948 } else {
1949 /* No rollover. Just increment and stop here. */
1950 elem = PyTuple_GET_ITEM(pool, indices[i]);
1951 Py_INCREF(elem);
1952 oldelem = PyTuple_GET_ITEM(result, i);
1953 PyTuple_SET_ITEM(result, i, elem);
1954 Py_DECREF(oldelem);
1955 break;
1956 }
1957 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00001958
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001959 /* If i is negative, then the indices have all rolled-over
1960 and we're done. */
1961 if (i < 0)
1962 goto empty;
1963 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001964
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001965 Py_INCREF(result);
1966 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001967
1968empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001969 lz->stopped = 1;
1970 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001971}
1972
1973PyDoc_STRVAR(product_doc,
1974"product(*iterables) --> product object\n\
1975\n\
Raymond Hettinger73d79632008-02-23 02:20:41 +00001976Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001977For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1978The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1979cycle in a manner similar to an odometer (with the rightmost element changing\n\
1980on every iteration).\n\n\
Raymond Hettinger94034ea2009-02-09 18:39:41 +00001981To compute the product of an iterable with itself, specify the number\n\
1982of repetitions with the optional repeat keyword argument. For example,\n\
1983product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001984product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1985product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1986
1987static PyTypeObject product_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001988 PyVarObject_HEAD_INIT(NULL, 0)
1989 "itertools.product", /* tp_name */
1990 sizeof(productobject), /* tp_basicsize */
1991 0, /* tp_itemsize */
1992 /* methods */
1993 (destructor)product_dealloc, /* tp_dealloc */
1994 0, /* tp_print */
1995 0, /* tp_getattr */
1996 0, /* tp_setattr */
1997 0, /* tp_compare */
1998 0, /* tp_repr */
1999 0, /* tp_as_number */
2000 0, /* tp_as_sequence */
2001 0, /* tp_as_mapping */
2002 0, /* tp_hash */
2003 0, /* tp_call */
2004 0, /* tp_str */
2005 PyObject_GenericGetAttr, /* tp_getattro */
2006 0, /* tp_setattro */
2007 0, /* tp_as_buffer */
2008 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2009 Py_TPFLAGS_BASETYPE, /* tp_flags */
2010 product_doc, /* tp_doc */
2011 (traverseproc)product_traverse, /* tp_traverse */
2012 0, /* tp_clear */
2013 0, /* tp_richcompare */
2014 0, /* tp_weaklistoffset */
2015 PyObject_SelfIter, /* tp_iter */
2016 (iternextfunc)product_next, /* tp_iternext */
2017 0, /* tp_methods */
2018 0, /* tp_members */
2019 0, /* tp_getset */
2020 0, /* tp_base */
2021 0, /* tp_dict */
2022 0, /* tp_descr_get */
2023 0, /* tp_descr_set */
2024 0, /* tp_dictoffset */
2025 0, /* tp_init */
2026 0, /* tp_alloc */
2027 product_new, /* tp_new */
2028 PyObject_GC_Del, /* tp_free */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002029};
2030
2031
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002032/* combinations object ************************************************************/
2033
2034typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002035 PyObject_HEAD
2036 PyObject *pool; /* input converted to a tuple */
2037 Py_ssize_t *indices; /* one index per result element */
2038 PyObject *result; /* most recently returned result tuple */
2039 Py_ssize_t r; /* size of result tuple */
2040 int stopped; /* set to 1 when the combinations iterator is exhausted */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002041} combinationsobject;
2042
2043static PyTypeObject combinations_type;
2044
2045static PyObject *
2046combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2047{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002048 combinationsobject *co;
2049 Py_ssize_t n;
2050 Py_ssize_t r;
2051 PyObject *pool = NULL;
2052 PyObject *iterable = NULL;
2053 Py_ssize_t *indices = NULL;
2054 Py_ssize_t i;
2055 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002056
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002057 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2058 &iterable, &r))
2059 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002060
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002061 pool = PySequence_Tuple(iterable);
2062 if (pool == NULL)
2063 goto error;
2064 n = PyTuple_GET_SIZE(pool);
2065 if (r < 0) {
2066 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2067 goto error;
2068 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002069
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002070 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2071 if (indices == NULL) {
2072 PyErr_NoMemory();
2073 goto error;
2074 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002075
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002076 for (i=0 ; i<r ; i++)
2077 indices[i] = i;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002078
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002079 /* create combinationsobject structure */
2080 co = (combinationsobject *)type->tp_alloc(type, 0);
2081 if (co == NULL)
2082 goto error;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002083
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002084 co->pool = pool;
2085 co->indices = indices;
2086 co->result = NULL;
2087 co->r = r;
2088 co->stopped = r > n ? 1 : 0;
2089
2090 return (PyObject *)co;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002091
2092error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002093 if (indices != NULL)
2094 PyMem_Free(indices);
2095 Py_XDECREF(pool);
2096 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002097}
2098
2099static void
2100combinations_dealloc(combinationsobject *co)
2101{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002102 PyObject_GC_UnTrack(co);
2103 Py_XDECREF(co->pool);
2104 Py_XDECREF(co->result);
2105 if (co->indices != NULL)
2106 PyMem_Free(co->indices);
2107 Py_TYPE(co)->tp_free(co);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002108}
2109
2110static int
2111combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2112{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002113 Py_VISIT(co->pool);
2114 Py_VISIT(co->result);
2115 return 0;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002116}
2117
2118static PyObject *
2119combinations_next(combinationsobject *co)
2120{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002121 PyObject *elem;
2122 PyObject *oldelem;
2123 PyObject *pool = co->pool;
2124 Py_ssize_t *indices = co->indices;
2125 PyObject *result = co->result;
2126 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2127 Py_ssize_t r = co->r;
2128 Py_ssize_t i, j, index;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002129
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002130 if (co->stopped)
2131 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002132
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002133 if (result == NULL) {
2134 /* On the first pass, initialize result tuple using the indices */
2135 result = PyTuple_New(r);
2136 if (result == NULL)
2137 goto empty;
2138 co->result = result;
2139 for (i=0; i<r ; i++) {
2140 index = indices[i];
2141 elem = PyTuple_GET_ITEM(pool, index);
2142 Py_INCREF(elem);
2143 PyTuple_SET_ITEM(result, i, elem);
2144 }
2145 } else {
2146 /* Copy the previous result tuple or re-use it if available */
2147 if (Py_REFCNT(result) > 1) {
2148 PyObject *old_result = result;
2149 result = PyTuple_New(r);
2150 if (result == NULL)
2151 goto empty;
2152 co->result = result;
2153 for (i=0; i<r ; i++) {
2154 elem = PyTuple_GET_ITEM(old_result, i);
2155 Py_INCREF(elem);
2156 PyTuple_SET_ITEM(result, i, elem);
2157 }
2158 Py_DECREF(old_result);
2159 }
2160 /* Now, we've got the only copy so we can update it in-place
2161 * CPython's empty tuple is a singleton and cached in
2162 * PyTuple's freelist.
2163 */
2164 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002165
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002166 /* Scan indices right-to-left until finding one that is not
2167 at its maximum (i + n - r). */
2168 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2169 ;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002170
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002171 /* If i is negative, then the indices are all at
2172 their maximum value and we're done. */
2173 if (i < 0)
2174 goto empty;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002175
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002176 /* Increment the current index which we know is not at its
2177 maximum. Then move back to the right setting each index
2178 to its lowest possible value (one higher than the index
2179 to its left -- this maintains the sort order invariant). */
2180 indices[i]++;
2181 for (j=i+1 ; j<r ; j++)
2182 indices[j] = indices[j-1] + 1;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002183
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002184 /* Update the result tuple for the new indices
2185 starting with i, the leftmost index that changed */
2186 for ( ; i<r ; i++) {
2187 index = indices[i];
2188 elem = PyTuple_GET_ITEM(pool, index);
2189 Py_INCREF(elem);
2190 oldelem = PyTuple_GET_ITEM(result, i);
2191 PyTuple_SET_ITEM(result, i, elem);
2192 Py_DECREF(oldelem);
2193 }
2194 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002195
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002196 Py_INCREF(result);
2197 return result;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002198
2199empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002200 co->stopped = 1;
2201 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002202}
2203
2204PyDoc_STRVAR(combinations_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002205"combinations(iterable, r) --> combinations object\n\
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002206\n\
2207Return successive r-length combinations of elements in the iterable.\n\n\
2208combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2209
2210static PyTypeObject combinations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002211 PyVarObject_HEAD_INIT(NULL, 0)
2212 "itertools.combinations", /* tp_name */
2213 sizeof(combinationsobject), /* tp_basicsize */
2214 0, /* tp_itemsize */
2215 /* methods */
2216 (destructor)combinations_dealloc, /* tp_dealloc */
2217 0, /* tp_print */
2218 0, /* tp_getattr */
2219 0, /* tp_setattr */
2220 0, /* tp_compare */
2221 0, /* tp_repr */
2222 0, /* tp_as_number */
2223 0, /* tp_as_sequence */
2224 0, /* tp_as_mapping */
2225 0, /* tp_hash */
2226 0, /* tp_call */
2227 0, /* tp_str */
2228 PyObject_GenericGetAttr, /* tp_getattro */
2229 0, /* tp_setattro */
2230 0, /* tp_as_buffer */
2231 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2232 Py_TPFLAGS_BASETYPE, /* tp_flags */
2233 combinations_doc, /* tp_doc */
2234 (traverseproc)combinations_traverse, /* tp_traverse */
2235 0, /* tp_clear */
2236 0, /* tp_richcompare */
2237 0, /* tp_weaklistoffset */
2238 PyObject_SelfIter, /* tp_iter */
2239 (iternextfunc)combinations_next, /* tp_iternext */
2240 0, /* tp_methods */
2241 0, /* tp_members */
2242 0, /* tp_getset */
2243 0, /* tp_base */
2244 0, /* tp_dict */
2245 0, /* tp_descr_get */
2246 0, /* tp_descr_set */
2247 0, /* tp_dictoffset */
2248 0, /* tp_init */
2249 0, /* tp_alloc */
2250 combinations_new, /* tp_new */
2251 PyObject_GC_Del, /* tp_free */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002252};
2253
2254
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002255/* combinations with replacement object *******************************************/
2256
2257/* Equivalent to:
2258
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002259 def combinations_with_replacement(iterable, r):
2260 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2261 # number items returned: (n+r-1)! / r! / (n-1)!
2262 pool = tuple(iterable)
2263 n = len(pool)
2264 indices = [0] * r
2265 yield tuple(pool[i] for i in indices)
2266 while 1:
2267 for i in reversed(range(r)):
2268 if indices[i] != n - 1:
2269 break
2270 else:
2271 return
2272 indices[i:] = [indices[i] + 1] * (r - i)
2273 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002274
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002275 def combinations_with_replacement2(iterable, r):
2276 'Alternate version that filters from product()'
2277 pool = tuple(iterable)
2278 n = len(pool)
2279 for indices in product(range(n), repeat=r):
2280 if sorted(indices) == list(indices):
2281 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002282*/
2283typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002284 PyObject_HEAD
2285 PyObject *pool; /* input converted to a tuple */
2286 Py_ssize_t *indices; /* one index per result element */
2287 PyObject *result; /* most recently returned result tuple */
2288 Py_ssize_t r; /* size of result tuple */
2289 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002290} cwrobject;
2291
2292static PyTypeObject cwr_type;
2293
2294static PyObject *
2295cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2296{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002297 cwrobject *co;
2298 Py_ssize_t n;
2299 Py_ssize_t r;
2300 PyObject *pool = NULL;
2301 PyObject *iterable = NULL;
2302 Py_ssize_t *indices = NULL;
2303 Py_ssize_t i;
2304 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002305
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002306 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2307 &iterable, &r))
2308 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002309
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002310 pool = PySequence_Tuple(iterable);
2311 if (pool == NULL)
2312 goto error;
2313 n = PyTuple_GET_SIZE(pool);
2314 if (r < 0) {
2315 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2316 goto error;
2317 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002318
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002319 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2320 if (indices == NULL) {
2321 PyErr_NoMemory();
2322 goto error;
2323 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002324
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002325 for (i=0 ; i<r ; i++)
2326 indices[i] = 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002327
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002328 /* create cwrobject structure */
2329 co = (cwrobject *)type->tp_alloc(type, 0);
2330 if (co == NULL)
2331 goto error;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002332
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002333 co->pool = pool;
2334 co->indices = indices;
2335 co->result = NULL;
2336 co->r = r;
2337 co->stopped = !n && r;
2338
2339 return (PyObject *)co;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002340
2341error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002342 if (indices != NULL)
2343 PyMem_Free(indices);
2344 Py_XDECREF(pool);
2345 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002346}
2347
2348static void
2349cwr_dealloc(cwrobject *co)
2350{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002351 PyObject_GC_UnTrack(co);
2352 Py_XDECREF(co->pool);
2353 Py_XDECREF(co->result);
2354 if (co->indices != NULL)
2355 PyMem_Free(co->indices);
2356 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002357}
2358
2359static int
2360cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2361{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002362 Py_VISIT(co->pool);
2363 Py_VISIT(co->result);
2364 return 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002365}
2366
2367static PyObject *
2368cwr_next(cwrobject *co)
2369{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002370 PyObject *elem;
2371 PyObject *oldelem;
2372 PyObject *pool = co->pool;
2373 Py_ssize_t *indices = co->indices;
2374 PyObject *result = co->result;
2375 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2376 Py_ssize_t r = co->r;
2377 Py_ssize_t i, j, index;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002378
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002379 if (co->stopped)
2380 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002381
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002382 if (result == NULL) {
2383 /* On the first pass, initialize result tuple using the indices */
2384 result = PyTuple_New(r);
2385 if (result == NULL)
2386 goto empty;
2387 co->result = result;
2388 for (i=0; i<r ; i++) {
2389 index = indices[i];
2390 elem = PyTuple_GET_ITEM(pool, index);
2391 Py_INCREF(elem);
2392 PyTuple_SET_ITEM(result, i, elem);
2393 }
2394 } else {
2395 /* Copy the previous result tuple or re-use it if available */
2396 if (Py_REFCNT(result) > 1) {
2397 PyObject *old_result = result;
2398 result = PyTuple_New(r);
2399 if (result == NULL)
2400 goto empty;
2401 co->result = result;
2402 for (i=0; i<r ; i++) {
2403 elem = PyTuple_GET_ITEM(old_result, i);
2404 Py_INCREF(elem);
2405 PyTuple_SET_ITEM(result, i, elem);
2406 }
2407 Py_DECREF(old_result);
2408 }
2409 /* Now, we've got the only copy so we can update it in-place CPython's
2410 empty tuple is a singleton and cached in PyTuple's freelist. */
2411 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002412
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002413 /* Scan indices right-to-left until finding one that is not
2414 * at its maximum (n-1). */
2415 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2416 ;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002417
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002418 /* If i is negative, then the indices are all at
2419 their maximum value and we're done. */
2420 if (i < 0)
2421 goto empty;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002422
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002423 /* Increment the current index which we know is not at its
2424 maximum. Then set all to the right to the same value. */
2425 indices[i]++;
2426 for (j=i+1 ; j<r ; j++)
2427 indices[j] = indices[j-1];
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002428
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002429 /* Update the result tuple for the new indices
2430 starting with i, the leftmost index that changed */
2431 for ( ; i<r ; i++) {
2432 index = indices[i];
2433 elem = PyTuple_GET_ITEM(pool, index);
2434 Py_INCREF(elem);
2435 oldelem = PyTuple_GET_ITEM(result, i);
2436 PyTuple_SET_ITEM(result, i, elem);
2437 Py_DECREF(oldelem);
2438 }
2439 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002440
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002441 Py_INCREF(result);
2442 return result;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002443
2444empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002445 co->stopped = 1;
2446 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002447}
2448
2449PyDoc_STRVAR(cwr_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002450"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002451\n\
2452Return successive r-length combinations of elements in the iterable\n\
2453allowing individual elements to have successive repeats.\n\
2454combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2455
2456static PyTypeObject cwr_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002457 PyVarObject_HEAD_INIT(NULL, 0)
2458 "itertools.combinations_with_replacement", /* tp_name */
2459 sizeof(cwrobject), /* tp_basicsize */
2460 0, /* tp_itemsize */
2461 /* methods */
2462 (destructor)cwr_dealloc, /* tp_dealloc */
2463 0, /* tp_print */
2464 0, /* tp_getattr */
2465 0, /* tp_setattr */
2466 0, /* tp_compare */
2467 0, /* tp_repr */
2468 0, /* tp_as_number */
2469 0, /* tp_as_sequence */
2470 0, /* tp_as_mapping */
2471 0, /* tp_hash */
2472 0, /* tp_call */
2473 0, /* tp_str */
2474 PyObject_GenericGetAttr, /* tp_getattro */
2475 0, /* tp_setattro */
2476 0, /* tp_as_buffer */
2477 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2478 Py_TPFLAGS_BASETYPE, /* tp_flags */
2479 cwr_doc, /* tp_doc */
2480 (traverseproc)cwr_traverse, /* tp_traverse */
2481 0, /* tp_clear */
2482 0, /* tp_richcompare */
2483 0, /* tp_weaklistoffset */
2484 PyObject_SelfIter, /* tp_iter */
2485 (iternextfunc)cwr_next, /* tp_iternext */
2486 0, /* tp_methods */
2487 0, /* tp_members */
2488 0, /* tp_getset */
2489 0, /* tp_base */
2490 0, /* tp_dict */
2491 0, /* tp_descr_get */
2492 0, /* tp_descr_set */
2493 0, /* tp_dictoffset */
2494 0, /* tp_init */
2495 0, /* tp_alloc */
2496 cwr_new, /* tp_new */
2497 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002498};
2499
2500
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002501/* permutations object ************************************************************
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002502
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002503def permutations(iterable, r=None):
2504 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2505 pool = tuple(iterable)
2506 n = len(pool)
2507 r = n if r is None else r
2508 indices = range(n)
2509 cycles = range(n-r+1, n+1)[::-1]
2510 yield tuple(pool[i] for i in indices[:r])
2511 while n:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002512 for i in reversed(range(r)):
2513 cycles[i] -= 1
2514 if cycles[i] == 0:
2515 indices[i:] = indices[i+1:] + indices[i:i+1]
2516 cycles[i] = n - i
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002517 else:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002518 j = cycles[i]
2519 indices[i], indices[-j] = indices[-j], indices[i]
2520 yield tuple(pool[i] for i in indices[:r])
2521 break
2522 else:
2523 return
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002524*/
2525
2526typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002527 PyObject_HEAD
2528 PyObject *pool; /* input converted to a tuple */
2529 Py_ssize_t *indices; /* one index per element in the pool */
2530 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2531 PyObject *result; /* most recently returned result tuple */
2532 Py_ssize_t r; /* size of result tuple */
2533 int stopped; /* set to 1 when the permutations iterator is exhausted */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002534} permutationsobject;
2535
2536static PyTypeObject permutations_type;
2537
2538static PyObject *
2539permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2540{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002541 permutationsobject *po;
2542 Py_ssize_t n;
2543 Py_ssize_t r;
2544 PyObject *robj = Py_None;
2545 PyObject *pool = NULL;
2546 PyObject *iterable = NULL;
2547 Py_ssize_t *indices = NULL;
2548 Py_ssize_t *cycles = NULL;
2549 Py_ssize_t i;
2550 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002551
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002552 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2553 &iterable, &robj))
2554 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002555
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002556 pool = PySequence_Tuple(iterable);
2557 if (pool == NULL)
2558 goto error;
2559 n = PyTuple_GET_SIZE(pool);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002560
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002561 r = n;
2562 if (robj != Py_None) {
2563 r = PyInt_AsSsize_t(robj);
2564 if (r == -1 && PyErr_Occurred())
2565 goto error;
2566 }
2567 if (r < 0) {
2568 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2569 goto error;
2570 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002571
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002572 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2573 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2574 if (indices == NULL || cycles == NULL) {
2575 PyErr_NoMemory();
2576 goto error;
2577 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002578
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002579 for (i=0 ; i<n ; i++)
2580 indices[i] = i;
2581 for (i=0 ; i<r ; i++)
2582 cycles[i] = n - i;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002583
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002584 /* create permutationsobject structure */
2585 po = (permutationsobject *)type->tp_alloc(type, 0);
2586 if (po == NULL)
2587 goto error;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002588
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002589 po->pool = pool;
2590 po->indices = indices;
2591 po->cycles = cycles;
2592 po->result = NULL;
2593 po->r = r;
2594 po->stopped = r > n ? 1 : 0;
2595
2596 return (PyObject *)po;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002597
2598error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002599 if (indices != NULL)
2600 PyMem_Free(indices);
2601 if (cycles != NULL)
2602 PyMem_Free(cycles);
2603 Py_XDECREF(pool);
2604 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002605}
2606
2607static void
2608permutations_dealloc(permutationsobject *po)
2609{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002610 PyObject_GC_UnTrack(po);
2611 Py_XDECREF(po->pool);
2612 Py_XDECREF(po->result);
2613 PyMem_Free(po->indices);
2614 PyMem_Free(po->cycles);
2615 Py_TYPE(po)->tp_free(po);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002616}
2617
2618static int
2619permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2620{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002621 Py_VISIT(po->pool);
2622 Py_VISIT(po->result);
2623 return 0;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002624}
2625
2626static PyObject *
2627permutations_next(permutationsobject *po)
2628{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002629 PyObject *elem;
2630 PyObject *oldelem;
2631 PyObject *pool = po->pool;
2632 Py_ssize_t *indices = po->indices;
2633 Py_ssize_t *cycles = po->cycles;
2634 PyObject *result = po->result;
2635 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2636 Py_ssize_t r = po->r;
2637 Py_ssize_t i, j, k, index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002638
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002639 if (po->stopped)
2640 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002641
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002642 if (result == NULL) {
2643 /* On the first pass, initialize result tuple using the indices */
2644 result = PyTuple_New(r);
2645 if (result == NULL)
2646 goto empty;
2647 po->result = result;
2648 for (i=0; i<r ; i++) {
2649 index = indices[i];
2650 elem = PyTuple_GET_ITEM(pool, index);
2651 Py_INCREF(elem);
2652 PyTuple_SET_ITEM(result, i, elem);
2653 }
2654 } else {
2655 if (n == 0)
2656 goto empty;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002657
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002658 /* Copy the previous result tuple or re-use it if available */
2659 if (Py_REFCNT(result) > 1) {
2660 PyObject *old_result = result;
2661 result = PyTuple_New(r);
2662 if (result == NULL)
2663 goto empty;
2664 po->result = result;
2665 for (i=0; i<r ; i++) {
2666 elem = PyTuple_GET_ITEM(old_result, i);
2667 Py_INCREF(elem);
2668 PyTuple_SET_ITEM(result, i, elem);
2669 }
2670 Py_DECREF(old_result);
2671 }
2672 /* Now, we've got the only copy so we can update it in-place */
2673 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002674
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002675 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2676 for (i=r-1 ; i>=0 ; i--) {
2677 cycles[i] -= 1;
2678 if (cycles[i] == 0) {
2679 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2680 index = indices[i];
2681 for (j=i ; j<n-1 ; j++)
2682 indices[j] = indices[j+1];
2683 indices[n-1] = index;
2684 cycles[i] = n - i;
2685 } else {
2686 j = cycles[i];
2687 index = indices[i];
2688 indices[i] = indices[n-j];
2689 indices[n-j] = index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002690
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002691 for (k=i; k<r ; k++) {
2692 /* start with i, the leftmost element that changed */
2693 /* yield tuple(pool[k] for k in indices[:r]) */
2694 index = indices[k];
2695 elem = PyTuple_GET_ITEM(pool, index);
2696 Py_INCREF(elem);
2697 oldelem = PyTuple_GET_ITEM(result, k);
2698 PyTuple_SET_ITEM(result, k, elem);
2699 Py_DECREF(oldelem);
2700 }
2701 break;
2702 }
2703 }
2704 /* If i is negative, then the cycles have all
2705 rolled-over and we're done. */
2706 if (i < 0)
2707 goto empty;
2708 }
2709 Py_INCREF(result);
2710 return result;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002711
2712empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002713 po->stopped = 1;
2714 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002715}
2716
2717PyDoc_STRVAR(permutations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +00002718"permutations(iterable[, r]) --> permutations object\n\
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002719\n\
2720Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl8550d692008-06-10 12:46:39 +00002721permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002722
2723static PyTypeObject permutations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002724 PyVarObject_HEAD_INIT(NULL, 0)
2725 "itertools.permutations", /* tp_name */
2726 sizeof(permutationsobject), /* tp_basicsize */
2727 0, /* tp_itemsize */
2728 /* methods */
2729 (destructor)permutations_dealloc, /* tp_dealloc */
2730 0, /* tp_print */
2731 0, /* tp_getattr */
2732 0, /* tp_setattr */
2733 0, /* tp_compare */
2734 0, /* tp_repr */
2735 0, /* tp_as_number */
2736 0, /* tp_as_sequence */
2737 0, /* tp_as_mapping */
2738 0, /* tp_hash */
2739 0, /* tp_call */
2740 0, /* tp_str */
2741 PyObject_GenericGetAttr, /* tp_getattro */
2742 0, /* tp_setattro */
2743 0, /* tp_as_buffer */
2744 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2745 Py_TPFLAGS_BASETYPE, /* tp_flags */
2746 permutations_doc, /* tp_doc */
2747 (traverseproc)permutations_traverse, /* tp_traverse */
2748 0, /* tp_clear */
2749 0, /* tp_richcompare */
2750 0, /* tp_weaklistoffset */
2751 PyObject_SelfIter, /* tp_iter */
2752 (iternextfunc)permutations_next, /* tp_iternext */
2753 0, /* tp_methods */
2754 0, /* tp_members */
2755 0, /* tp_getset */
2756 0, /* tp_base */
2757 0, /* tp_dict */
2758 0, /* tp_descr_get */
2759 0, /* tp_descr_set */
2760 0, /* tp_dictoffset */
2761 0, /* tp_init */
2762 0, /* tp_alloc */
2763 permutations_new, /* tp_new */
2764 PyObject_GC_Del, /* tp_free */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002765};
2766
2767
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002768/* compress object ************************************************************/
2769
2770/* Equivalent to:
2771
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002772 def compress(data, selectors):
2773 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2774 return (d for d, s in izip(data, selectors) if s)
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002775*/
2776
2777typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002778 PyObject_HEAD
2779 PyObject *data;
2780 PyObject *selectors;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002781} compressobject;
2782
2783static PyTypeObject compress_type;
2784
2785static PyObject *
2786compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2787{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002788 PyObject *seq1, *seq2;
2789 PyObject *data=NULL, *selectors=NULL;
2790 compressobject *lz;
2791 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002792
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002793 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2794 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002795
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002796 data = PyObject_GetIter(seq1);
2797 if (data == NULL)
2798 goto fail;
2799 selectors = PyObject_GetIter(seq2);
2800 if (selectors == NULL)
2801 goto fail;
2802
2803 /* create compressobject structure */
2804 lz = (compressobject *)type->tp_alloc(type, 0);
2805 if (lz == NULL)
2806 goto fail;
2807 lz->data = data;
2808 lz->selectors = selectors;
2809 return (PyObject *)lz;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002810
2811fail:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002812 Py_XDECREF(data);
2813 Py_XDECREF(selectors);
2814 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002815}
2816
2817static void
2818compress_dealloc(compressobject *lz)
2819{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002820 PyObject_GC_UnTrack(lz);
2821 Py_XDECREF(lz->data);
2822 Py_XDECREF(lz->selectors);
2823 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002824}
2825
2826static int
2827compress_traverse(compressobject *lz, visitproc visit, void *arg)
2828{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002829 Py_VISIT(lz->data);
2830 Py_VISIT(lz->selectors);
2831 return 0;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002832}
2833
2834static PyObject *
2835compress_next(compressobject *lz)
2836{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002837 PyObject *data = lz->data, *selectors = lz->selectors;
2838 PyObject *datum, *selector;
2839 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2840 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2841 int ok;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002842
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002843 while (1) {
2844 /* Steps: get datum, get selector, evaluate selector.
2845 Order is important (to match the pure python version
2846 in terms of which input gets a chance to raise an
2847 exception first).
2848 */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002849
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002850 datum = datanext(data);
2851 if (datum == NULL)
2852 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002853
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002854 selector = selectornext(selectors);
2855 if (selector == NULL) {
2856 Py_DECREF(datum);
2857 return NULL;
2858 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002859
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002860 ok = PyObject_IsTrue(selector);
2861 Py_DECREF(selector);
2862 if (ok == 1)
2863 return datum;
2864 Py_DECREF(datum);
2865 if (ok == -1)
2866 return NULL;
2867 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002868}
2869
2870PyDoc_STRVAR(compress_doc,
Raymond Hettinger2e2909f2009-02-19 02:15:14 +00002871"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002872\n\
2873Return data elements corresponding to true selector elements.\n\
2874Forms a shorter iterator from selected data elements using the\n\
2875selectors to choose the data elements.");
2876
2877static PyTypeObject compress_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002878 PyVarObject_HEAD_INIT(NULL, 0)
2879 "itertools.compress", /* tp_name */
2880 sizeof(compressobject), /* tp_basicsize */
2881 0, /* tp_itemsize */
2882 /* methods */
2883 (destructor)compress_dealloc, /* tp_dealloc */
2884 0, /* tp_print */
2885 0, /* tp_getattr */
2886 0, /* tp_setattr */
2887 0, /* tp_compare */
2888 0, /* tp_repr */
2889 0, /* tp_as_number */
2890 0, /* tp_as_sequence */
2891 0, /* tp_as_mapping */
2892 0, /* tp_hash */
2893 0, /* tp_call */
2894 0, /* tp_str */
2895 PyObject_GenericGetAttr, /* tp_getattro */
2896 0, /* tp_setattro */
2897 0, /* tp_as_buffer */
2898 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2899 Py_TPFLAGS_BASETYPE, /* tp_flags */
2900 compress_doc, /* tp_doc */
2901 (traverseproc)compress_traverse, /* tp_traverse */
2902 0, /* tp_clear */
2903 0, /* tp_richcompare */
2904 0, /* tp_weaklistoffset */
2905 PyObject_SelfIter, /* tp_iter */
2906 (iternextfunc)compress_next, /* tp_iternext */
2907 0, /* tp_methods */
2908 0, /* tp_members */
2909 0, /* tp_getset */
2910 0, /* tp_base */
2911 0, /* tp_dict */
2912 0, /* tp_descr_get */
2913 0, /* tp_descr_set */
2914 0, /* tp_dictoffset */
2915 0, /* tp_init */
2916 0, /* tp_alloc */
2917 compress_new, /* tp_new */
2918 PyObject_GC_Del, /* tp_free */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002919};
2920
2921
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002922/* ifilter object ************************************************************/
2923
2924typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002925 PyObject_HEAD
2926 PyObject *func;
2927 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002928} ifilterobject;
2929
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002930static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002931
2932static PyObject *
2933ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2934{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002935 PyObject *func, *seq;
2936 PyObject *it;
2937 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002938
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002939 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
2940 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002941
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002942 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
2943 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002944
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002945 /* Get iterator. */
2946 it = PyObject_GetIter(seq);
2947 if (it == NULL)
2948 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002949
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002950 /* create ifilterobject structure */
2951 lz = (ifilterobject *)type->tp_alloc(type, 0);
2952 if (lz == NULL) {
2953 Py_DECREF(it);
2954 return NULL;
2955 }
2956 Py_INCREF(func);
2957 lz->func = func;
2958 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002959
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002960 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002961}
2962
2963static void
2964ifilter_dealloc(ifilterobject *lz)
2965{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002966 PyObject_GC_UnTrack(lz);
2967 Py_XDECREF(lz->func);
2968 Py_XDECREF(lz->it);
2969 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002970}
2971
2972static int
2973ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
2974{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002975 Py_VISIT(lz->it);
2976 Py_VISIT(lz->func);
2977 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002978}
2979
2980static PyObject *
2981ifilter_next(ifilterobject *lz)
2982{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002983 PyObject *item;
2984 PyObject *it = lz->it;
2985 long ok;
2986 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002987
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002988 iternext = *Py_TYPE(it)->tp_iternext;
2989 for (;;) {
2990 item = iternext(it);
2991 if (item == NULL)
2992 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002993
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002994 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2995 ok = PyObject_IsTrue(item);
2996 } else {
2997 PyObject *good;
2998 good = PyObject_CallFunctionObjArgs(lz->func,
2999 item, NULL);
3000 if (good == NULL) {
3001 Py_DECREF(item);
3002 return NULL;
3003 }
3004 ok = PyObject_IsTrue(good);
3005 Py_DECREF(good);
3006 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003007 if (ok > 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003008 return item;
3009 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003010 if (ok < 0)
3011 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003012 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003013}
3014
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003015PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003016"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003017\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003018Return those items of sequence for which function(item) is true.\n\
3019If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003020
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003021static PyTypeObject ifilter_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003022 PyVarObject_HEAD_INIT(NULL, 0)
3023 "itertools.ifilter", /* tp_name */
3024 sizeof(ifilterobject), /* tp_basicsize */
3025 0, /* tp_itemsize */
3026 /* methods */
3027 (destructor)ifilter_dealloc, /* tp_dealloc */
3028 0, /* tp_print */
3029 0, /* tp_getattr */
3030 0, /* tp_setattr */
3031 0, /* tp_compare */
3032 0, /* tp_repr */
3033 0, /* tp_as_number */
3034 0, /* tp_as_sequence */
3035 0, /* tp_as_mapping */
3036 0, /* tp_hash */
3037 0, /* tp_call */
3038 0, /* tp_str */
3039 PyObject_GenericGetAttr, /* tp_getattro */
3040 0, /* tp_setattro */
3041 0, /* tp_as_buffer */
3042 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3043 Py_TPFLAGS_BASETYPE, /* tp_flags */
3044 ifilter_doc, /* tp_doc */
3045 (traverseproc)ifilter_traverse, /* tp_traverse */
3046 0, /* tp_clear */
3047 0, /* tp_richcompare */
3048 0, /* tp_weaklistoffset */
3049 PyObject_SelfIter, /* tp_iter */
3050 (iternextfunc)ifilter_next, /* tp_iternext */
3051 0, /* tp_methods */
3052 0, /* tp_members */
3053 0, /* tp_getset */
3054 0, /* tp_base */
3055 0, /* tp_dict */
3056 0, /* tp_descr_get */
3057 0, /* tp_descr_set */
3058 0, /* tp_dictoffset */
3059 0, /* tp_init */
3060 0, /* tp_alloc */
3061 ifilter_new, /* tp_new */
3062 PyObject_GC_Del, /* tp_free */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003063};
3064
3065
3066/* ifilterfalse object ************************************************************/
3067
3068typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003069 PyObject_HEAD
3070 PyObject *func;
3071 PyObject *it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003072} ifilterfalseobject;
3073
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003074static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003075
3076static PyObject *
3077ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3078{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003079 PyObject *func, *seq;
3080 PyObject *it;
3081 ifilterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003082
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003083 if (type == &ifilterfalse_type &&
3084 !_PyArg_NoKeywords("ifilterfalse()", kwds))
3085 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003086
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003087 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
3088 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003089
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003090 /* Get iterator. */
3091 it = PyObject_GetIter(seq);
3092 if (it == NULL)
3093 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003094
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003095 /* create ifilterfalseobject structure */
3096 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
3097 if (lz == NULL) {
3098 Py_DECREF(it);
3099 return NULL;
3100 }
3101 Py_INCREF(func);
3102 lz->func = func;
3103 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003104
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003105 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003106}
3107
3108static void
3109ifilterfalse_dealloc(ifilterfalseobject *lz)
3110{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003111 PyObject_GC_UnTrack(lz);
3112 Py_XDECREF(lz->func);
3113 Py_XDECREF(lz->it);
3114 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003115}
3116
3117static int
3118ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
3119{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003120 Py_VISIT(lz->it);
3121 Py_VISIT(lz->func);
3122 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003123}
3124
3125static PyObject *
3126ifilterfalse_next(ifilterfalseobject *lz)
3127{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003128 PyObject *item;
3129 PyObject *it = lz->it;
3130 long ok;
3131 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003132
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003133 iternext = *Py_TYPE(it)->tp_iternext;
3134 for (;;) {
3135 item = iternext(it);
3136 if (item == NULL)
3137 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003138
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003139 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3140 ok = PyObject_IsTrue(item);
3141 } else {
3142 PyObject *good;
3143 good = PyObject_CallFunctionObjArgs(lz->func,
3144 item, NULL);
3145 if (good == NULL) {
3146 Py_DECREF(item);
3147 return NULL;
3148 }
3149 ok = PyObject_IsTrue(good);
3150 Py_DECREF(good);
3151 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003152 if (ok == 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003153 return item;
3154 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003155 if (ok < 0)
3156 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003157 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003158}
3159
Raymond Hettinger60eca932003-02-09 06:40:58 +00003160PyDoc_STRVAR(ifilterfalse_doc,
3161"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
3162\n\
3163Return those items of sequence for which function(item) is false.\n\
3164If function is None, return the items that are false.");
3165
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003166static PyTypeObject ifilterfalse_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003167 PyVarObject_HEAD_INIT(NULL, 0)
3168 "itertools.ifilterfalse", /* tp_name */
3169 sizeof(ifilterfalseobject), /* tp_basicsize */
3170 0, /* tp_itemsize */
3171 /* methods */
3172 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
3173 0, /* tp_print */
3174 0, /* tp_getattr */
3175 0, /* tp_setattr */
3176 0, /* tp_compare */
3177 0, /* tp_repr */
3178 0, /* tp_as_number */
3179 0, /* tp_as_sequence */
3180 0, /* tp_as_mapping */
3181 0, /* tp_hash */
3182 0, /* tp_call */
3183 0, /* tp_str */
3184 PyObject_GenericGetAttr, /* tp_getattro */
3185 0, /* tp_setattro */
3186 0, /* tp_as_buffer */
3187 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3188 Py_TPFLAGS_BASETYPE, /* tp_flags */
3189 ifilterfalse_doc, /* tp_doc */
3190 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
3191 0, /* tp_clear */
3192 0, /* tp_richcompare */
3193 0, /* tp_weaklistoffset */
3194 PyObject_SelfIter, /* tp_iter */
3195 (iternextfunc)ifilterfalse_next, /* tp_iternext */
3196 0, /* tp_methods */
3197 0, /* tp_members */
3198 0, /* tp_getset */
3199 0, /* tp_base */
3200 0, /* tp_dict */
3201 0, /* tp_descr_get */
3202 0, /* tp_descr_set */
3203 0, /* tp_dictoffset */
3204 0, /* tp_init */
3205 0, /* tp_alloc */
3206 ifilterfalse_new, /* tp_new */
3207 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003208};
3209
3210
3211/* count object ************************************************************/
3212
3213typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003214 PyObject_HEAD
3215 Py_ssize_t cnt;
3216 PyObject *long_cnt;
3217 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003218} countobject;
3219
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003220/* Counting logic and invariants:
3221
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003222fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003223
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003224 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3225 Advances with: cnt += 1
3226 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003227
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003228slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003229
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003230 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3231 All counting is done with python objects (no overflows or underflows).
3232 Advances with: long_cnt += long_step
3233 Step may be zero -- effectively a slow version of repeat(cnt).
3234 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003235*/
3236
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003237static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003238
3239static PyObject *
3240count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3241{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003242 countobject *lz;
3243 int slow_mode = 0;
3244 Py_ssize_t cnt = 0;
3245 PyObject *long_cnt = NULL;
3246 PyObject *long_step = NULL;
3247 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003248
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003249 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3250 kwlist, &long_cnt, &long_step))
3251 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003252
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003253 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3254 (long_step != NULL && !PyNumber_Check(long_step))) {
3255 PyErr_SetString(PyExc_TypeError, "a number is required");
3256 return NULL;
3257 }
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003258
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003259 if (long_cnt != NULL) {
3260 cnt = PyInt_AsSsize_t(long_cnt);
3261 if ((cnt == -1 && PyErr_Occurred()) || !PyInt_Check(long_cnt)) {
3262 PyErr_Clear();
3263 slow_mode = 1;
3264 }
3265 Py_INCREF(long_cnt);
3266 } else {
3267 cnt = 0;
3268 long_cnt = PyInt_FromLong(0);
3269 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003270
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003271 /* If not specified, step defaults to 1 */
3272 if (long_step == NULL) {
3273 long_step = PyInt_FromLong(1);
3274 if (long_step == NULL) {
3275 Py_DECREF(long_cnt);
3276 return NULL;
3277 }
3278 } else
3279 Py_INCREF(long_step);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003280
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003281 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003282
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003283 /* Fast mode only works when the step is 1 */
3284 if (!PyInt_Check(long_step) ||
3285 PyInt_AS_LONG(long_step) != 1) {
3286 slow_mode = 1;
3287 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003288
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003289 if (slow_mode)
3290 cnt = PY_SSIZE_T_MAX;
3291 else
3292 Py_CLEAR(long_cnt);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003293
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003294 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3295 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3296 assert(slow_mode ||
3297 (PyInt_Check(long_step) && PyInt_AS_LONG(long_step) == 1));
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003298
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003299 /* create countobject structure */
3300 lz = (countobject *)type->tp_alloc(type, 0);
3301 if (lz == NULL) {
3302 Py_XDECREF(long_cnt);
3303 return NULL;
3304 }
3305 lz->cnt = cnt;
3306 lz->long_cnt = long_cnt;
3307 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003308
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003309 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003310}
3311
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003312static void
3313count_dealloc(countobject *lz)
3314{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003315 PyObject_GC_UnTrack(lz);
3316 Py_XDECREF(lz->long_cnt);
3317 Py_XDECREF(lz->long_step);
3318 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003319}
3320
Benjamin Peterson062a7c32009-02-16 21:07:52 +00003321static int
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003322count_traverse(countobject *lz, visitproc visit, void *arg)
3323{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003324 Py_VISIT(lz->long_cnt);
3325 Py_VISIT(lz->long_step);
3326 return 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003327}
3328
3329static PyObject *
3330count_nextlong(countobject *lz)
3331{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003332 PyObject *long_cnt;
3333 PyObject *stepped_up;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003334
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003335 long_cnt = lz->long_cnt;
3336 if (long_cnt == NULL) {
3337 /* Switch to slow_mode */
3338 long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
3339 if (long_cnt == NULL)
3340 return NULL;
3341 }
3342 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003343
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003344 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3345 if (stepped_up == NULL)
3346 return NULL;
3347 lz->long_cnt = stepped_up;
3348 return long_cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003349}
3350
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003351static PyObject *
3352count_next(countobject *lz)
3353{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003354 if (lz->cnt == PY_SSIZE_T_MAX)
3355 return count_nextlong(lz);
3356 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003357}
3358
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003359static PyObject *
3360count_repr(countobject *lz)
3361{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003362 PyObject *cnt_repr, *step_repr = NULL;
3363 PyObject *result = NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003364
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003365 if (lz->cnt != PY_SSIZE_T_MAX)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003366 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003367
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003368 cnt_repr = PyObject_Repr(lz->long_cnt);
3369 if (cnt_repr == NULL)
3370 return NULL;
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003371
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003372 if (PyInt_Check(lz->long_step) && PyInt_AS_LONG(lz->long_step) == 1) {
3373 /* Don't display step when it is an integer equal to 1 */
3374 result = PyString_FromFormat("count(%s)",
3375 PyString_AS_STRING(cnt_repr));
3376 } else {
3377 step_repr = PyObject_Repr(lz->long_step);
3378 if (step_repr != NULL)
3379 result = PyString_FromFormat("count(%s, %s)",
3380 PyString_AS_STRING(cnt_repr),
3381 PyString_AS_STRING(step_repr));
3382 }
3383 Py_DECREF(cnt_repr);
3384 Py_XDECREF(step_repr);
3385 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003386}
3387
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003388static PyObject *
3389count_reduce(countobject *lz)
3390{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003391 if (lz->cnt == PY_SSIZE_T_MAX)
3392 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3393 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003394}
3395
3396PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3397
3398static PyMethodDef count_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003399 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3400 count_reduce_doc},
3401 {NULL, NULL} /* sentinel */
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003402};
3403
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003404PyDoc_STRVAR(count_doc,
Ezio Melottifdc1e0d2010-06-11 02:21:25 +00003405 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003406\n\
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003407Return a count object whose .next() method returns consecutive values.\n\
3408Equivalent to:\n\n\
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003409 def count(firstval=0, step=1):\n\
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003410 x = firstval\n\
3411 while 1:\n\
3412 yield x\n\
3413 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003414
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003415static PyTypeObject count_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003416 PyVarObject_HEAD_INIT(NULL, 0)
3417 "itertools.count", /* tp_name */
3418 sizeof(countobject), /* tp_basicsize */
3419 0, /* tp_itemsize */
3420 /* methods */
3421 (destructor)count_dealloc, /* tp_dealloc */
3422 0, /* tp_print */
3423 0, /* tp_getattr */
3424 0, /* tp_setattr */
3425 0, /* tp_compare */
3426 (reprfunc)count_repr, /* tp_repr */
3427 0, /* tp_as_number */
3428 0, /* tp_as_sequence */
3429 0, /* tp_as_mapping */
3430 0, /* tp_hash */
3431 0, /* tp_call */
3432 0, /* tp_str */
3433 PyObject_GenericGetAttr, /* tp_getattro */
3434 0, /* tp_setattro */
3435 0, /* tp_as_buffer */
3436 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3437 Py_TPFLAGS_BASETYPE, /* tp_flags */
3438 count_doc, /* tp_doc */
3439 (traverseproc)count_traverse, /* tp_traverse */
3440 0, /* tp_clear */
3441 0, /* tp_richcompare */
3442 0, /* tp_weaklistoffset */
3443 PyObject_SelfIter, /* tp_iter */
3444 (iternextfunc)count_next, /* tp_iternext */
3445 count_methods, /* tp_methods */
3446 0, /* tp_members */
3447 0, /* tp_getset */
3448 0, /* tp_base */
3449 0, /* tp_dict */
3450 0, /* tp_descr_get */
3451 0, /* tp_descr_set */
3452 0, /* tp_dictoffset */
3453 0, /* tp_init */
3454 0, /* tp_alloc */
3455 count_new, /* tp_new */
3456 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003457};
3458
3459
3460/* izip object ************************************************************/
3461
3462#include "Python.h"
3463
3464typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003465 PyObject_HEAD
3466 Py_ssize_t tuplesize;
3467 PyObject *ittuple; /* tuple of iterators */
3468 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003469} izipobject;
3470
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003471static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003472
3473static PyObject *
3474izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3475{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003476 izipobject *lz;
3477 Py_ssize_t i;
3478 PyObject *ittuple; /* tuple of iterators */
3479 PyObject *result;
3480 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003481
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003482 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
3483 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003484
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003485 /* args must be a tuple */
3486 assert(PyTuple_Check(args));
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003487
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003488 /* obtain iterators */
3489 ittuple = PyTuple_New(tuplesize);
3490 if (ittuple == NULL)
3491 return NULL;
3492 for (i=0; i < tuplesize; ++i) {
3493 PyObject *item = PyTuple_GET_ITEM(args, i);
3494 PyObject *it = PyObject_GetIter(item);
3495 if (it == NULL) {
3496 if (PyErr_ExceptionMatches(PyExc_TypeError))
3497 PyErr_Format(PyExc_TypeError,
3498 "izip argument #%zd must support iteration",
3499 i+1);
3500 Py_DECREF(ittuple);
3501 return NULL;
3502 }
3503 PyTuple_SET_ITEM(ittuple, i, it);
3504 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003505
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003506 /* create a result holder */
3507 result = PyTuple_New(tuplesize);
3508 if (result == NULL) {
3509 Py_DECREF(ittuple);
3510 return NULL;
3511 }
3512 for (i=0 ; i < tuplesize ; i++) {
3513 Py_INCREF(Py_None);
3514 PyTuple_SET_ITEM(result, i, Py_None);
3515 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003516
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003517 /* create izipobject structure */
3518 lz = (izipobject *)type->tp_alloc(type, 0);
3519 if (lz == NULL) {
3520 Py_DECREF(ittuple);
3521 Py_DECREF(result);
3522 return NULL;
3523 }
3524 lz->ittuple = ittuple;
3525 lz->tuplesize = tuplesize;
3526 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003527
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003528 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003529}
3530
3531static void
3532izip_dealloc(izipobject *lz)
3533{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003534 PyObject_GC_UnTrack(lz);
3535 Py_XDECREF(lz->ittuple);
3536 Py_XDECREF(lz->result);
3537 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003538}
3539
3540static int
3541izip_traverse(izipobject *lz, visitproc visit, void *arg)
3542{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003543 Py_VISIT(lz->ittuple);
3544 Py_VISIT(lz->result);
3545 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003546}
3547
3548static PyObject *
3549izip_next(izipobject *lz)
3550{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003551 Py_ssize_t i;
3552 Py_ssize_t tuplesize = lz->tuplesize;
3553 PyObject *result = lz->result;
3554 PyObject *it;
3555 PyObject *item;
3556 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003557
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003558 if (tuplesize == 0)
3559 return NULL;
3560 if (Py_REFCNT(result) == 1) {
3561 Py_INCREF(result);
3562 for (i=0 ; i < tuplesize ; i++) {
3563 it = PyTuple_GET_ITEM(lz->ittuple, i);
3564 item = (*Py_TYPE(it)->tp_iternext)(it);
3565 if (item == NULL) {
3566 Py_DECREF(result);
3567 return NULL;
3568 }
3569 olditem = PyTuple_GET_ITEM(result, i);
3570 PyTuple_SET_ITEM(result, i, item);
3571 Py_DECREF(olditem);
3572 }
3573 } else {
3574 result = PyTuple_New(tuplesize);
3575 if (result == NULL)
3576 return NULL;
3577 for (i=0 ; i < tuplesize ; i++) {
3578 it = PyTuple_GET_ITEM(lz->ittuple, i);
3579 item = (*Py_TYPE(it)->tp_iternext)(it);
3580 if (item == NULL) {
3581 Py_DECREF(result);
3582 return NULL;
3583 }
3584 PyTuple_SET_ITEM(result, i, item);
3585 }
3586 }
3587 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003588}
3589
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003590PyDoc_STRVAR(izip_doc,
3591"izip(iter1 [,iter2 [...]]) --> izip object\n\
3592\n\
3593Return a izip object whose .next() method returns a tuple where\n\
3594the i-th element comes from the i-th iterable argument. The .next()\n\
3595method continues until the shortest iterable in the argument sequence\n\
3596is exhausted and then it raises StopIteration. Works like the zip()\n\
3597function but consumes less memory by returning an iterator instead of\n\
3598a list.");
3599
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003600static PyTypeObject izip_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003601 PyVarObject_HEAD_INIT(NULL, 0)
3602 "itertools.izip", /* tp_name */
3603 sizeof(izipobject), /* tp_basicsize */
3604 0, /* tp_itemsize */
3605 /* methods */
3606 (destructor)izip_dealloc, /* tp_dealloc */
3607 0, /* tp_print */
3608 0, /* tp_getattr */
3609 0, /* tp_setattr */
3610 0, /* tp_compare */
3611 0, /* tp_repr */
3612 0, /* tp_as_number */
3613 0, /* tp_as_sequence */
3614 0, /* tp_as_mapping */
3615 0, /* tp_hash */
3616 0, /* tp_call */
3617 0, /* tp_str */
3618 PyObject_GenericGetAttr, /* tp_getattro */
3619 0, /* tp_setattro */
3620 0, /* tp_as_buffer */
3621 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3622 Py_TPFLAGS_BASETYPE, /* tp_flags */
3623 izip_doc, /* tp_doc */
3624 (traverseproc)izip_traverse, /* tp_traverse */
3625 0, /* tp_clear */
3626 0, /* tp_richcompare */
3627 0, /* tp_weaklistoffset */
3628 PyObject_SelfIter, /* tp_iter */
3629 (iternextfunc)izip_next, /* tp_iternext */
3630 0, /* tp_methods */
3631 0, /* tp_members */
3632 0, /* tp_getset */
3633 0, /* tp_base */
3634 0, /* tp_dict */
3635 0, /* tp_descr_get */
3636 0, /* tp_descr_set */
3637 0, /* tp_dictoffset */
3638 0, /* tp_init */
3639 0, /* tp_alloc */
3640 izip_new, /* tp_new */
3641 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003642};
3643
3644
3645/* repeat object ************************************************************/
3646
3647typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003648 PyObject_HEAD
3649 PyObject *element;
3650 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003651} repeatobject;
3652
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003653static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003654
3655static PyObject *
3656repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3657{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003658 repeatobject *ro;
3659 PyObject *element;
3660 Py_ssize_t cnt = -1;
3661 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003662
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003663 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3664 &element, &cnt))
3665 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003666
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003667 if (PyTuple_Size(args) == 2 && cnt < 0)
3668 cnt = 0;
3669
3670 ro = (repeatobject *)type->tp_alloc(type, 0);
3671 if (ro == NULL)
3672 return NULL;
3673 Py_INCREF(element);
3674 ro->element = element;
3675 ro->cnt = cnt;
3676 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003677}
3678
3679static void
3680repeat_dealloc(repeatobject *ro)
3681{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003682 PyObject_GC_UnTrack(ro);
3683 Py_XDECREF(ro->element);
3684 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003685}
3686
3687static int
3688repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3689{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003690 Py_VISIT(ro->element);
3691 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003692}
3693
3694static PyObject *
3695repeat_next(repeatobject *ro)
3696{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003697 if (ro->cnt == 0)
3698 return NULL;
3699 if (ro->cnt > 0)
3700 ro->cnt--;
3701 Py_INCREF(ro->element);
3702 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003703}
3704
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003705static PyObject *
3706repeat_repr(repeatobject *ro)
3707{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003708 PyObject *result, *objrepr;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003709
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003710 objrepr = PyObject_Repr(ro->element);
3711 if (objrepr == NULL)
3712 return NULL;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003713
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003714 if (ro->cnt == -1)
3715 result = PyString_FromFormat("repeat(%s)",
3716 PyString_AS_STRING(objrepr));
3717 else
3718 result = PyString_FromFormat("repeat(%s, %zd)",
3719 PyString_AS_STRING(objrepr), ro->cnt);
3720 Py_DECREF(objrepr);
3721 return result;
3722}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003723
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003724static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003725repeat_len(repeatobject *ro)
3726{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003727 if (ro->cnt == -1) {
3728 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3729 return NULL;
3730 }
3731 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003732}
3733
Armin Rigof5b3e362006-02-11 21:32:43 +00003734PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003735
3736static PyMethodDef repeat_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003737 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3738 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003739};
3740
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003741PyDoc_STRVAR(repeat_doc,
Raymond Hettinger182edae2009-02-19 02:38:25 +00003742"repeat(object [,times]) -> create an iterator which returns the object\n\
3743for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003744endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003745
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003746static PyTypeObject repeat_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003747 PyVarObject_HEAD_INIT(NULL, 0)
3748 "itertools.repeat", /* tp_name */
3749 sizeof(repeatobject), /* tp_basicsize */
3750 0, /* tp_itemsize */
3751 /* methods */
3752 (destructor)repeat_dealloc, /* tp_dealloc */
3753 0, /* tp_print */
3754 0, /* tp_getattr */
3755 0, /* tp_setattr */
3756 0, /* tp_compare */
3757 (reprfunc)repeat_repr, /* tp_repr */
3758 0, /* tp_as_number */
3759 0, /* tp_as_sequence */
3760 0, /* tp_as_mapping */
3761 0, /* tp_hash */
3762 0, /* tp_call */
3763 0, /* tp_str */
3764 PyObject_GenericGetAttr, /* tp_getattro */
3765 0, /* tp_setattro */
3766 0, /* tp_as_buffer */
3767 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3768 Py_TPFLAGS_BASETYPE, /* tp_flags */
3769 repeat_doc, /* tp_doc */
3770 (traverseproc)repeat_traverse, /* tp_traverse */
3771 0, /* tp_clear */
3772 0, /* tp_richcompare */
3773 0, /* tp_weaklistoffset */
3774 PyObject_SelfIter, /* tp_iter */
3775 (iternextfunc)repeat_next, /* tp_iternext */
3776 repeat_methods, /* tp_methods */
3777 0, /* tp_members */
3778 0, /* tp_getset */
3779 0, /* tp_base */
3780 0, /* tp_dict */
3781 0, /* tp_descr_get */
3782 0, /* tp_descr_set */
3783 0, /* tp_dictoffset */
3784 0, /* tp_init */
3785 0, /* tp_alloc */
3786 repeat_new, /* tp_new */
3787 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003788};
3789
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003790/* iziplongest object ************************************************************/
3791
3792#include "Python.h"
3793
3794typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003795 PyObject_HEAD
3796 Py_ssize_t tuplesize;
3797 Py_ssize_t numactive;
3798 PyObject *ittuple; /* tuple of iterators */
3799 PyObject *result;
3800 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003801} iziplongestobject;
3802
3803static PyTypeObject iziplongest_type;
3804
3805static PyObject *
3806izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3807{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003808 iziplongestobject *lz;
3809 Py_ssize_t i;
3810 PyObject *ittuple; /* tuple of iterators */
3811 PyObject *result;
3812 PyObject *fillvalue = Py_None;
3813 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003814
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003815 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3816 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3817 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3818 PyErr_SetString(PyExc_TypeError,
3819 "izip_longest() got an unexpected keyword argument");
3820 return NULL;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003821 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003822 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003823
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003824 /* args must be a tuple */
3825 assert(PyTuple_Check(args));
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003826
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003827 /* obtain iterators */
3828 ittuple = PyTuple_New(tuplesize);
3829 if (ittuple == NULL)
3830 return NULL;
3831 for (i=0; i < tuplesize; ++i) {
3832 PyObject *item = PyTuple_GET_ITEM(args, i);
3833 PyObject *it = PyObject_GetIter(item);
3834 if (it == NULL) {
3835 if (PyErr_ExceptionMatches(PyExc_TypeError))
3836 PyErr_Format(PyExc_TypeError,
3837 "izip_longest argument #%zd must support iteration",
3838 i+1);
3839 Py_DECREF(ittuple);
3840 return NULL;
3841 }
3842 PyTuple_SET_ITEM(ittuple, i, it);
3843 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003844
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003845 /* create a result holder */
3846 result = PyTuple_New(tuplesize);
3847 if (result == NULL) {
3848 Py_DECREF(ittuple);
3849 return NULL;
3850 }
3851 for (i=0 ; i < tuplesize ; i++) {
3852 Py_INCREF(Py_None);
3853 PyTuple_SET_ITEM(result, i, Py_None);
3854 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003855
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003856 /* create iziplongestobject structure */
3857 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3858 if (lz == NULL) {
3859 Py_DECREF(ittuple);
3860 Py_DECREF(result);
3861 return NULL;
3862 }
3863 lz->ittuple = ittuple;
3864 lz->tuplesize = tuplesize;
3865 lz->numactive = tuplesize;
3866 lz->result = result;
3867 Py_INCREF(fillvalue);
3868 lz->fillvalue = fillvalue;
3869 return (PyObject *)lz;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003870}
3871
3872static void
3873izip_longest_dealloc(iziplongestobject *lz)
3874{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003875 PyObject_GC_UnTrack(lz);
3876 Py_XDECREF(lz->ittuple);
3877 Py_XDECREF(lz->result);
3878 Py_XDECREF(lz->fillvalue);
3879 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003880}
3881
3882static int
3883izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3884{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003885 Py_VISIT(lz->ittuple);
3886 Py_VISIT(lz->result);
3887 Py_VISIT(lz->fillvalue);
3888 return 0;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003889}
3890
3891static PyObject *
3892izip_longest_next(iziplongestobject *lz)
3893{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003894 Py_ssize_t i;
3895 Py_ssize_t tuplesize = lz->tuplesize;
3896 PyObject *result = lz->result;
3897 PyObject *it;
3898 PyObject *item;
3899 PyObject *olditem;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003900
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003901 if (tuplesize == 0)
3902 return NULL;
3903 if (lz->numactive == 0)
3904 return NULL;
3905 if (Py_REFCNT(result) == 1) {
3906 Py_INCREF(result);
3907 for (i=0 ; i < tuplesize ; i++) {
3908 it = PyTuple_GET_ITEM(lz->ittuple, i);
3909 if (it == NULL) {
3910 Py_INCREF(lz->fillvalue);
3911 item = lz->fillvalue;
3912 } else {
3913 item = PyIter_Next(it);
3914 if (item == NULL) {
3915 lz->numactive -= 1;
3916 if (lz->numactive == 0 || PyErr_Occurred()) {
3917 lz->numactive = 0;
3918 Py_DECREF(result);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003919 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003920 } else {
3921 Py_INCREF(lz->fillvalue);
3922 item = lz->fillvalue;
3923 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3924 Py_DECREF(it);
3925 }
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003926 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003927 }
3928 olditem = PyTuple_GET_ITEM(result, i);
3929 PyTuple_SET_ITEM(result, i, item);
3930 Py_DECREF(olditem);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003931 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003932 } else {
3933 result = PyTuple_New(tuplesize);
3934 if (result == NULL)
3935 return NULL;
3936 for (i=0 ; i < tuplesize ; i++) {
3937 it = PyTuple_GET_ITEM(lz->ittuple, i);
3938 if (it == NULL) {
3939 Py_INCREF(lz->fillvalue);
3940 item = lz->fillvalue;
3941 } else {
3942 item = PyIter_Next(it);
3943 if (item == NULL) {
3944 lz->numactive -= 1;
3945 if (lz->numactive == 0 || PyErr_Occurred()) {
3946 lz->numactive = 0;
3947 Py_DECREF(result);
3948 return NULL;
3949 } else {
3950 Py_INCREF(lz->fillvalue);
3951 item = lz->fillvalue;
3952 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3953 Py_DECREF(it);
3954 }
3955 }
3956 }
3957 PyTuple_SET_ITEM(result, i, item);
3958 }
3959 }
3960 return result;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003961}
3962
3963PyDoc_STRVAR(izip_longest_doc,
3964"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3965\n\
3966Return an izip_longest object whose .next() method returns a tuple where\n\
3967the i-th element comes from the i-th iterable argument. The .next()\n\
3968method continues until the longest iterable in the argument sequence\n\
3969is exhausted and then it raises StopIteration. When the shorter iterables\n\
3970are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3971defaults to None or can be specified by a keyword argument.\n\
3972");
3973
3974static PyTypeObject iziplongest_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003975 PyVarObject_HEAD_INIT(NULL, 0)
3976 "itertools.izip_longest", /* tp_name */
3977 sizeof(iziplongestobject), /* tp_basicsize */
3978 0, /* tp_itemsize */
3979 /* methods */
3980 (destructor)izip_longest_dealloc, /* tp_dealloc */
3981 0, /* tp_print */
3982 0, /* tp_getattr */
3983 0, /* tp_setattr */
3984 0, /* tp_compare */
3985 0, /* tp_repr */
3986 0, /* tp_as_number */
3987 0, /* tp_as_sequence */
3988 0, /* tp_as_mapping */
3989 0, /* tp_hash */
3990 0, /* tp_call */
3991 0, /* tp_str */
3992 PyObject_GenericGetAttr, /* tp_getattro */
3993 0, /* tp_setattro */
3994 0, /* tp_as_buffer */
3995 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3996 Py_TPFLAGS_BASETYPE, /* tp_flags */
3997 izip_longest_doc, /* tp_doc */
3998 (traverseproc)izip_longest_traverse, /* tp_traverse */
3999 0, /* tp_clear */
4000 0, /* tp_richcompare */
4001 0, /* tp_weaklistoffset */
4002 PyObject_SelfIter, /* tp_iter */
4003 (iternextfunc)izip_longest_next, /* tp_iternext */
4004 0, /* tp_methods */
4005 0, /* tp_members */
4006 0, /* tp_getset */
4007 0, /* tp_base */
4008 0, /* tp_dict */
4009 0, /* tp_descr_get */
4010 0, /* tp_descr_set */
4011 0, /* tp_dictoffset */
4012 0, /* tp_init */
4013 0, /* tp_alloc */
4014 izip_longest_new, /* tp_new */
4015 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd36862c2007-02-21 05:20:38 +00004016};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004017
4018/* module level code ********************************************************/
4019
4020PyDoc_STRVAR(module_doc,
4021"Functional tools for creating and using iterators.\n\
4022\n\
4023Infinite iterators:\n\
4024count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004025cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004026repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004027\n\
4028Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004029chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettingere76816b2009-01-29 03:43:44 +00004030compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004031dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4032groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00004033ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
4034ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004035islice(seq, [start,] stop [, step]) --> elements from\n\
4036 seq[start:stop:step]\n\
4037imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
4038starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004039tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004040takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004041izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4042izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4043\n\
4044Combinatoric generators:\n\
4045product(p, q, ... [repeat=1]) --> cartesian product\n\
4046permutations(p[, r])\n\
Raymond Hettinger9eac1192009-11-19 01:22:04 +00004047combinations(p, r)\n\
4048combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004049");
4050
4051
Raymond Hettingerad983e72003-11-12 14:32:26 +00004052static PyMethodDef module_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004053 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4054 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004055};
4056
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004057PyMODINIT_FUNC
4058inititertools(void)
4059{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004060 int i;
4061 PyObject *m;
4062 char *name;
4063 PyTypeObject *typelist[] = {
4064 &combinations_type,
4065 &cwr_type,
4066 &cycle_type,
4067 &dropwhile_type,
4068 &takewhile_type,
4069 &islice_type,
4070 &starmap_type,
4071 &imap_type,
4072 &chain_type,
4073 &compress_type,
4074 &ifilter_type,
4075 &ifilterfalse_type,
4076 &count_type,
4077 &izip_type,
4078 &iziplongest_type,
4079 &permutations_type,
4080 &product_type,
4081 &repeat_type,
4082 &groupby_type,
4083 NULL
4084 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004085
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004086 Py_TYPE(&teedataobject_type) = &PyType_Type;
4087 m = Py_InitModule3("itertools", module_methods, module_doc);
4088 if (m == NULL)
4089 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004090
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004091 for (i=0 ; typelist[i] != NULL ; i++) {
4092 if (PyType_Ready(typelist[i]) < 0)
4093 return;
4094 name = strchr(typelist[i]->tp_name, '.');
4095 assert (name != NULL);
4096 Py_INCREF(typelist[i]);
4097 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4098 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004099
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004100 if (PyType_Ready(&teedataobject_type) < 0)
4101 return;
4102 if (PyType_Ready(&tee_type) < 0)
4103 return;
4104 if (PyType_Ready(&_grouper_type) < 0)
4105 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004106}