blob: b51ccf96436c315b8c376cb7db78ac39536cc2a6 [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);
906 if (!ok) {
907 lz->start = 1;
908 return item;
909 }
910 Py_DECREF(item);
911 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000912}
913
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000914PyDoc_STRVAR(dropwhile_doc,
915"dropwhile(predicate, iterable) --> dropwhile object\n\
916\n\
917Drop items from the iterable while predicate(item) is true.\n\
918Afterwards, return every element until the iterable is exhausted.");
919
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000920static PyTypeObject dropwhile_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000921 PyVarObject_HEAD_INIT(NULL, 0)
922 "itertools.dropwhile", /* tp_name */
923 sizeof(dropwhileobject), /* tp_basicsize */
924 0, /* tp_itemsize */
925 /* methods */
926 (destructor)dropwhile_dealloc, /* tp_dealloc */
927 0, /* tp_print */
928 0, /* tp_getattr */
929 0, /* tp_setattr */
930 0, /* tp_compare */
931 0, /* tp_repr */
932 0, /* tp_as_number */
933 0, /* tp_as_sequence */
934 0, /* tp_as_mapping */
935 0, /* tp_hash */
936 0, /* tp_call */
937 0, /* tp_str */
938 PyObject_GenericGetAttr, /* tp_getattro */
939 0, /* tp_setattro */
940 0, /* tp_as_buffer */
941 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
942 Py_TPFLAGS_BASETYPE, /* tp_flags */
943 dropwhile_doc, /* tp_doc */
944 (traverseproc)dropwhile_traverse, /* tp_traverse */
945 0, /* tp_clear */
946 0, /* tp_richcompare */
947 0, /* tp_weaklistoffset */
948 PyObject_SelfIter, /* tp_iter */
949 (iternextfunc)dropwhile_next, /* tp_iternext */
950 0, /* tp_methods */
951 0, /* tp_members */
952 0, /* tp_getset */
953 0, /* tp_base */
954 0, /* tp_dict */
955 0, /* tp_descr_get */
956 0, /* tp_descr_set */
957 0, /* tp_dictoffset */
958 0, /* tp_init */
959 0, /* tp_alloc */
960 dropwhile_new, /* tp_new */
961 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000962};
963
964
965/* takewhile object **********************************************************/
966
967typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000968 PyObject_HEAD
969 PyObject *func;
970 PyObject *it;
971 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000972} takewhileobject;
973
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000974static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000975
976static PyObject *
977takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
978{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000979 PyObject *func, *seq;
980 PyObject *it;
981 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000982
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000983 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
984 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000985
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000986 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
987 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000988
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000989 /* Get iterator. */
990 it = PyObject_GetIter(seq);
991 if (it == NULL)
992 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000993
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000994 /* create takewhileobject structure */
995 lz = (takewhileobject *)type->tp_alloc(type, 0);
996 if (lz == NULL) {
997 Py_DECREF(it);
998 return NULL;
999 }
1000 Py_INCREF(func);
1001 lz->func = func;
1002 lz->it = it;
1003 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001004
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001005 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001006}
1007
1008static void
1009takewhile_dealloc(takewhileobject *lz)
1010{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001011 PyObject_GC_UnTrack(lz);
1012 Py_XDECREF(lz->func);
1013 Py_XDECREF(lz->it);
1014 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001015}
1016
1017static int
1018takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1019{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001020 Py_VISIT(lz->it);
1021 Py_VISIT(lz->func);
1022 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001023}
1024
1025static PyObject *
1026takewhile_next(takewhileobject *lz)
1027{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001028 PyObject *item, *good;
1029 PyObject *it = lz->it;
1030 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001031
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001032 if (lz->stop == 1)
1033 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001034
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001035 item = (*Py_TYPE(it)->tp_iternext)(it);
1036 if (item == NULL)
1037 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001038
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001039 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1040 if (good == NULL) {
1041 Py_DECREF(item);
1042 return NULL;
1043 }
1044 ok = PyObject_IsTrue(good);
1045 Py_DECREF(good);
1046 if (ok)
1047 return item;
1048 Py_DECREF(item);
1049 lz->stop = 1;
1050 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001051}
1052
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053PyDoc_STRVAR(takewhile_doc,
1054"takewhile(predicate, iterable) --> takewhile object\n\
1055\n\
1056Return successive entries from an iterable as long as the \n\
1057predicate evaluates to true for each entry.");
1058
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001059static PyTypeObject takewhile_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001060 PyVarObject_HEAD_INIT(NULL, 0)
1061 "itertools.takewhile", /* tp_name */
1062 sizeof(takewhileobject), /* tp_basicsize */
1063 0, /* tp_itemsize */
1064 /* methods */
1065 (destructor)takewhile_dealloc, /* tp_dealloc */
1066 0, /* tp_print */
1067 0, /* tp_getattr */
1068 0, /* tp_setattr */
1069 0, /* tp_compare */
1070 0, /* tp_repr */
1071 0, /* tp_as_number */
1072 0, /* tp_as_sequence */
1073 0, /* tp_as_mapping */
1074 0, /* tp_hash */
1075 0, /* tp_call */
1076 0, /* tp_str */
1077 PyObject_GenericGetAttr, /* tp_getattro */
1078 0, /* tp_setattro */
1079 0, /* tp_as_buffer */
1080 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1081 Py_TPFLAGS_BASETYPE, /* tp_flags */
1082 takewhile_doc, /* tp_doc */
1083 (traverseproc)takewhile_traverse, /* tp_traverse */
1084 0, /* tp_clear */
1085 0, /* tp_richcompare */
1086 0, /* tp_weaklistoffset */
1087 PyObject_SelfIter, /* tp_iter */
1088 (iternextfunc)takewhile_next, /* tp_iternext */
1089 0, /* tp_methods */
1090 0, /* tp_members */
1091 0, /* tp_getset */
1092 0, /* tp_base */
1093 0, /* tp_dict */
1094 0, /* tp_descr_get */
1095 0, /* tp_descr_set */
1096 0, /* tp_dictoffset */
1097 0, /* tp_init */
1098 0, /* tp_alloc */
1099 takewhile_new, /* tp_new */
1100 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001101};
1102
1103
1104/* islice object ************************************************************/
1105
1106typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001107 PyObject_HEAD
1108 PyObject *it;
1109 Py_ssize_t next;
1110 Py_ssize_t stop;
1111 Py_ssize_t step;
1112 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001113} isliceobject;
1114
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001115static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116
1117static PyObject *
1118islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1119{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001120 PyObject *seq;
1121 Py_ssize_t start=0, stop=-1, step=1;
1122 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1123 Py_ssize_t numargs;
1124 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001125
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001126 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1127 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001128
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001129 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1130 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001131
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001132 numargs = PyTuple_Size(args);
1133 if (numargs == 2) {
1134 if (a1 != Py_None) {
1135 stop = PyInt_AsSsize_t(a1);
1136 if (stop == -1) {
1137 if (PyErr_Occurred())
1138 PyErr_Clear();
1139 PyErr_SetString(PyExc_ValueError,
1140 "Stop argument for islice() must be None or an integer: 0 <= x <= maxint.");
1141 return NULL;
1142 }
1143 }
1144 } else {
1145 if (a1 != Py_None)
1146 start = PyInt_AsSsize_t(a1);
1147 if (start == -1 && PyErr_Occurred())
1148 PyErr_Clear();
1149 if (a2 != Py_None) {
1150 stop = PyInt_AsSsize_t(a2);
1151 if (stop == -1) {
1152 if (PyErr_Occurred())
1153 PyErr_Clear();
1154 PyErr_SetString(PyExc_ValueError,
1155 "Stop argument for islice() must be None or an integer: 0 <= x <= maxint.");
1156 return NULL;
1157 }
1158 }
1159 }
1160 if (start<0 || stop<-1) {
1161 PyErr_SetString(PyExc_ValueError,
1162 "Indices for islice() must be None or an integer: 0 <= x <= maxint.");
1163 return NULL;
1164 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001165
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001166 if (a3 != NULL) {
1167 if (a3 != Py_None)
1168 step = PyInt_AsSsize_t(a3);
1169 if (step == -1 && PyErr_Occurred())
1170 PyErr_Clear();
1171 }
1172 if (step<1) {
1173 PyErr_SetString(PyExc_ValueError,
1174 "Step for islice() must be a positive integer or None.");
1175 return NULL;
1176 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001177
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001178 /* Get iterator. */
1179 it = PyObject_GetIter(seq);
1180 if (it == NULL)
1181 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001182
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001183 /* create isliceobject structure */
1184 lz = (isliceobject *)type->tp_alloc(type, 0);
1185 if (lz == NULL) {
1186 Py_DECREF(it);
1187 return NULL;
1188 }
1189 lz->it = it;
1190 lz->next = start;
1191 lz->stop = stop;
1192 lz->step = step;
1193 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001194
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001195 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196}
1197
1198static void
1199islice_dealloc(isliceobject *lz)
1200{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001201 PyObject_GC_UnTrack(lz);
1202 Py_XDECREF(lz->it);
1203 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204}
1205
1206static int
1207islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1208{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001209 Py_VISIT(lz->it);
1210 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001211}
1212
1213static PyObject *
1214islice_next(isliceobject *lz)
1215{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001216 PyObject *item;
1217 PyObject *it = lz->it;
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001218 Py_ssize_t stop = lz->stop;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001219 Py_ssize_t oldnext;
1220 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001221
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001222 iternext = *Py_TYPE(it)->tp_iternext;
1223 while (lz->cnt < lz->next) {
1224 item = iternext(it);
1225 if (item == NULL)
1226 return NULL;
1227 Py_DECREF(item);
1228 lz->cnt++;
1229 }
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001230 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001231 return NULL;
1232 item = iternext(it);
1233 if (item == NULL)
1234 return NULL;
1235 lz->cnt++;
1236 oldnext = lz->next;
Mark Dickinsona96b0d12011-09-24 09:01:16 +01001237 /* The (size_t) cast below avoids the danger of undefined
1238 behaviour from signed integer overflow. */
1239 lz->next += (size_t)lz->step;
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001240 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1241 lz->next = stop;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001242 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001243}
1244
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001245PyDoc_STRVAR(islice_doc,
1246"islice(iterable, [start,] stop [, step]) --> islice object\n\
1247\n\
1248Return an iterator whose next() method returns selected values from an\n\
1249iterable. If start is specified, will skip all preceding elements;\n\
1250otherwise, start defaults to zero. Step defaults to one. If\n\
1251specified as another value, step determines how many values are \n\
1252skipped between successive calls. Works like a slice() on a list\n\
1253but returns an iterator.");
1254
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001255static PyTypeObject islice_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001256 PyVarObject_HEAD_INIT(NULL, 0)
1257 "itertools.islice", /* tp_name */
1258 sizeof(isliceobject), /* tp_basicsize */
1259 0, /* tp_itemsize */
1260 /* methods */
1261 (destructor)islice_dealloc, /* tp_dealloc */
1262 0, /* tp_print */
1263 0, /* tp_getattr */
1264 0, /* tp_setattr */
1265 0, /* tp_compare */
1266 0, /* tp_repr */
1267 0, /* tp_as_number */
1268 0, /* tp_as_sequence */
1269 0, /* tp_as_mapping */
1270 0, /* tp_hash */
1271 0, /* tp_call */
1272 0, /* tp_str */
1273 PyObject_GenericGetAttr, /* tp_getattro */
1274 0, /* tp_setattro */
1275 0, /* tp_as_buffer */
1276 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1277 Py_TPFLAGS_BASETYPE, /* tp_flags */
1278 islice_doc, /* tp_doc */
1279 (traverseproc)islice_traverse, /* tp_traverse */
1280 0, /* tp_clear */
1281 0, /* tp_richcompare */
1282 0, /* tp_weaklistoffset */
1283 PyObject_SelfIter, /* tp_iter */
1284 (iternextfunc)islice_next, /* tp_iternext */
1285 0, /* tp_methods */
1286 0, /* tp_members */
1287 0, /* tp_getset */
1288 0, /* tp_base */
1289 0, /* tp_dict */
1290 0, /* tp_descr_get */
1291 0, /* tp_descr_set */
1292 0, /* tp_dictoffset */
1293 0, /* tp_init */
1294 0, /* tp_alloc */
1295 islice_new, /* tp_new */
1296 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001297};
1298
1299
1300/* starmap object ************************************************************/
1301
1302typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001303 PyObject_HEAD
1304 PyObject *func;
1305 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001306} starmapobject;
1307
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001308static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001309
1310static PyObject *
1311starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1312{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001313 PyObject *func, *seq;
1314 PyObject *it;
1315 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001316
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001317 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1318 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001319
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001320 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1321 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001322
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001323 /* Get iterator. */
1324 it = PyObject_GetIter(seq);
1325 if (it == NULL)
1326 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001327
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001328 /* create starmapobject structure */
1329 lz = (starmapobject *)type->tp_alloc(type, 0);
1330 if (lz == NULL) {
1331 Py_DECREF(it);
1332 return NULL;
1333 }
1334 Py_INCREF(func);
1335 lz->func = func;
1336 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001337
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001338 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001339}
1340
1341static void
1342starmap_dealloc(starmapobject *lz)
1343{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001344 PyObject_GC_UnTrack(lz);
1345 Py_XDECREF(lz->func);
1346 Py_XDECREF(lz->it);
1347 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001348}
1349
1350static int
1351starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1352{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001353 Py_VISIT(lz->it);
1354 Py_VISIT(lz->func);
1355 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001356}
1357
1358static PyObject *
1359starmap_next(starmapobject *lz)
1360{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001361 PyObject *args;
1362 PyObject *result;
1363 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001364
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001365 args = (*Py_TYPE(it)->tp_iternext)(it);
1366 if (args == NULL)
1367 return NULL;
1368 if (!PyTuple_CheckExact(args)) {
1369 PyObject *newargs = PySequence_Tuple(args);
1370 Py_DECREF(args);
1371 if (newargs == NULL)
1372 return NULL;
1373 args = newargs;
1374 }
1375 result = PyObject_Call(lz->func, args, NULL);
1376 Py_DECREF(args);
1377 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001378}
1379
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001380PyDoc_STRVAR(starmap_doc,
1381"starmap(function, sequence) --> starmap object\n\
1382\n\
1383Return an iterator whose values are returned from the function evaluated\n\
1384with a argument tuple taken from the given sequence.");
1385
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001386static PyTypeObject starmap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001387 PyVarObject_HEAD_INIT(NULL, 0)
1388 "itertools.starmap", /* tp_name */
1389 sizeof(starmapobject), /* tp_basicsize */
1390 0, /* tp_itemsize */
1391 /* methods */
1392 (destructor)starmap_dealloc, /* tp_dealloc */
1393 0, /* tp_print */
1394 0, /* tp_getattr */
1395 0, /* tp_setattr */
1396 0, /* tp_compare */
1397 0, /* tp_repr */
1398 0, /* tp_as_number */
1399 0, /* tp_as_sequence */
1400 0, /* tp_as_mapping */
1401 0, /* tp_hash */
1402 0, /* tp_call */
1403 0, /* tp_str */
1404 PyObject_GenericGetAttr, /* tp_getattro */
1405 0, /* tp_setattro */
1406 0, /* tp_as_buffer */
1407 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1408 Py_TPFLAGS_BASETYPE, /* tp_flags */
1409 starmap_doc, /* tp_doc */
1410 (traverseproc)starmap_traverse, /* tp_traverse */
1411 0, /* tp_clear */
1412 0, /* tp_richcompare */
1413 0, /* tp_weaklistoffset */
1414 PyObject_SelfIter, /* tp_iter */
1415 (iternextfunc)starmap_next, /* tp_iternext */
1416 0, /* tp_methods */
1417 0, /* tp_members */
1418 0, /* tp_getset */
1419 0, /* tp_base */
1420 0, /* tp_dict */
1421 0, /* tp_descr_get */
1422 0, /* tp_descr_set */
1423 0, /* tp_dictoffset */
1424 0, /* tp_init */
1425 0, /* tp_alloc */
1426 starmap_new, /* tp_new */
1427 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001428};
1429
1430
1431/* imap object ************************************************************/
1432
1433typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001434 PyObject_HEAD
1435 PyObject *iters;
1436 PyObject *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001437} imapobject;
1438
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001439static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001440
1441static PyObject *
1442imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1443{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001444 PyObject *it, *iters, *func;
1445 imapobject *lz;
1446 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001447
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001448 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
1449 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001450
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001451 numargs = PyTuple_Size(args);
1452 if (numargs < 2) {
1453 PyErr_SetString(PyExc_TypeError,
1454 "imap() must have at least two arguments.");
1455 return NULL;
1456 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001458 iters = PyTuple_New(numargs-1);
1459 if (iters == NULL)
1460 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001461
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001462 for (i=1 ; i<numargs ; i++) {
1463 /* Get iterator. */
1464 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1465 if (it == NULL) {
1466 Py_DECREF(iters);
1467 return NULL;
1468 }
1469 PyTuple_SET_ITEM(iters, i-1, it);
1470 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001471
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001472 /* create imapobject structure */
1473 lz = (imapobject *)type->tp_alloc(type, 0);
1474 if (lz == NULL) {
1475 Py_DECREF(iters);
1476 return NULL;
1477 }
1478 lz->iters = iters;
1479 func = PyTuple_GET_ITEM(args, 0);
1480 Py_INCREF(func);
1481 lz->func = func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001482
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001483 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001484}
1485
1486static void
1487imap_dealloc(imapobject *lz)
1488{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001489 PyObject_GC_UnTrack(lz);
1490 Py_XDECREF(lz->iters);
1491 Py_XDECREF(lz->func);
1492 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001493}
1494
1495static int
1496imap_traverse(imapobject *lz, visitproc visit, void *arg)
1497{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001498 Py_VISIT(lz->iters);
1499 Py_VISIT(lz->func);
1500 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001501}
1502
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001503/*
Raymond Hettinger2012f172003-02-07 05:32:58 +00001504imap() is an iterator version of __builtins__.map() except that it does
1505not have the None fill-in feature. That was intentionally left out for
1506the following reasons:
1507
1508 1) Itertools are designed to be easily combined and chained together.
1509 Having all tools stop with the shortest input is a unifying principle
1510 that makes it easier to combine finite iterators (supplying data) with
1511 infinite iterators like count() and repeat() (for supplying sequential
1512 or constant arguments to a function).
1513
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001514 2) In typical use cases for combining itertools, having one finite data
Raymond Hettinger2012f172003-02-07 05:32:58 +00001515 supplier run out before another is likely to be an error condition which
1516 should not pass silently by automatically supplying None.
1517
1518 3) The use cases for automatic None fill-in are rare -- not many functions
1519 do something useful when a parameter suddenly switches type and becomes
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001520 None.
Raymond Hettinger2012f172003-02-07 05:32:58 +00001521
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001522 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001523 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001524
1525 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1526*/
1527
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001528static PyObject *
1529imap_next(imapobject *lz)
1530{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001531 PyObject *val;
1532 PyObject *argtuple;
1533 PyObject *result;
1534 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001535
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001536 numargs = PyTuple_Size(lz->iters);
1537 argtuple = PyTuple_New(numargs);
1538 if (argtuple == NULL)
1539 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001540
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001541 for (i=0 ; i<numargs ; i++) {
1542 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1543 if (val == NULL) {
1544 Py_DECREF(argtuple);
1545 return NULL;
1546 }
1547 PyTuple_SET_ITEM(argtuple, i, val);
1548 }
1549 if (lz->func == Py_None)
1550 return argtuple;
1551 result = PyObject_Call(lz->func, argtuple, NULL);
1552 Py_DECREF(argtuple);
1553 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001554}
1555
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001556PyDoc_STRVAR(imap_doc,
1557"imap(func, *iterables) --> imap object\n\
1558\n\
1559Make an iterator that computes the function using arguments from\n\
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001560each of the iterables. Like map() except that it returns\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001561an iterator instead of a list and that it stops when the shortest\n\
1562iterable is exhausted instead of filling in None for shorter\n\
1563iterables.");
1564
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001565static PyTypeObject imap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001566 PyVarObject_HEAD_INIT(NULL, 0)
1567 "itertools.imap", /* tp_name */
1568 sizeof(imapobject), /* tp_basicsize */
1569 0, /* tp_itemsize */
1570 /* methods */
1571 (destructor)imap_dealloc, /* tp_dealloc */
1572 0, /* tp_print */
1573 0, /* tp_getattr */
1574 0, /* tp_setattr */
1575 0, /* tp_compare */
1576 0, /* tp_repr */
1577 0, /* tp_as_number */
1578 0, /* tp_as_sequence */
1579 0, /* tp_as_mapping */
1580 0, /* tp_hash */
1581 0, /* tp_call */
1582 0, /* tp_str */
1583 PyObject_GenericGetAttr, /* tp_getattro */
1584 0, /* tp_setattro */
1585 0, /* tp_as_buffer */
1586 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1587 Py_TPFLAGS_BASETYPE, /* tp_flags */
1588 imap_doc, /* tp_doc */
1589 (traverseproc)imap_traverse, /* tp_traverse */
1590 0, /* tp_clear */
1591 0, /* tp_richcompare */
1592 0, /* tp_weaklistoffset */
1593 PyObject_SelfIter, /* tp_iter */
1594 (iternextfunc)imap_next, /* tp_iternext */
1595 0, /* tp_methods */
1596 0, /* tp_members */
1597 0, /* tp_getset */
1598 0, /* tp_base */
1599 0, /* tp_dict */
1600 0, /* tp_descr_get */
1601 0, /* tp_descr_set */
1602 0, /* tp_dictoffset */
1603 0, /* tp_init */
1604 0, /* tp_alloc */
1605 imap_new, /* tp_new */
1606 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001607};
1608
1609
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001610/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001611
1612typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001613 PyObject_HEAD
1614 PyObject *source; /* Iterator over input iterables */
1615 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001616} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001617
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001618static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001619
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001620static PyObject *
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001621chain_new_internal(PyTypeObject *type, PyObject *source)
1622{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001623 chainobject *lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001624
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001625 lz = (chainobject *)type->tp_alloc(type, 0);
1626 if (lz == NULL) {
1627 Py_DECREF(source);
1628 return NULL;
1629 }
1630
1631 lz->source = source;
1632 lz->active = NULL;
1633 return (PyObject *)lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001634}
1635
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001636static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001637chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001638{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001639 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001640
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001641 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1642 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001643
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001644 source = PyObject_GetIter(args);
1645 if (source == NULL)
1646 return NULL;
1647
1648 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001649}
1650
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001651static PyObject *
1652chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1653{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001654 PyObject *source;
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001655
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001656 source = PyObject_GetIter(arg);
1657 if (source == NULL)
1658 return NULL;
1659
1660 return chain_new_internal(type, source);
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001661}
1662
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001663static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001664chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001665{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001666 PyObject_GC_UnTrack(lz);
1667 Py_XDECREF(lz->active);
1668 Py_XDECREF(lz->source);
1669 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001670}
1671
Raymond Hettinger2012f172003-02-07 05:32:58 +00001672static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001673chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001674{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001675 Py_VISIT(lz->source);
1676 Py_VISIT(lz->active);
1677 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001678}
1679
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001680static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001681chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001682{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001683 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001684
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001685 if (lz->source == NULL)
1686 return NULL; /* already stopped */
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001687
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001688 if (lz->active == NULL) {
1689 PyObject *iterable = PyIter_Next(lz->source);
1690 if (iterable == NULL) {
1691 Py_CLEAR(lz->source);
1692 return NULL; /* no more input sources */
1693 }
1694 lz->active = PyObject_GetIter(iterable);
1695 Py_DECREF(iterable);
1696 if (lz->active == NULL) {
1697 Py_CLEAR(lz->source);
1698 return NULL; /* input not iterable */
1699 }
1700 }
1701 item = PyIter_Next(lz->active);
1702 if (item != NULL)
1703 return item;
1704 if (PyErr_Occurred()) {
1705 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1706 PyErr_Clear();
1707 else
1708 return NULL; /* input raised an exception */
1709 }
1710 Py_CLEAR(lz->active);
1711 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001712}
1713
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001714PyDoc_STRVAR(chain_doc,
1715"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001716\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001717Return a chain object whose .next() method returns elements from the\n\
1718first iterable until it is exhausted, then elements from the next\n\
1719iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001720
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001721PyDoc_STRVAR(chain_from_iterable_doc,
1722"chain.from_iterable(iterable) --> chain object\n\
1723\n\
1724Alternate chain() contructor taking a single iterable argument\n\
1725that evaluates lazily.");
1726
1727static PyMethodDef chain_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001728 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1729 chain_from_iterable_doc},
1730 {NULL, NULL} /* sentinel */
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001731};
1732
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001733static PyTypeObject chain_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001734 PyVarObject_HEAD_INIT(NULL, 0)
1735 "itertools.chain", /* tp_name */
1736 sizeof(chainobject), /* tp_basicsize */
1737 0, /* tp_itemsize */
1738 /* methods */
1739 (destructor)chain_dealloc, /* tp_dealloc */
1740 0, /* tp_print */
1741 0, /* tp_getattr */
1742 0, /* tp_setattr */
1743 0, /* tp_compare */
1744 0, /* tp_repr */
1745 0, /* tp_as_number */
1746 0, /* tp_as_sequence */
1747 0, /* tp_as_mapping */
1748 0, /* tp_hash */
1749 0, /* tp_call */
1750 0, /* tp_str */
1751 PyObject_GenericGetAttr, /* tp_getattro */
1752 0, /* tp_setattro */
1753 0, /* tp_as_buffer */
1754 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1755 Py_TPFLAGS_BASETYPE, /* tp_flags */
1756 chain_doc, /* tp_doc */
1757 (traverseproc)chain_traverse, /* tp_traverse */
1758 0, /* tp_clear */
1759 0, /* tp_richcompare */
1760 0, /* tp_weaklistoffset */
1761 PyObject_SelfIter, /* tp_iter */
1762 (iternextfunc)chain_next, /* tp_iternext */
1763 chain_methods, /* tp_methods */
1764 0, /* tp_members */
1765 0, /* tp_getset */
1766 0, /* tp_base */
1767 0, /* tp_dict */
1768 0, /* tp_descr_get */
1769 0, /* tp_descr_set */
1770 0, /* tp_dictoffset */
1771 0, /* tp_init */
1772 0, /* tp_alloc */
1773 chain_new, /* tp_new */
1774 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001775};
1776
1777
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001778/* product object ************************************************************/
1779
1780typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001781 PyObject_HEAD
1782 PyObject *pools; /* tuple of pool tuples */
1783 Py_ssize_t *indices; /* one index per pool */
1784 PyObject *result; /* most recently returned result tuple */
1785 int stopped; /* set to 1 when the product iterator is exhausted */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001786} productobject;
1787
1788static PyTypeObject product_type;
1789
1790static PyObject *
1791product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1792{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001793 productobject *lz;
1794 Py_ssize_t nargs, npools, repeat=1;
1795 PyObject *pools = NULL;
1796 Py_ssize_t *indices = NULL;
1797 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001798
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001799 if (kwds != NULL) {
1800 char *kwlist[] = {"repeat", 0};
1801 PyObject *tmpargs = PyTuple_New(0);
1802 if (tmpargs == NULL)
1803 return NULL;
1804 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1805 Py_DECREF(tmpargs);
1806 return NULL;
1807 }
1808 Py_DECREF(tmpargs);
1809 if (repeat < 0) {
1810 PyErr_SetString(PyExc_ValueError,
1811 "repeat argument cannot be negative");
1812 return NULL;
1813 }
1814 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001815
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001816 assert(PyTuple_Check(args));
1817 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1818 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001819
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001820 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1821 if (indices == NULL) {
1822 PyErr_NoMemory();
1823 goto error;
1824 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001825
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001826 pools = PyTuple_New(npools);
1827 if (pools == NULL)
1828 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001829
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001830 for (i=0; i < nargs ; ++i) {
1831 PyObject *item = PyTuple_GET_ITEM(args, i);
1832 PyObject *pool = PySequence_Tuple(item);
1833 if (pool == NULL)
1834 goto error;
1835 PyTuple_SET_ITEM(pools, i, pool);
1836 indices[i] = 0;
1837 }
1838 for ( ; i < npools; ++i) {
1839 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1840 Py_INCREF(pool);
1841 PyTuple_SET_ITEM(pools, i, pool);
1842 indices[i] = 0;
1843 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001844
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001845 /* create productobject structure */
1846 lz = (productobject *)type->tp_alloc(type, 0);
1847 if (lz == NULL)
1848 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001849
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001850 lz->pools = pools;
1851 lz->indices = indices;
1852 lz->result = NULL;
1853 lz->stopped = 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001854
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001855 return (PyObject *)lz;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001856
1857error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001858 if (indices != NULL)
1859 PyMem_Free(indices);
1860 Py_XDECREF(pools);
1861 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001862}
1863
1864static void
1865product_dealloc(productobject *lz)
1866{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001867 PyObject_GC_UnTrack(lz);
1868 Py_XDECREF(lz->pools);
1869 Py_XDECREF(lz->result);
1870 if (lz->indices != NULL)
1871 PyMem_Free(lz->indices);
1872 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001873}
1874
1875static int
1876product_traverse(productobject *lz, visitproc visit, void *arg)
1877{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001878 Py_VISIT(lz->pools);
1879 Py_VISIT(lz->result);
1880 return 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001881}
1882
1883static PyObject *
1884product_next(productobject *lz)
1885{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001886 PyObject *pool;
1887 PyObject *elem;
1888 PyObject *oldelem;
1889 PyObject *pools = lz->pools;
1890 PyObject *result = lz->result;
1891 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1892 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001893
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001894 if (lz->stopped)
1895 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001896
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001897 if (result == NULL) {
1898 /* On the first pass, return an initial tuple filled with the
1899 first element from each pool. */
1900 result = PyTuple_New(npools);
1901 if (result == NULL)
1902 goto empty;
1903 lz->result = result;
1904 for (i=0; i < npools; i++) {
1905 pool = PyTuple_GET_ITEM(pools, i);
1906 if (PyTuple_GET_SIZE(pool) == 0)
1907 goto empty;
1908 elem = PyTuple_GET_ITEM(pool, 0);
1909 Py_INCREF(elem);
1910 PyTuple_SET_ITEM(result, i, elem);
1911 }
1912 } else {
1913 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001914
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001915 /* Copy the previous result tuple or re-use it if available */
1916 if (Py_REFCNT(result) > 1) {
1917 PyObject *old_result = result;
1918 result = PyTuple_New(npools);
1919 if (result == NULL)
1920 goto empty;
1921 lz->result = result;
1922 for (i=0; i < npools; i++) {
1923 elem = PyTuple_GET_ITEM(old_result, i);
1924 Py_INCREF(elem);
1925 PyTuple_SET_ITEM(result, i, elem);
1926 }
1927 Py_DECREF(old_result);
1928 }
1929 /* Now, we've got the only copy so we can update it in-place */
1930 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001931
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001932 /* Update the pool indices right-to-left. Only advance to the
1933 next pool when the previous one rolls-over */
1934 for (i=npools-1 ; i >= 0 ; i--) {
1935 pool = PyTuple_GET_ITEM(pools, i);
1936 indices[i]++;
1937 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1938 /* Roll-over and advance to next pool */
1939 indices[i] = 0;
1940 elem = PyTuple_GET_ITEM(pool, 0);
1941 Py_INCREF(elem);
1942 oldelem = PyTuple_GET_ITEM(result, i);
1943 PyTuple_SET_ITEM(result, i, elem);
1944 Py_DECREF(oldelem);
1945 } else {
1946 /* No rollover. Just increment and stop here. */
1947 elem = PyTuple_GET_ITEM(pool, indices[i]);
1948 Py_INCREF(elem);
1949 oldelem = PyTuple_GET_ITEM(result, i);
1950 PyTuple_SET_ITEM(result, i, elem);
1951 Py_DECREF(oldelem);
1952 break;
1953 }
1954 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00001955
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001956 /* If i is negative, then the indices have all rolled-over
1957 and we're done. */
1958 if (i < 0)
1959 goto empty;
1960 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001961
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001962 Py_INCREF(result);
1963 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001964
1965empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001966 lz->stopped = 1;
1967 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001968}
1969
1970PyDoc_STRVAR(product_doc,
1971"product(*iterables) --> product object\n\
1972\n\
Raymond Hettinger73d79632008-02-23 02:20:41 +00001973Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001974For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1975The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1976cycle in a manner similar to an odometer (with the rightmost element changing\n\
1977on every iteration).\n\n\
Raymond Hettinger94034ea2009-02-09 18:39:41 +00001978To compute the product of an iterable with itself, specify the number\n\
1979of repetitions with the optional repeat keyword argument. For example,\n\
1980product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001981product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1982product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1983
1984static PyTypeObject product_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001985 PyVarObject_HEAD_INIT(NULL, 0)
1986 "itertools.product", /* tp_name */
1987 sizeof(productobject), /* tp_basicsize */
1988 0, /* tp_itemsize */
1989 /* methods */
1990 (destructor)product_dealloc, /* tp_dealloc */
1991 0, /* tp_print */
1992 0, /* tp_getattr */
1993 0, /* tp_setattr */
1994 0, /* tp_compare */
1995 0, /* tp_repr */
1996 0, /* tp_as_number */
1997 0, /* tp_as_sequence */
1998 0, /* tp_as_mapping */
1999 0, /* tp_hash */
2000 0, /* tp_call */
2001 0, /* tp_str */
2002 PyObject_GenericGetAttr, /* tp_getattro */
2003 0, /* tp_setattro */
2004 0, /* tp_as_buffer */
2005 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2006 Py_TPFLAGS_BASETYPE, /* tp_flags */
2007 product_doc, /* tp_doc */
2008 (traverseproc)product_traverse, /* tp_traverse */
2009 0, /* tp_clear */
2010 0, /* tp_richcompare */
2011 0, /* tp_weaklistoffset */
2012 PyObject_SelfIter, /* tp_iter */
2013 (iternextfunc)product_next, /* tp_iternext */
2014 0, /* tp_methods */
2015 0, /* tp_members */
2016 0, /* tp_getset */
2017 0, /* tp_base */
2018 0, /* tp_dict */
2019 0, /* tp_descr_get */
2020 0, /* tp_descr_set */
2021 0, /* tp_dictoffset */
2022 0, /* tp_init */
2023 0, /* tp_alloc */
2024 product_new, /* tp_new */
2025 PyObject_GC_Del, /* tp_free */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002026};
2027
2028
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002029/* combinations object ************************************************************/
2030
2031typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002032 PyObject_HEAD
2033 PyObject *pool; /* input converted to a tuple */
2034 Py_ssize_t *indices; /* one index per result element */
2035 PyObject *result; /* most recently returned result tuple */
2036 Py_ssize_t r; /* size of result tuple */
2037 int stopped; /* set to 1 when the combinations iterator is exhausted */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002038} combinationsobject;
2039
2040static PyTypeObject combinations_type;
2041
2042static PyObject *
2043combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2044{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002045 combinationsobject *co;
2046 Py_ssize_t n;
2047 Py_ssize_t r;
2048 PyObject *pool = NULL;
2049 PyObject *iterable = NULL;
2050 Py_ssize_t *indices = NULL;
2051 Py_ssize_t i;
2052 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002053
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002054 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2055 &iterable, &r))
2056 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002057
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002058 pool = PySequence_Tuple(iterable);
2059 if (pool == NULL)
2060 goto error;
2061 n = PyTuple_GET_SIZE(pool);
2062 if (r < 0) {
2063 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2064 goto error;
2065 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002066
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002067 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2068 if (indices == NULL) {
2069 PyErr_NoMemory();
2070 goto error;
2071 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002072
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002073 for (i=0 ; i<r ; i++)
2074 indices[i] = i;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002075
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002076 /* create combinationsobject structure */
2077 co = (combinationsobject *)type->tp_alloc(type, 0);
2078 if (co == NULL)
2079 goto error;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002080
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002081 co->pool = pool;
2082 co->indices = indices;
2083 co->result = NULL;
2084 co->r = r;
2085 co->stopped = r > n ? 1 : 0;
2086
2087 return (PyObject *)co;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002088
2089error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002090 if (indices != NULL)
2091 PyMem_Free(indices);
2092 Py_XDECREF(pool);
2093 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002094}
2095
2096static void
2097combinations_dealloc(combinationsobject *co)
2098{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002099 PyObject_GC_UnTrack(co);
2100 Py_XDECREF(co->pool);
2101 Py_XDECREF(co->result);
2102 if (co->indices != NULL)
2103 PyMem_Free(co->indices);
2104 Py_TYPE(co)->tp_free(co);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002105}
2106
2107static int
2108combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2109{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002110 Py_VISIT(co->pool);
2111 Py_VISIT(co->result);
2112 return 0;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002113}
2114
2115static PyObject *
2116combinations_next(combinationsobject *co)
2117{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002118 PyObject *elem;
2119 PyObject *oldelem;
2120 PyObject *pool = co->pool;
2121 Py_ssize_t *indices = co->indices;
2122 PyObject *result = co->result;
2123 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2124 Py_ssize_t r = co->r;
2125 Py_ssize_t i, j, index;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002126
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002127 if (co->stopped)
2128 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002129
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002130 if (result == NULL) {
2131 /* On the first pass, initialize result tuple using the indices */
2132 result = PyTuple_New(r);
2133 if (result == NULL)
2134 goto empty;
2135 co->result = result;
2136 for (i=0; i<r ; i++) {
2137 index = indices[i];
2138 elem = PyTuple_GET_ITEM(pool, index);
2139 Py_INCREF(elem);
2140 PyTuple_SET_ITEM(result, i, elem);
2141 }
2142 } else {
2143 /* Copy the previous result tuple or re-use it if available */
2144 if (Py_REFCNT(result) > 1) {
2145 PyObject *old_result = result;
2146 result = PyTuple_New(r);
2147 if (result == NULL)
2148 goto empty;
2149 co->result = result;
2150 for (i=0; i<r ; i++) {
2151 elem = PyTuple_GET_ITEM(old_result, i);
2152 Py_INCREF(elem);
2153 PyTuple_SET_ITEM(result, i, elem);
2154 }
2155 Py_DECREF(old_result);
2156 }
2157 /* Now, we've got the only copy so we can update it in-place
2158 * CPython's empty tuple is a singleton and cached in
2159 * PyTuple's freelist.
2160 */
2161 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002162
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002163 /* Scan indices right-to-left until finding one that is not
2164 at its maximum (i + n - r). */
2165 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2166 ;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002167
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002168 /* If i is negative, then the indices are all at
2169 their maximum value and we're done. */
2170 if (i < 0)
2171 goto empty;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002172
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002173 /* Increment the current index which we know is not at its
2174 maximum. Then move back to the right setting each index
2175 to its lowest possible value (one higher than the index
2176 to its left -- this maintains the sort order invariant). */
2177 indices[i]++;
2178 for (j=i+1 ; j<r ; j++)
2179 indices[j] = indices[j-1] + 1;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002180
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002181 /* Update the result tuple for the new indices
2182 starting with i, the leftmost index that changed */
2183 for ( ; i<r ; i++) {
2184 index = indices[i];
2185 elem = PyTuple_GET_ITEM(pool, index);
2186 Py_INCREF(elem);
2187 oldelem = PyTuple_GET_ITEM(result, i);
2188 PyTuple_SET_ITEM(result, i, elem);
2189 Py_DECREF(oldelem);
2190 }
2191 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002192
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002193 Py_INCREF(result);
2194 return result;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002195
2196empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002197 co->stopped = 1;
2198 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002199}
2200
2201PyDoc_STRVAR(combinations_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002202"combinations(iterable, r) --> combinations object\n\
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002203\n\
2204Return successive r-length combinations of elements in the iterable.\n\n\
2205combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2206
2207static PyTypeObject combinations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002208 PyVarObject_HEAD_INIT(NULL, 0)
2209 "itertools.combinations", /* tp_name */
2210 sizeof(combinationsobject), /* tp_basicsize */
2211 0, /* tp_itemsize */
2212 /* methods */
2213 (destructor)combinations_dealloc, /* tp_dealloc */
2214 0, /* tp_print */
2215 0, /* tp_getattr */
2216 0, /* tp_setattr */
2217 0, /* tp_compare */
2218 0, /* tp_repr */
2219 0, /* tp_as_number */
2220 0, /* tp_as_sequence */
2221 0, /* tp_as_mapping */
2222 0, /* tp_hash */
2223 0, /* tp_call */
2224 0, /* tp_str */
2225 PyObject_GenericGetAttr, /* tp_getattro */
2226 0, /* tp_setattro */
2227 0, /* tp_as_buffer */
2228 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2229 Py_TPFLAGS_BASETYPE, /* tp_flags */
2230 combinations_doc, /* tp_doc */
2231 (traverseproc)combinations_traverse, /* tp_traverse */
2232 0, /* tp_clear */
2233 0, /* tp_richcompare */
2234 0, /* tp_weaklistoffset */
2235 PyObject_SelfIter, /* tp_iter */
2236 (iternextfunc)combinations_next, /* tp_iternext */
2237 0, /* tp_methods */
2238 0, /* tp_members */
2239 0, /* tp_getset */
2240 0, /* tp_base */
2241 0, /* tp_dict */
2242 0, /* tp_descr_get */
2243 0, /* tp_descr_set */
2244 0, /* tp_dictoffset */
2245 0, /* tp_init */
2246 0, /* tp_alloc */
2247 combinations_new, /* tp_new */
2248 PyObject_GC_Del, /* tp_free */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002249};
2250
2251
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002252/* combinations with replacement object *******************************************/
2253
2254/* Equivalent to:
2255
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002256 def combinations_with_replacement(iterable, r):
2257 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2258 # number items returned: (n+r-1)! / r! / (n-1)!
2259 pool = tuple(iterable)
2260 n = len(pool)
2261 indices = [0] * r
2262 yield tuple(pool[i] for i in indices)
2263 while 1:
2264 for i in reversed(range(r)):
2265 if indices[i] != n - 1:
2266 break
2267 else:
2268 return
2269 indices[i:] = [indices[i] + 1] * (r - i)
2270 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002271
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002272 def combinations_with_replacement2(iterable, r):
2273 'Alternate version that filters from product()'
2274 pool = tuple(iterable)
2275 n = len(pool)
2276 for indices in product(range(n), repeat=r):
2277 if sorted(indices) == list(indices):
2278 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002279*/
2280typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002281 PyObject_HEAD
2282 PyObject *pool; /* input converted to a tuple */
2283 Py_ssize_t *indices; /* one index per result element */
2284 PyObject *result; /* most recently returned result tuple */
2285 Py_ssize_t r; /* size of result tuple */
2286 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002287} cwrobject;
2288
2289static PyTypeObject cwr_type;
2290
2291static PyObject *
2292cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2293{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002294 cwrobject *co;
2295 Py_ssize_t n;
2296 Py_ssize_t r;
2297 PyObject *pool = NULL;
2298 PyObject *iterable = NULL;
2299 Py_ssize_t *indices = NULL;
2300 Py_ssize_t i;
2301 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002302
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002303 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2304 &iterable, &r))
2305 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002306
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002307 pool = PySequence_Tuple(iterable);
2308 if (pool == NULL)
2309 goto error;
2310 n = PyTuple_GET_SIZE(pool);
2311 if (r < 0) {
2312 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2313 goto error;
2314 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002315
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002316 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2317 if (indices == NULL) {
2318 PyErr_NoMemory();
2319 goto error;
2320 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002321
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002322 for (i=0 ; i<r ; i++)
2323 indices[i] = 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002324
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002325 /* create cwrobject structure */
2326 co = (cwrobject *)type->tp_alloc(type, 0);
2327 if (co == NULL)
2328 goto error;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002329
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002330 co->pool = pool;
2331 co->indices = indices;
2332 co->result = NULL;
2333 co->r = r;
2334 co->stopped = !n && r;
2335
2336 return (PyObject *)co;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002337
2338error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002339 if (indices != NULL)
2340 PyMem_Free(indices);
2341 Py_XDECREF(pool);
2342 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002343}
2344
2345static void
2346cwr_dealloc(cwrobject *co)
2347{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002348 PyObject_GC_UnTrack(co);
2349 Py_XDECREF(co->pool);
2350 Py_XDECREF(co->result);
2351 if (co->indices != NULL)
2352 PyMem_Free(co->indices);
2353 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002354}
2355
2356static int
2357cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2358{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002359 Py_VISIT(co->pool);
2360 Py_VISIT(co->result);
2361 return 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002362}
2363
2364static PyObject *
2365cwr_next(cwrobject *co)
2366{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002367 PyObject *elem;
2368 PyObject *oldelem;
2369 PyObject *pool = co->pool;
2370 Py_ssize_t *indices = co->indices;
2371 PyObject *result = co->result;
2372 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2373 Py_ssize_t r = co->r;
2374 Py_ssize_t i, j, index;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002375
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002376 if (co->stopped)
2377 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002378
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002379 if (result == NULL) {
2380 /* On the first pass, initialize result tuple using the indices */
2381 result = PyTuple_New(r);
2382 if (result == NULL)
2383 goto empty;
2384 co->result = result;
2385 for (i=0; i<r ; i++) {
2386 index = indices[i];
2387 elem = PyTuple_GET_ITEM(pool, index);
2388 Py_INCREF(elem);
2389 PyTuple_SET_ITEM(result, i, elem);
2390 }
2391 } else {
2392 /* Copy the previous result tuple or re-use it if available */
2393 if (Py_REFCNT(result) > 1) {
2394 PyObject *old_result = result;
2395 result = PyTuple_New(r);
2396 if (result == NULL)
2397 goto empty;
2398 co->result = result;
2399 for (i=0; i<r ; i++) {
2400 elem = PyTuple_GET_ITEM(old_result, i);
2401 Py_INCREF(elem);
2402 PyTuple_SET_ITEM(result, i, elem);
2403 }
2404 Py_DECREF(old_result);
2405 }
2406 /* Now, we've got the only copy so we can update it in-place CPython's
2407 empty tuple is a singleton and cached in PyTuple's freelist. */
2408 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002409
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002410 /* Scan indices right-to-left until finding one that is not
2411 * at its maximum (n-1). */
2412 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2413 ;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002414
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002415 /* If i is negative, then the indices are all at
2416 their maximum value and we're done. */
2417 if (i < 0)
2418 goto empty;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002419
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002420 /* Increment the current index which we know is not at its
2421 maximum. Then set all to the right to the same value. */
2422 indices[i]++;
2423 for (j=i+1 ; j<r ; j++)
2424 indices[j] = indices[j-1];
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002425
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002426 /* Update the result tuple for the new indices
2427 starting with i, the leftmost index that changed */
2428 for ( ; i<r ; i++) {
2429 index = indices[i];
2430 elem = PyTuple_GET_ITEM(pool, index);
2431 Py_INCREF(elem);
2432 oldelem = PyTuple_GET_ITEM(result, i);
2433 PyTuple_SET_ITEM(result, i, elem);
2434 Py_DECREF(oldelem);
2435 }
2436 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002437
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002438 Py_INCREF(result);
2439 return result;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002440
2441empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002442 co->stopped = 1;
2443 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002444}
2445
2446PyDoc_STRVAR(cwr_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002447"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002448\n\
2449Return successive r-length combinations of elements in the iterable\n\
2450allowing individual elements to have successive repeats.\n\
2451combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2452
2453static PyTypeObject cwr_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002454 PyVarObject_HEAD_INIT(NULL, 0)
2455 "itertools.combinations_with_replacement", /* tp_name */
2456 sizeof(cwrobject), /* tp_basicsize */
2457 0, /* tp_itemsize */
2458 /* methods */
2459 (destructor)cwr_dealloc, /* tp_dealloc */
2460 0, /* tp_print */
2461 0, /* tp_getattr */
2462 0, /* tp_setattr */
2463 0, /* tp_compare */
2464 0, /* tp_repr */
2465 0, /* tp_as_number */
2466 0, /* tp_as_sequence */
2467 0, /* tp_as_mapping */
2468 0, /* tp_hash */
2469 0, /* tp_call */
2470 0, /* tp_str */
2471 PyObject_GenericGetAttr, /* tp_getattro */
2472 0, /* tp_setattro */
2473 0, /* tp_as_buffer */
2474 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2475 Py_TPFLAGS_BASETYPE, /* tp_flags */
2476 cwr_doc, /* tp_doc */
2477 (traverseproc)cwr_traverse, /* tp_traverse */
2478 0, /* tp_clear */
2479 0, /* tp_richcompare */
2480 0, /* tp_weaklistoffset */
2481 PyObject_SelfIter, /* tp_iter */
2482 (iternextfunc)cwr_next, /* tp_iternext */
2483 0, /* tp_methods */
2484 0, /* tp_members */
2485 0, /* tp_getset */
2486 0, /* tp_base */
2487 0, /* tp_dict */
2488 0, /* tp_descr_get */
2489 0, /* tp_descr_set */
2490 0, /* tp_dictoffset */
2491 0, /* tp_init */
2492 0, /* tp_alloc */
2493 cwr_new, /* tp_new */
2494 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002495};
2496
2497
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002498/* permutations object ************************************************************
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002499
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002500def permutations(iterable, r=None):
2501 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2502 pool = tuple(iterable)
2503 n = len(pool)
2504 r = n if r is None else r
2505 indices = range(n)
2506 cycles = range(n-r+1, n+1)[::-1]
2507 yield tuple(pool[i] for i in indices[:r])
2508 while n:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002509 for i in reversed(range(r)):
2510 cycles[i] -= 1
2511 if cycles[i] == 0:
2512 indices[i:] = indices[i+1:] + indices[i:i+1]
2513 cycles[i] = n - i
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002514 else:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002515 j = cycles[i]
2516 indices[i], indices[-j] = indices[-j], indices[i]
2517 yield tuple(pool[i] for i in indices[:r])
2518 break
2519 else:
2520 return
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002521*/
2522
2523typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002524 PyObject_HEAD
2525 PyObject *pool; /* input converted to a tuple */
2526 Py_ssize_t *indices; /* one index per element in the pool */
2527 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2528 PyObject *result; /* most recently returned result tuple */
2529 Py_ssize_t r; /* size of result tuple */
2530 int stopped; /* set to 1 when the permutations iterator is exhausted */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002531} permutationsobject;
2532
2533static PyTypeObject permutations_type;
2534
2535static PyObject *
2536permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2537{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002538 permutationsobject *po;
2539 Py_ssize_t n;
2540 Py_ssize_t r;
2541 PyObject *robj = Py_None;
2542 PyObject *pool = NULL;
2543 PyObject *iterable = NULL;
2544 Py_ssize_t *indices = NULL;
2545 Py_ssize_t *cycles = NULL;
2546 Py_ssize_t i;
2547 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002548
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002549 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2550 &iterable, &robj))
2551 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002552
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002553 pool = PySequence_Tuple(iterable);
2554 if (pool == NULL)
2555 goto error;
2556 n = PyTuple_GET_SIZE(pool);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002557
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002558 r = n;
2559 if (robj != Py_None) {
2560 r = PyInt_AsSsize_t(robj);
2561 if (r == -1 && PyErr_Occurred())
2562 goto error;
2563 }
2564 if (r < 0) {
2565 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2566 goto error;
2567 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002568
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002569 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2570 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2571 if (indices == NULL || cycles == NULL) {
2572 PyErr_NoMemory();
2573 goto error;
2574 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002575
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002576 for (i=0 ; i<n ; i++)
2577 indices[i] = i;
2578 for (i=0 ; i<r ; i++)
2579 cycles[i] = n - i;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002580
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002581 /* create permutationsobject structure */
2582 po = (permutationsobject *)type->tp_alloc(type, 0);
2583 if (po == NULL)
2584 goto error;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002585
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002586 po->pool = pool;
2587 po->indices = indices;
2588 po->cycles = cycles;
2589 po->result = NULL;
2590 po->r = r;
2591 po->stopped = r > n ? 1 : 0;
2592
2593 return (PyObject *)po;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002594
2595error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002596 if (indices != NULL)
2597 PyMem_Free(indices);
2598 if (cycles != NULL)
2599 PyMem_Free(cycles);
2600 Py_XDECREF(pool);
2601 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002602}
2603
2604static void
2605permutations_dealloc(permutationsobject *po)
2606{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002607 PyObject_GC_UnTrack(po);
2608 Py_XDECREF(po->pool);
2609 Py_XDECREF(po->result);
2610 PyMem_Free(po->indices);
2611 PyMem_Free(po->cycles);
2612 Py_TYPE(po)->tp_free(po);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002613}
2614
2615static int
2616permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2617{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002618 Py_VISIT(po->pool);
2619 Py_VISIT(po->result);
2620 return 0;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002621}
2622
2623static PyObject *
2624permutations_next(permutationsobject *po)
2625{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002626 PyObject *elem;
2627 PyObject *oldelem;
2628 PyObject *pool = po->pool;
2629 Py_ssize_t *indices = po->indices;
2630 Py_ssize_t *cycles = po->cycles;
2631 PyObject *result = po->result;
2632 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2633 Py_ssize_t r = po->r;
2634 Py_ssize_t i, j, k, index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002635
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002636 if (po->stopped)
2637 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002638
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002639 if (result == NULL) {
2640 /* On the first pass, initialize result tuple using the indices */
2641 result = PyTuple_New(r);
2642 if (result == NULL)
2643 goto empty;
2644 po->result = result;
2645 for (i=0; i<r ; i++) {
2646 index = indices[i];
2647 elem = PyTuple_GET_ITEM(pool, index);
2648 Py_INCREF(elem);
2649 PyTuple_SET_ITEM(result, i, elem);
2650 }
2651 } else {
2652 if (n == 0)
2653 goto empty;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002654
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002655 /* Copy the previous result tuple or re-use it if available */
2656 if (Py_REFCNT(result) > 1) {
2657 PyObject *old_result = result;
2658 result = PyTuple_New(r);
2659 if (result == NULL)
2660 goto empty;
2661 po->result = result;
2662 for (i=0; i<r ; i++) {
2663 elem = PyTuple_GET_ITEM(old_result, i);
2664 Py_INCREF(elem);
2665 PyTuple_SET_ITEM(result, i, elem);
2666 }
2667 Py_DECREF(old_result);
2668 }
2669 /* Now, we've got the only copy so we can update it in-place */
2670 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002671
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002672 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2673 for (i=r-1 ; i>=0 ; i--) {
2674 cycles[i] -= 1;
2675 if (cycles[i] == 0) {
2676 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2677 index = indices[i];
2678 for (j=i ; j<n-1 ; j++)
2679 indices[j] = indices[j+1];
2680 indices[n-1] = index;
2681 cycles[i] = n - i;
2682 } else {
2683 j = cycles[i];
2684 index = indices[i];
2685 indices[i] = indices[n-j];
2686 indices[n-j] = index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002687
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002688 for (k=i; k<r ; k++) {
2689 /* start with i, the leftmost element that changed */
2690 /* yield tuple(pool[k] for k in indices[:r]) */
2691 index = indices[k];
2692 elem = PyTuple_GET_ITEM(pool, index);
2693 Py_INCREF(elem);
2694 oldelem = PyTuple_GET_ITEM(result, k);
2695 PyTuple_SET_ITEM(result, k, elem);
2696 Py_DECREF(oldelem);
2697 }
2698 break;
2699 }
2700 }
2701 /* If i is negative, then the cycles have all
2702 rolled-over and we're done. */
2703 if (i < 0)
2704 goto empty;
2705 }
2706 Py_INCREF(result);
2707 return result;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002708
2709empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002710 po->stopped = 1;
2711 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002712}
2713
2714PyDoc_STRVAR(permutations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +00002715"permutations(iterable[, r]) --> permutations object\n\
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002716\n\
2717Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl8550d692008-06-10 12:46:39 +00002718permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002719
2720static PyTypeObject permutations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002721 PyVarObject_HEAD_INIT(NULL, 0)
2722 "itertools.permutations", /* tp_name */
2723 sizeof(permutationsobject), /* tp_basicsize */
2724 0, /* tp_itemsize */
2725 /* methods */
2726 (destructor)permutations_dealloc, /* tp_dealloc */
2727 0, /* tp_print */
2728 0, /* tp_getattr */
2729 0, /* tp_setattr */
2730 0, /* tp_compare */
2731 0, /* tp_repr */
2732 0, /* tp_as_number */
2733 0, /* tp_as_sequence */
2734 0, /* tp_as_mapping */
2735 0, /* tp_hash */
2736 0, /* tp_call */
2737 0, /* tp_str */
2738 PyObject_GenericGetAttr, /* tp_getattro */
2739 0, /* tp_setattro */
2740 0, /* tp_as_buffer */
2741 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2742 Py_TPFLAGS_BASETYPE, /* tp_flags */
2743 permutations_doc, /* tp_doc */
2744 (traverseproc)permutations_traverse, /* tp_traverse */
2745 0, /* tp_clear */
2746 0, /* tp_richcompare */
2747 0, /* tp_weaklistoffset */
2748 PyObject_SelfIter, /* tp_iter */
2749 (iternextfunc)permutations_next, /* tp_iternext */
2750 0, /* tp_methods */
2751 0, /* tp_members */
2752 0, /* tp_getset */
2753 0, /* tp_base */
2754 0, /* tp_dict */
2755 0, /* tp_descr_get */
2756 0, /* tp_descr_set */
2757 0, /* tp_dictoffset */
2758 0, /* tp_init */
2759 0, /* tp_alloc */
2760 permutations_new, /* tp_new */
2761 PyObject_GC_Del, /* tp_free */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002762};
2763
2764
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002765/* compress object ************************************************************/
2766
2767/* Equivalent to:
2768
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002769 def compress(data, selectors):
2770 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2771 return (d for d, s in izip(data, selectors) if s)
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002772*/
2773
2774typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002775 PyObject_HEAD
2776 PyObject *data;
2777 PyObject *selectors;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002778} compressobject;
2779
2780static PyTypeObject compress_type;
2781
2782static PyObject *
2783compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2784{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002785 PyObject *seq1, *seq2;
2786 PyObject *data=NULL, *selectors=NULL;
2787 compressobject *lz;
2788 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002789
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002790 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2791 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002792
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002793 data = PyObject_GetIter(seq1);
2794 if (data == NULL)
2795 goto fail;
2796 selectors = PyObject_GetIter(seq2);
2797 if (selectors == NULL)
2798 goto fail;
2799
2800 /* create compressobject structure */
2801 lz = (compressobject *)type->tp_alloc(type, 0);
2802 if (lz == NULL)
2803 goto fail;
2804 lz->data = data;
2805 lz->selectors = selectors;
2806 return (PyObject *)lz;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002807
2808fail:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002809 Py_XDECREF(data);
2810 Py_XDECREF(selectors);
2811 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002812}
2813
2814static void
2815compress_dealloc(compressobject *lz)
2816{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002817 PyObject_GC_UnTrack(lz);
2818 Py_XDECREF(lz->data);
2819 Py_XDECREF(lz->selectors);
2820 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002821}
2822
2823static int
2824compress_traverse(compressobject *lz, visitproc visit, void *arg)
2825{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002826 Py_VISIT(lz->data);
2827 Py_VISIT(lz->selectors);
2828 return 0;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002829}
2830
2831static PyObject *
2832compress_next(compressobject *lz)
2833{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002834 PyObject *data = lz->data, *selectors = lz->selectors;
2835 PyObject *datum, *selector;
2836 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2837 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2838 int ok;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002839
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002840 while (1) {
2841 /* Steps: get datum, get selector, evaluate selector.
2842 Order is important (to match the pure python version
2843 in terms of which input gets a chance to raise an
2844 exception first).
2845 */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002846
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002847 datum = datanext(data);
2848 if (datum == NULL)
2849 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002850
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002851 selector = selectornext(selectors);
2852 if (selector == NULL) {
2853 Py_DECREF(datum);
2854 return NULL;
2855 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002856
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002857 ok = PyObject_IsTrue(selector);
2858 Py_DECREF(selector);
2859 if (ok == 1)
2860 return datum;
2861 Py_DECREF(datum);
2862 if (ok == -1)
2863 return NULL;
2864 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002865}
2866
2867PyDoc_STRVAR(compress_doc,
Raymond Hettinger2e2909f2009-02-19 02:15:14 +00002868"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002869\n\
2870Return data elements corresponding to true selector elements.\n\
2871Forms a shorter iterator from selected data elements using the\n\
2872selectors to choose the data elements.");
2873
2874static PyTypeObject compress_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002875 PyVarObject_HEAD_INIT(NULL, 0)
2876 "itertools.compress", /* tp_name */
2877 sizeof(compressobject), /* tp_basicsize */
2878 0, /* tp_itemsize */
2879 /* methods */
2880 (destructor)compress_dealloc, /* tp_dealloc */
2881 0, /* tp_print */
2882 0, /* tp_getattr */
2883 0, /* tp_setattr */
2884 0, /* tp_compare */
2885 0, /* tp_repr */
2886 0, /* tp_as_number */
2887 0, /* tp_as_sequence */
2888 0, /* tp_as_mapping */
2889 0, /* tp_hash */
2890 0, /* tp_call */
2891 0, /* tp_str */
2892 PyObject_GenericGetAttr, /* tp_getattro */
2893 0, /* tp_setattro */
2894 0, /* tp_as_buffer */
2895 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2896 Py_TPFLAGS_BASETYPE, /* tp_flags */
2897 compress_doc, /* tp_doc */
2898 (traverseproc)compress_traverse, /* tp_traverse */
2899 0, /* tp_clear */
2900 0, /* tp_richcompare */
2901 0, /* tp_weaklistoffset */
2902 PyObject_SelfIter, /* tp_iter */
2903 (iternextfunc)compress_next, /* tp_iternext */
2904 0, /* tp_methods */
2905 0, /* tp_members */
2906 0, /* tp_getset */
2907 0, /* tp_base */
2908 0, /* tp_dict */
2909 0, /* tp_descr_get */
2910 0, /* tp_descr_set */
2911 0, /* tp_dictoffset */
2912 0, /* tp_init */
2913 0, /* tp_alloc */
2914 compress_new, /* tp_new */
2915 PyObject_GC_Del, /* tp_free */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002916};
2917
2918
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002919/* ifilter object ************************************************************/
2920
2921typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002922 PyObject_HEAD
2923 PyObject *func;
2924 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002925} ifilterobject;
2926
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002927static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002928
2929static PyObject *
2930ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2931{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002932 PyObject *func, *seq;
2933 PyObject *it;
2934 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002935
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002936 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
2937 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002938
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002939 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
2940 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002941
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002942 /* Get iterator. */
2943 it = PyObject_GetIter(seq);
2944 if (it == NULL)
2945 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002946
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002947 /* create ifilterobject structure */
2948 lz = (ifilterobject *)type->tp_alloc(type, 0);
2949 if (lz == NULL) {
2950 Py_DECREF(it);
2951 return NULL;
2952 }
2953 Py_INCREF(func);
2954 lz->func = func;
2955 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002956
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002957 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002958}
2959
2960static void
2961ifilter_dealloc(ifilterobject *lz)
2962{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002963 PyObject_GC_UnTrack(lz);
2964 Py_XDECREF(lz->func);
2965 Py_XDECREF(lz->it);
2966 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002967}
2968
2969static int
2970ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
2971{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002972 Py_VISIT(lz->it);
2973 Py_VISIT(lz->func);
2974 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002975}
2976
2977static PyObject *
2978ifilter_next(ifilterobject *lz)
2979{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002980 PyObject *item;
2981 PyObject *it = lz->it;
2982 long ok;
2983 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002984
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002985 iternext = *Py_TYPE(it)->tp_iternext;
2986 for (;;) {
2987 item = iternext(it);
2988 if (item == NULL)
2989 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002990
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002991 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2992 ok = PyObject_IsTrue(item);
2993 } else {
2994 PyObject *good;
2995 good = PyObject_CallFunctionObjArgs(lz->func,
2996 item, NULL);
2997 if (good == NULL) {
2998 Py_DECREF(item);
2999 return NULL;
3000 }
3001 ok = PyObject_IsTrue(good);
3002 Py_DECREF(good);
3003 }
3004 if (ok)
3005 return item;
3006 Py_DECREF(item);
3007 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003008}
3009
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003010PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003011"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003012\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003013Return those items of sequence for which function(item) is true.\n\
3014If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003015
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003016static PyTypeObject ifilter_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003017 PyVarObject_HEAD_INIT(NULL, 0)
3018 "itertools.ifilter", /* tp_name */
3019 sizeof(ifilterobject), /* tp_basicsize */
3020 0, /* tp_itemsize */
3021 /* methods */
3022 (destructor)ifilter_dealloc, /* tp_dealloc */
3023 0, /* tp_print */
3024 0, /* tp_getattr */
3025 0, /* tp_setattr */
3026 0, /* tp_compare */
3027 0, /* tp_repr */
3028 0, /* tp_as_number */
3029 0, /* tp_as_sequence */
3030 0, /* tp_as_mapping */
3031 0, /* tp_hash */
3032 0, /* tp_call */
3033 0, /* tp_str */
3034 PyObject_GenericGetAttr, /* tp_getattro */
3035 0, /* tp_setattro */
3036 0, /* tp_as_buffer */
3037 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3038 Py_TPFLAGS_BASETYPE, /* tp_flags */
3039 ifilter_doc, /* tp_doc */
3040 (traverseproc)ifilter_traverse, /* tp_traverse */
3041 0, /* tp_clear */
3042 0, /* tp_richcompare */
3043 0, /* tp_weaklistoffset */
3044 PyObject_SelfIter, /* tp_iter */
3045 (iternextfunc)ifilter_next, /* tp_iternext */
3046 0, /* tp_methods */
3047 0, /* tp_members */
3048 0, /* tp_getset */
3049 0, /* tp_base */
3050 0, /* tp_dict */
3051 0, /* tp_descr_get */
3052 0, /* tp_descr_set */
3053 0, /* tp_dictoffset */
3054 0, /* tp_init */
3055 0, /* tp_alloc */
3056 ifilter_new, /* tp_new */
3057 PyObject_GC_Del, /* tp_free */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003058};
3059
3060
3061/* ifilterfalse object ************************************************************/
3062
3063typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003064 PyObject_HEAD
3065 PyObject *func;
3066 PyObject *it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003067} ifilterfalseobject;
3068
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003069static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003070
3071static PyObject *
3072ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3073{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003074 PyObject *func, *seq;
3075 PyObject *it;
3076 ifilterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003077
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003078 if (type == &ifilterfalse_type &&
3079 !_PyArg_NoKeywords("ifilterfalse()", kwds))
3080 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003081
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003082 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
3083 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003084
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003085 /* Get iterator. */
3086 it = PyObject_GetIter(seq);
3087 if (it == NULL)
3088 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003089
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003090 /* create ifilterfalseobject structure */
3091 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
3092 if (lz == NULL) {
3093 Py_DECREF(it);
3094 return NULL;
3095 }
3096 Py_INCREF(func);
3097 lz->func = func;
3098 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003099
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003100 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003101}
3102
3103static void
3104ifilterfalse_dealloc(ifilterfalseobject *lz)
3105{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003106 PyObject_GC_UnTrack(lz);
3107 Py_XDECREF(lz->func);
3108 Py_XDECREF(lz->it);
3109 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003110}
3111
3112static int
3113ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
3114{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003115 Py_VISIT(lz->it);
3116 Py_VISIT(lz->func);
3117 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003118}
3119
3120static PyObject *
3121ifilterfalse_next(ifilterfalseobject *lz)
3122{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003123 PyObject *item;
3124 PyObject *it = lz->it;
3125 long ok;
3126 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003127
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003128 iternext = *Py_TYPE(it)->tp_iternext;
3129 for (;;) {
3130 item = iternext(it);
3131 if (item == NULL)
3132 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003133
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003134 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3135 ok = PyObject_IsTrue(item);
3136 } else {
3137 PyObject *good;
3138 good = PyObject_CallFunctionObjArgs(lz->func,
3139 item, NULL);
3140 if (good == NULL) {
3141 Py_DECREF(item);
3142 return NULL;
3143 }
3144 ok = PyObject_IsTrue(good);
3145 Py_DECREF(good);
3146 }
3147 if (!ok)
3148 return item;
3149 Py_DECREF(item);
3150 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003151}
3152
Raymond Hettinger60eca932003-02-09 06:40:58 +00003153PyDoc_STRVAR(ifilterfalse_doc,
3154"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
3155\n\
3156Return those items of sequence for which function(item) is false.\n\
3157If function is None, return the items that are false.");
3158
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003159static PyTypeObject ifilterfalse_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003160 PyVarObject_HEAD_INIT(NULL, 0)
3161 "itertools.ifilterfalse", /* tp_name */
3162 sizeof(ifilterfalseobject), /* tp_basicsize */
3163 0, /* tp_itemsize */
3164 /* methods */
3165 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
3166 0, /* tp_print */
3167 0, /* tp_getattr */
3168 0, /* tp_setattr */
3169 0, /* tp_compare */
3170 0, /* tp_repr */
3171 0, /* tp_as_number */
3172 0, /* tp_as_sequence */
3173 0, /* tp_as_mapping */
3174 0, /* tp_hash */
3175 0, /* tp_call */
3176 0, /* tp_str */
3177 PyObject_GenericGetAttr, /* tp_getattro */
3178 0, /* tp_setattro */
3179 0, /* tp_as_buffer */
3180 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3181 Py_TPFLAGS_BASETYPE, /* tp_flags */
3182 ifilterfalse_doc, /* tp_doc */
3183 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
3184 0, /* tp_clear */
3185 0, /* tp_richcompare */
3186 0, /* tp_weaklistoffset */
3187 PyObject_SelfIter, /* tp_iter */
3188 (iternextfunc)ifilterfalse_next, /* tp_iternext */
3189 0, /* tp_methods */
3190 0, /* tp_members */
3191 0, /* tp_getset */
3192 0, /* tp_base */
3193 0, /* tp_dict */
3194 0, /* tp_descr_get */
3195 0, /* tp_descr_set */
3196 0, /* tp_dictoffset */
3197 0, /* tp_init */
3198 0, /* tp_alloc */
3199 ifilterfalse_new, /* tp_new */
3200 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003201};
3202
3203
3204/* count object ************************************************************/
3205
3206typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003207 PyObject_HEAD
3208 Py_ssize_t cnt;
3209 PyObject *long_cnt;
3210 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003211} countobject;
3212
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003213/* Counting logic and invariants:
3214
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003215fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003216
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003217 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3218 Advances with: cnt += 1
3219 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003220
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003221slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003222
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003223 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3224 All counting is done with python objects (no overflows or underflows).
3225 Advances with: long_cnt += long_step
3226 Step may be zero -- effectively a slow version of repeat(cnt).
3227 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003228*/
3229
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003230static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003231
3232static PyObject *
3233count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3234{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003235 countobject *lz;
3236 int slow_mode = 0;
3237 Py_ssize_t cnt = 0;
3238 PyObject *long_cnt = NULL;
3239 PyObject *long_step = NULL;
3240 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003241
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003242 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3243 kwlist, &long_cnt, &long_step))
3244 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003245
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003246 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3247 (long_step != NULL && !PyNumber_Check(long_step))) {
3248 PyErr_SetString(PyExc_TypeError, "a number is required");
3249 return NULL;
3250 }
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003251
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003252 if (long_cnt != NULL) {
3253 cnt = PyInt_AsSsize_t(long_cnt);
3254 if ((cnt == -1 && PyErr_Occurred()) || !PyInt_Check(long_cnt)) {
3255 PyErr_Clear();
3256 slow_mode = 1;
3257 }
3258 Py_INCREF(long_cnt);
3259 } else {
3260 cnt = 0;
3261 long_cnt = PyInt_FromLong(0);
3262 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003263
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003264 /* If not specified, step defaults to 1 */
3265 if (long_step == NULL) {
3266 long_step = PyInt_FromLong(1);
3267 if (long_step == NULL) {
3268 Py_DECREF(long_cnt);
3269 return NULL;
3270 }
3271 } else
3272 Py_INCREF(long_step);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003273
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003274 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003275
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003276 /* Fast mode only works when the step is 1 */
3277 if (!PyInt_Check(long_step) ||
3278 PyInt_AS_LONG(long_step) != 1) {
3279 slow_mode = 1;
3280 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003281
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003282 if (slow_mode)
3283 cnt = PY_SSIZE_T_MAX;
3284 else
3285 Py_CLEAR(long_cnt);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003286
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003287 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3288 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3289 assert(slow_mode ||
3290 (PyInt_Check(long_step) && PyInt_AS_LONG(long_step) == 1));
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003291
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003292 /* create countobject structure */
3293 lz = (countobject *)type->tp_alloc(type, 0);
3294 if (lz == NULL) {
3295 Py_XDECREF(long_cnt);
3296 return NULL;
3297 }
3298 lz->cnt = cnt;
3299 lz->long_cnt = long_cnt;
3300 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003301
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003302 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003303}
3304
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003305static void
3306count_dealloc(countobject *lz)
3307{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003308 PyObject_GC_UnTrack(lz);
3309 Py_XDECREF(lz->long_cnt);
3310 Py_XDECREF(lz->long_step);
3311 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003312}
3313
Benjamin Peterson062a7c32009-02-16 21:07:52 +00003314static int
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003315count_traverse(countobject *lz, visitproc visit, void *arg)
3316{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003317 Py_VISIT(lz->long_cnt);
3318 Py_VISIT(lz->long_step);
3319 return 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003320}
3321
3322static PyObject *
3323count_nextlong(countobject *lz)
3324{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003325 PyObject *long_cnt;
3326 PyObject *stepped_up;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003327
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003328 long_cnt = lz->long_cnt;
3329 if (long_cnt == NULL) {
3330 /* Switch to slow_mode */
3331 long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
3332 if (long_cnt == NULL)
3333 return NULL;
3334 }
3335 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003336
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003337 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3338 if (stepped_up == NULL)
3339 return NULL;
3340 lz->long_cnt = stepped_up;
3341 return long_cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003342}
3343
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003344static PyObject *
3345count_next(countobject *lz)
3346{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003347 if (lz->cnt == PY_SSIZE_T_MAX)
3348 return count_nextlong(lz);
3349 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003350}
3351
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003352static PyObject *
3353count_repr(countobject *lz)
3354{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003355 PyObject *cnt_repr, *step_repr = NULL;
3356 PyObject *result = NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003357
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003358 if (lz->cnt != PY_SSIZE_T_MAX)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003359 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003360
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003361 cnt_repr = PyObject_Repr(lz->long_cnt);
3362 if (cnt_repr == NULL)
3363 return NULL;
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003364
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003365 if (PyInt_Check(lz->long_step) && PyInt_AS_LONG(lz->long_step) == 1) {
3366 /* Don't display step when it is an integer equal to 1 */
3367 result = PyString_FromFormat("count(%s)",
3368 PyString_AS_STRING(cnt_repr));
3369 } else {
3370 step_repr = PyObject_Repr(lz->long_step);
3371 if (step_repr != NULL)
3372 result = PyString_FromFormat("count(%s, %s)",
3373 PyString_AS_STRING(cnt_repr),
3374 PyString_AS_STRING(step_repr));
3375 }
3376 Py_DECREF(cnt_repr);
3377 Py_XDECREF(step_repr);
3378 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003379}
3380
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003381static PyObject *
3382count_reduce(countobject *lz)
3383{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003384 if (lz->cnt == PY_SSIZE_T_MAX)
3385 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3386 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003387}
3388
3389PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3390
3391static PyMethodDef count_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003392 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3393 count_reduce_doc},
3394 {NULL, NULL} /* sentinel */
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003395};
3396
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003397PyDoc_STRVAR(count_doc,
Ezio Melottifdc1e0d2010-06-11 02:21:25 +00003398 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003399\n\
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003400Return a count object whose .next() method returns consecutive values.\n\
3401Equivalent to:\n\n\
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003402 def count(firstval=0, step=1):\n\
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003403 x = firstval\n\
3404 while 1:\n\
3405 yield x\n\
3406 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003407
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003408static PyTypeObject count_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003409 PyVarObject_HEAD_INIT(NULL, 0)
3410 "itertools.count", /* tp_name */
3411 sizeof(countobject), /* tp_basicsize */
3412 0, /* tp_itemsize */
3413 /* methods */
3414 (destructor)count_dealloc, /* tp_dealloc */
3415 0, /* tp_print */
3416 0, /* tp_getattr */
3417 0, /* tp_setattr */
3418 0, /* tp_compare */
3419 (reprfunc)count_repr, /* tp_repr */
3420 0, /* tp_as_number */
3421 0, /* tp_as_sequence */
3422 0, /* tp_as_mapping */
3423 0, /* tp_hash */
3424 0, /* tp_call */
3425 0, /* tp_str */
3426 PyObject_GenericGetAttr, /* tp_getattro */
3427 0, /* tp_setattro */
3428 0, /* tp_as_buffer */
3429 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3430 Py_TPFLAGS_BASETYPE, /* tp_flags */
3431 count_doc, /* tp_doc */
3432 (traverseproc)count_traverse, /* tp_traverse */
3433 0, /* tp_clear */
3434 0, /* tp_richcompare */
3435 0, /* tp_weaklistoffset */
3436 PyObject_SelfIter, /* tp_iter */
3437 (iternextfunc)count_next, /* tp_iternext */
3438 count_methods, /* tp_methods */
3439 0, /* tp_members */
3440 0, /* tp_getset */
3441 0, /* tp_base */
3442 0, /* tp_dict */
3443 0, /* tp_descr_get */
3444 0, /* tp_descr_set */
3445 0, /* tp_dictoffset */
3446 0, /* tp_init */
3447 0, /* tp_alloc */
3448 count_new, /* tp_new */
3449 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003450};
3451
3452
3453/* izip object ************************************************************/
3454
3455#include "Python.h"
3456
3457typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003458 PyObject_HEAD
3459 Py_ssize_t tuplesize;
3460 PyObject *ittuple; /* tuple of iterators */
3461 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003462} izipobject;
3463
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003464static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003465
3466static PyObject *
3467izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3468{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003469 izipobject *lz;
3470 Py_ssize_t i;
3471 PyObject *ittuple; /* tuple of iterators */
3472 PyObject *result;
3473 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003474
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003475 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
3476 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003477
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003478 /* args must be a tuple */
3479 assert(PyTuple_Check(args));
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003480
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003481 /* obtain iterators */
3482 ittuple = PyTuple_New(tuplesize);
3483 if (ittuple == NULL)
3484 return NULL;
3485 for (i=0; i < tuplesize; ++i) {
3486 PyObject *item = PyTuple_GET_ITEM(args, i);
3487 PyObject *it = PyObject_GetIter(item);
3488 if (it == NULL) {
3489 if (PyErr_ExceptionMatches(PyExc_TypeError))
3490 PyErr_Format(PyExc_TypeError,
3491 "izip argument #%zd must support iteration",
3492 i+1);
3493 Py_DECREF(ittuple);
3494 return NULL;
3495 }
3496 PyTuple_SET_ITEM(ittuple, i, it);
3497 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003498
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003499 /* create a result holder */
3500 result = PyTuple_New(tuplesize);
3501 if (result == NULL) {
3502 Py_DECREF(ittuple);
3503 return NULL;
3504 }
3505 for (i=0 ; i < tuplesize ; i++) {
3506 Py_INCREF(Py_None);
3507 PyTuple_SET_ITEM(result, i, Py_None);
3508 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003509
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003510 /* create izipobject structure */
3511 lz = (izipobject *)type->tp_alloc(type, 0);
3512 if (lz == NULL) {
3513 Py_DECREF(ittuple);
3514 Py_DECREF(result);
3515 return NULL;
3516 }
3517 lz->ittuple = ittuple;
3518 lz->tuplesize = tuplesize;
3519 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003520
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003521 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003522}
3523
3524static void
3525izip_dealloc(izipobject *lz)
3526{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003527 PyObject_GC_UnTrack(lz);
3528 Py_XDECREF(lz->ittuple);
3529 Py_XDECREF(lz->result);
3530 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003531}
3532
3533static int
3534izip_traverse(izipobject *lz, visitproc visit, void *arg)
3535{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003536 Py_VISIT(lz->ittuple);
3537 Py_VISIT(lz->result);
3538 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003539}
3540
3541static PyObject *
3542izip_next(izipobject *lz)
3543{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003544 Py_ssize_t i;
3545 Py_ssize_t tuplesize = lz->tuplesize;
3546 PyObject *result = lz->result;
3547 PyObject *it;
3548 PyObject *item;
3549 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003550
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003551 if (tuplesize == 0)
3552 return NULL;
3553 if (Py_REFCNT(result) == 1) {
3554 Py_INCREF(result);
3555 for (i=0 ; i < tuplesize ; i++) {
3556 it = PyTuple_GET_ITEM(lz->ittuple, i);
3557 item = (*Py_TYPE(it)->tp_iternext)(it);
3558 if (item == NULL) {
3559 Py_DECREF(result);
3560 return NULL;
3561 }
3562 olditem = PyTuple_GET_ITEM(result, i);
3563 PyTuple_SET_ITEM(result, i, item);
3564 Py_DECREF(olditem);
3565 }
3566 } else {
3567 result = PyTuple_New(tuplesize);
3568 if (result == NULL)
3569 return NULL;
3570 for (i=0 ; i < tuplesize ; i++) {
3571 it = PyTuple_GET_ITEM(lz->ittuple, i);
3572 item = (*Py_TYPE(it)->tp_iternext)(it);
3573 if (item == NULL) {
3574 Py_DECREF(result);
3575 return NULL;
3576 }
3577 PyTuple_SET_ITEM(result, i, item);
3578 }
3579 }
3580 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003581}
3582
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003583PyDoc_STRVAR(izip_doc,
3584"izip(iter1 [,iter2 [...]]) --> izip object\n\
3585\n\
3586Return a izip object whose .next() method returns a tuple where\n\
3587the i-th element comes from the i-th iterable argument. The .next()\n\
3588method continues until the shortest iterable in the argument sequence\n\
3589is exhausted and then it raises StopIteration. Works like the zip()\n\
3590function but consumes less memory by returning an iterator instead of\n\
3591a list.");
3592
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003593static PyTypeObject izip_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003594 PyVarObject_HEAD_INIT(NULL, 0)
3595 "itertools.izip", /* tp_name */
3596 sizeof(izipobject), /* tp_basicsize */
3597 0, /* tp_itemsize */
3598 /* methods */
3599 (destructor)izip_dealloc, /* tp_dealloc */
3600 0, /* tp_print */
3601 0, /* tp_getattr */
3602 0, /* tp_setattr */
3603 0, /* tp_compare */
3604 0, /* tp_repr */
3605 0, /* tp_as_number */
3606 0, /* tp_as_sequence */
3607 0, /* tp_as_mapping */
3608 0, /* tp_hash */
3609 0, /* tp_call */
3610 0, /* tp_str */
3611 PyObject_GenericGetAttr, /* tp_getattro */
3612 0, /* tp_setattro */
3613 0, /* tp_as_buffer */
3614 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3615 Py_TPFLAGS_BASETYPE, /* tp_flags */
3616 izip_doc, /* tp_doc */
3617 (traverseproc)izip_traverse, /* tp_traverse */
3618 0, /* tp_clear */
3619 0, /* tp_richcompare */
3620 0, /* tp_weaklistoffset */
3621 PyObject_SelfIter, /* tp_iter */
3622 (iternextfunc)izip_next, /* tp_iternext */
3623 0, /* tp_methods */
3624 0, /* tp_members */
3625 0, /* tp_getset */
3626 0, /* tp_base */
3627 0, /* tp_dict */
3628 0, /* tp_descr_get */
3629 0, /* tp_descr_set */
3630 0, /* tp_dictoffset */
3631 0, /* tp_init */
3632 0, /* tp_alloc */
3633 izip_new, /* tp_new */
3634 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003635};
3636
3637
3638/* repeat object ************************************************************/
3639
3640typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003641 PyObject_HEAD
3642 PyObject *element;
3643 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003644} repeatobject;
3645
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003646static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003647
3648static PyObject *
3649repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3650{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003651 repeatobject *ro;
3652 PyObject *element;
3653 Py_ssize_t cnt = -1;
3654 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003655
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003656 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3657 &element, &cnt))
3658 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003659
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003660 if (PyTuple_Size(args) == 2 && cnt < 0)
3661 cnt = 0;
3662
3663 ro = (repeatobject *)type->tp_alloc(type, 0);
3664 if (ro == NULL)
3665 return NULL;
3666 Py_INCREF(element);
3667 ro->element = element;
3668 ro->cnt = cnt;
3669 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003670}
3671
3672static void
3673repeat_dealloc(repeatobject *ro)
3674{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003675 PyObject_GC_UnTrack(ro);
3676 Py_XDECREF(ro->element);
3677 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003678}
3679
3680static int
3681repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3682{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003683 Py_VISIT(ro->element);
3684 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003685}
3686
3687static PyObject *
3688repeat_next(repeatobject *ro)
3689{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003690 if (ro->cnt == 0)
3691 return NULL;
3692 if (ro->cnt > 0)
3693 ro->cnt--;
3694 Py_INCREF(ro->element);
3695 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003696}
3697
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003698static PyObject *
3699repeat_repr(repeatobject *ro)
3700{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003701 PyObject *result, *objrepr;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003702
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003703 objrepr = PyObject_Repr(ro->element);
3704 if (objrepr == NULL)
3705 return NULL;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003706
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003707 if (ro->cnt == -1)
3708 result = PyString_FromFormat("repeat(%s)",
3709 PyString_AS_STRING(objrepr));
3710 else
3711 result = PyString_FromFormat("repeat(%s, %zd)",
3712 PyString_AS_STRING(objrepr), ro->cnt);
3713 Py_DECREF(objrepr);
3714 return result;
3715}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003716
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003717static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003718repeat_len(repeatobject *ro)
3719{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003720 if (ro->cnt == -1) {
3721 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3722 return NULL;
3723 }
3724 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003725}
3726
Armin Rigof5b3e362006-02-11 21:32:43 +00003727PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003728
3729static PyMethodDef repeat_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003730 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3731 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003732};
3733
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003734PyDoc_STRVAR(repeat_doc,
Raymond Hettinger182edae2009-02-19 02:38:25 +00003735"repeat(object [,times]) -> create an iterator which returns the object\n\
3736for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003737endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003738
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003739static PyTypeObject repeat_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003740 PyVarObject_HEAD_INIT(NULL, 0)
3741 "itertools.repeat", /* tp_name */
3742 sizeof(repeatobject), /* tp_basicsize */
3743 0, /* tp_itemsize */
3744 /* methods */
3745 (destructor)repeat_dealloc, /* tp_dealloc */
3746 0, /* tp_print */
3747 0, /* tp_getattr */
3748 0, /* tp_setattr */
3749 0, /* tp_compare */
3750 (reprfunc)repeat_repr, /* tp_repr */
3751 0, /* tp_as_number */
3752 0, /* tp_as_sequence */
3753 0, /* tp_as_mapping */
3754 0, /* tp_hash */
3755 0, /* tp_call */
3756 0, /* tp_str */
3757 PyObject_GenericGetAttr, /* tp_getattro */
3758 0, /* tp_setattro */
3759 0, /* tp_as_buffer */
3760 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3761 Py_TPFLAGS_BASETYPE, /* tp_flags */
3762 repeat_doc, /* tp_doc */
3763 (traverseproc)repeat_traverse, /* tp_traverse */
3764 0, /* tp_clear */
3765 0, /* tp_richcompare */
3766 0, /* tp_weaklistoffset */
3767 PyObject_SelfIter, /* tp_iter */
3768 (iternextfunc)repeat_next, /* tp_iternext */
3769 repeat_methods, /* tp_methods */
3770 0, /* tp_members */
3771 0, /* tp_getset */
3772 0, /* tp_base */
3773 0, /* tp_dict */
3774 0, /* tp_descr_get */
3775 0, /* tp_descr_set */
3776 0, /* tp_dictoffset */
3777 0, /* tp_init */
3778 0, /* tp_alloc */
3779 repeat_new, /* tp_new */
3780 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003781};
3782
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003783/* iziplongest object ************************************************************/
3784
3785#include "Python.h"
3786
3787typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003788 PyObject_HEAD
3789 Py_ssize_t tuplesize;
3790 Py_ssize_t numactive;
3791 PyObject *ittuple; /* tuple of iterators */
3792 PyObject *result;
3793 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003794} iziplongestobject;
3795
3796static PyTypeObject iziplongest_type;
3797
3798static PyObject *
3799izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3800{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003801 iziplongestobject *lz;
3802 Py_ssize_t i;
3803 PyObject *ittuple; /* tuple of iterators */
3804 PyObject *result;
3805 PyObject *fillvalue = Py_None;
3806 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003807
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003808 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3809 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3810 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3811 PyErr_SetString(PyExc_TypeError,
3812 "izip_longest() got an unexpected keyword argument");
3813 return NULL;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003814 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003815 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003816
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003817 /* args must be a tuple */
3818 assert(PyTuple_Check(args));
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003819
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003820 /* obtain iterators */
3821 ittuple = PyTuple_New(tuplesize);
3822 if (ittuple == NULL)
3823 return NULL;
3824 for (i=0; i < tuplesize; ++i) {
3825 PyObject *item = PyTuple_GET_ITEM(args, i);
3826 PyObject *it = PyObject_GetIter(item);
3827 if (it == NULL) {
3828 if (PyErr_ExceptionMatches(PyExc_TypeError))
3829 PyErr_Format(PyExc_TypeError,
3830 "izip_longest argument #%zd must support iteration",
3831 i+1);
3832 Py_DECREF(ittuple);
3833 return NULL;
3834 }
3835 PyTuple_SET_ITEM(ittuple, i, it);
3836 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003837
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003838 /* create a result holder */
3839 result = PyTuple_New(tuplesize);
3840 if (result == NULL) {
3841 Py_DECREF(ittuple);
3842 return NULL;
3843 }
3844 for (i=0 ; i < tuplesize ; i++) {
3845 Py_INCREF(Py_None);
3846 PyTuple_SET_ITEM(result, i, Py_None);
3847 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003848
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003849 /* create iziplongestobject structure */
3850 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3851 if (lz == NULL) {
3852 Py_DECREF(ittuple);
3853 Py_DECREF(result);
3854 return NULL;
3855 }
3856 lz->ittuple = ittuple;
3857 lz->tuplesize = tuplesize;
3858 lz->numactive = tuplesize;
3859 lz->result = result;
3860 Py_INCREF(fillvalue);
3861 lz->fillvalue = fillvalue;
3862 return (PyObject *)lz;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003863}
3864
3865static void
3866izip_longest_dealloc(iziplongestobject *lz)
3867{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003868 PyObject_GC_UnTrack(lz);
3869 Py_XDECREF(lz->ittuple);
3870 Py_XDECREF(lz->result);
3871 Py_XDECREF(lz->fillvalue);
3872 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003873}
3874
3875static int
3876izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3877{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003878 Py_VISIT(lz->ittuple);
3879 Py_VISIT(lz->result);
3880 Py_VISIT(lz->fillvalue);
3881 return 0;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003882}
3883
3884static PyObject *
3885izip_longest_next(iziplongestobject *lz)
3886{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003887 Py_ssize_t i;
3888 Py_ssize_t tuplesize = lz->tuplesize;
3889 PyObject *result = lz->result;
3890 PyObject *it;
3891 PyObject *item;
3892 PyObject *olditem;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003893
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003894 if (tuplesize == 0)
3895 return NULL;
3896 if (lz->numactive == 0)
3897 return NULL;
3898 if (Py_REFCNT(result) == 1) {
3899 Py_INCREF(result);
3900 for (i=0 ; i < tuplesize ; i++) {
3901 it = PyTuple_GET_ITEM(lz->ittuple, i);
3902 if (it == NULL) {
3903 Py_INCREF(lz->fillvalue);
3904 item = lz->fillvalue;
3905 } else {
3906 item = PyIter_Next(it);
3907 if (item == NULL) {
3908 lz->numactive -= 1;
3909 if (lz->numactive == 0 || PyErr_Occurred()) {
3910 lz->numactive = 0;
3911 Py_DECREF(result);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003912 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003913 } else {
3914 Py_INCREF(lz->fillvalue);
3915 item = lz->fillvalue;
3916 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3917 Py_DECREF(it);
3918 }
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003919 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003920 }
3921 olditem = PyTuple_GET_ITEM(result, i);
3922 PyTuple_SET_ITEM(result, i, item);
3923 Py_DECREF(olditem);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003924 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003925 } else {
3926 result = PyTuple_New(tuplesize);
3927 if (result == NULL)
3928 return NULL;
3929 for (i=0 ; i < tuplesize ; i++) {
3930 it = PyTuple_GET_ITEM(lz->ittuple, i);
3931 if (it == NULL) {
3932 Py_INCREF(lz->fillvalue);
3933 item = lz->fillvalue;
3934 } else {
3935 item = PyIter_Next(it);
3936 if (item == NULL) {
3937 lz->numactive -= 1;
3938 if (lz->numactive == 0 || PyErr_Occurred()) {
3939 lz->numactive = 0;
3940 Py_DECREF(result);
3941 return NULL;
3942 } else {
3943 Py_INCREF(lz->fillvalue);
3944 item = lz->fillvalue;
3945 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3946 Py_DECREF(it);
3947 }
3948 }
3949 }
3950 PyTuple_SET_ITEM(result, i, item);
3951 }
3952 }
3953 return result;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003954}
3955
3956PyDoc_STRVAR(izip_longest_doc,
3957"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3958\n\
3959Return an izip_longest object whose .next() method returns a tuple where\n\
3960the i-th element comes from the i-th iterable argument. The .next()\n\
3961method continues until the longest iterable in the argument sequence\n\
3962is exhausted and then it raises StopIteration. When the shorter iterables\n\
3963are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3964defaults to None or can be specified by a keyword argument.\n\
3965");
3966
3967static PyTypeObject iziplongest_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003968 PyVarObject_HEAD_INIT(NULL, 0)
3969 "itertools.izip_longest", /* tp_name */
3970 sizeof(iziplongestobject), /* tp_basicsize */
3971 0, /* tp_itemsize */
3972 /* methods */
3973 (destructor)izip_longest_dealloc, /* tp_dealloc */
3974 0, /* tp_print */
3975 0, /* tp_getattr */
3976 0, /* tp_setattr */
3977 0, /* tp_compare */
3978 0, /* tp_repr */
3979 0, /* tp_as_number */
3980 0, /* tp_as_sequence */
3981 0, /* tp_as_mapping */
3982 0, /* tp_hash */
3983 0, /* tp_call */
3984 0, /* tp_str */
3985 PyObject_GenericGetAttr, /* tp_getattro */
3986 0, /* tp_setattro */
3987 0, /* tp_as_buffer */
3988 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3989 Py_TPFLAGS_BASETYPE, /* tp_flags */
3990 izip_longest_doc, /* tp_doc */
3991 (traverseproc)izip_longest_traverse, /* tp_traverse */
3992 0, /* tp_clear */
3993 0, /* tp_richcompare */
3994 0, /* tp_weaklistoffset */
3995 PyObject_SelfIter, /* tp_iter */
3996 (iternextfunc)izip_longest_next, /* tp_iternext */
3997 0, /* tp_methods */
3998 0, /* tp_members */
3999 0, /* tp_getset */
4000 0, /* tp_base */
4001 0, /* tp_dict */
4002 0, /* tp_descr_get */
4003 0, /* tp_descr_set */
4004 0, /* tp_dictoffset */
4005 0, /* tp_init */
4006 0, /* tp_alloc */
4007 izip_longest_new, /* tp_new */
4008 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd36862c2007-02-21 05:20:38 +00004009};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004010
4011/* module level code ********************************************************/
4012
4013PyDoc_STRVAR(module_doc,
4014"Functional tools for creating and using iterators.\n\
4015\n\
4016Infinite iterators:\n\
4017count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004018cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004019repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004020\n\
4021Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004022chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettingere76816b2009-01-29 03:43:44 +00004023compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004024dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4025groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00004026ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
4027ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004028islice(seq, [start,] stop [, step]) --> elements from\n\
4029 seq[start:stop:step]\n\
4030imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
4031starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004032tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004033takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004034izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4035izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4036\n\
4037Combinatoric generators:\n\
4038product(p, q, ... [repeat=1]) --> cartesian product\n\
4039permutations(p[, r])\n\
Raymond Hettinger9eac1192009-11-19 01:22:04 +00004040combinations(p, r)\n\
4041combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004042");
4043
4044
Raymond Hettingerad983e72003-11-12 14:32:26 +00004045static PyMethodDef module_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004046 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4047 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004048};
4049
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004050PyMODINIT_FUNC
4051inititertools(void)
4052{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004053 int i;
4054 PyObject *m;
4055 char *name;
4056 PyTypeObject *typelist[] = {
4057 &combinations_type,
4058 &cwr_type,
4059 &cycle_type,
4060 &dropwhile_type,
4061 &takewhile_type,
4062 &islice_type,
4063 &starmap_type,
4064 &imap_type,
4065 &chain_type,
4066 &compress_type,
4067 &ifilter_type,
4068 &ifilterfalse_type,
4069 &count_type,
4070 &izip_type,
4071 &iziplongest_type,
4072 &permutations_type,
4073 &product_type,
4074 &repeat_type,
4075 &groupby_type,
4076 NULL
4077 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004078
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004079 Py_TYPE(&teedataobject_type) = &PyType_Type;
4080 m = Py_InitModule3("itertools", module_methods, module_doc);
4081 if (m == NULL)
4082 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004083
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004084 for (i=0 ; typelist[i] != NULL ; i++) {
4085 if (PyType_Ready(typelist[i]) < 0)
4086 return;
4087 name = strchr(typelist[i]->tp_name, '.');
4088 assert (name != NULL);
4089 Py_INCREF(typelist[i]);
4090 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4091 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004092
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004093 if (PyType_Ready(&teedataobject_type) < 0)
4094 return;
4095 if (PyType_Ready(&tee_type) < 0)
4096 return;
4097 if (PyType_Ready(&_grouper_type) < 0)
4098 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004099}