blob: c6d92eafc9c7dd7252c5ba79c77ba5c2d6cffdf1 [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;
1218 Py_ssize_t oldnext;
1219 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001220
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001221 iternext = *Py_TYPE(it)->tp_iternext;
1222 while (lz->cnt < lz->next) {
1223 item = iternext(it);
1224 if (item == NULL)
1225 return NULL;
1226 Py_DECREF(item);
1227 lz->cnt++;
1228 }
1229 if (lz->stop != -1 && lz->cnt >= lz->stop)
1230 return NULL;
1231 item = iternext(it);
1232 if (item == NULL)
1233 return NULL;
1234 lz->cnt++;
1235 oldnext = lz->next;
1236 lz->next += lz->step;
1237 if (lz->next < oldnext) /* Check for overflow */
1238 lz->next = lz->stop;
1239 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001240}
1241
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001242PyDoc_STRVAR(islice_doc,
1243"islice(iterable, [start,] stop [, step]) --> islice object\n\
1244\n\
1245Return an iterator whose next() method returns selected values from an\n\
1246iterable. If start is specified, will skip all preceding elements;\n\
1247otherwise, start defaults to zero. Step defaults to one. If\n\
1248specified as another value, step determines how many values are \n\
1249skipped between successive calls. Works like a slice() on a list\n\
1250but returns an iterator.");
1251
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001252static PyTypeObject islice_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001253 PyVarObject_HEAD_INIT(NULL, 0)
1254 "itertools.islice", /* tp_name */
1255 sizeof(isliceobject), /* tp_basicsize */
1256 0, /* tp_itemsize */
1257 /* methods */
1258 (destructor)islice_dealloc, /* tp_dealloc */
1259 0, /* tp_print */
1260 0, /* tp_getattr */
1261 0, /* tp_setattr */
1262 0, /* tp_compare */
1263 0, /* tp_repr */
1264 0, /* tp_as_number */
1265 0, /* tp_as_sequence */
1266 0, /* tp_as_mapping */
1267 0, /* tp_hash */
1268 0, /* tp_call */
1269 0, /* tp_str */
1270 PyObject_GenericGetAttr, /* tp_getattro */
1271 0, /* tp_setattro */
1272 0, /* tp_as_buffer */
1273 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1274 Py_TPFLAGS_BASETYPE, /* tp_flags */
1275 islice_doc, /* tp_doc */
1276 (traverseproc)islice_traverse, /* tp_traverse */
1277 0, /* tp_clear */
1278 0, /* tp_richcompare */
1279 0, /* tp_weaklistoffset */
1280 PyObject_SelfIter, /* tp_iter */
1281 (iternextfunc)islice_next, /* tp_iternext */
1282 0, /* tp_methods */
1283 0, /* tp_members */
1284 0, /* tp_getset */
1285 0, /* tp_base */
1286 0, /* tp_dict */
1287 0, /* tp_descr_get */
1288 0, /* tp_descr_set */
1289 0, /* tp_dictoffset */
1290 0, /* tp_init */
1291 0, /* tp_alloc */
1292 islice_new, /* tp_new */
1293 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001294};
1295
1296
1297/* starmap object ************************************************************/
1298
1299typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001300 PyObject_HEAD
1301 PyObject *func;
1302 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001303} starmapobject;
1304
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001305static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001306
1307static PyObject *
1308starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1309{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001310 PyObject *func, *seq;
1311 PyObject *it;
1312 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001313
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001314 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1315 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001316
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001317 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1318 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001319
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001320 /* Get iterator. */
1321 it = PyObject_GetIter(seq);
1322 if (it == NULL)
1323 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001324
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001325 /* create starmapobject structure */
1326 lz = (starmapobject *)type->tp_alloc(type, 0);
1327 if (lz == NULL) {
1328 Py_DECREF(it);
1329 return NULL;
1330 }
1331 Py_INCREF(func);
1332 lz->func = func;
1333 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001334
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001335 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001336}
1337
1338static void
1339starmap_dealloc(starmapobject *lz)
1340{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001341 PyObject_GC_UnTrack(lz);
1342 Py_XDECREF(lz->func);
1343 Py_XDECREF(lz->it);
1344 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001345}
1346
1347static int
1348starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1349{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001350 Py_VISIT(lz->it);
1351 Py_VISIT(lz->func);
1352 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001353}
1354
1355static PyObject *
1356starmap_next(starmapobject *lz)
1357{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001358 PyObject *args;
1359 PyObject *result;
1360 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001361
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001362 args = (*Py_TYPE(it)->tp_iternext)(it);
1363 if (args == NULL)
1364 return NULL;
1365 if (!PyTuple_CheckExact(args)) {
1366 PyObject *newargs = PySequence_Tuple(args);
1367 Py_DECREF(args);
1368 if (newargs == NULL)
1369 return NULL;
1370 args = newargs;
1371 }
1372 result = PyObject_Call(lz->func, args, NULL);
1373 Py_DECREF(args);
1374 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001375}
1376
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001377PyDoc_STRVAR(starmap_doc,
1378"starmap(function, sequence) --> starmap object\n\
1379\n\
1380Return an iterator whose values are returned from the function evaluated\n\
1381with a argument tuple taken from the given sequence.");
1382
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001383static PyTypeObject starmap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001384 PyVarObject_HEAD_INIT(NULL, 0)
1385 "itertools.starmap", /* tp_name */
1386 sizeof(starmapobject), /* tp_basicsize */
1387 0, /* tp_itemsize */
1388 /* methods */
1389 (destructor)starmap_dealloc, /* tp_dealloc */
1390 0, /* tp_print */
1391 0, /* tp_getattr */
1392 0, /* tp_setattr */
1393 0, /* tp_compare */
1394 0, /* tp_repr */
1395 0, /* tp_as_number */
1396 0, /* tp_as_sequence */
1397 0, /* tp_as_mapping */
1398 0, /* tp_hash */
1399 0, /* tp_call */
1400 0, /* tp_str */
1401 PyObject_GenericGetAttr, /* tp_getattro */
1402 0, /* tp_setattro */
1403 0, /* tp_as_buffer */
1404 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1405 Py_TPFLAGS_BASETYPE, /* tp_flags */
1406 starmap_doc, /* tp_doc */
1407 (traverseproc)starmap_traverse, /* tp_traverse */
1408 0, /* tp_clear */
1409 0, /* tp_richcompare */
1410 0, /* tp_weaklistoffset */
1411 PyObject_SelfIter, /* tp_iter */
1412 (iternextfunc)starmap_next, /* tp_iternext */
1413 0, /* tp_methods */
1414 0, /* tp_members */
1415 0, /* tp_getset */
1416 0, /* tp_base */
1417 0, /* tp_dict */
1418 0, /* tp_descr_get */
1419 0, /* tp_descr_set */
1420 0, /* tp_dictoffset */
1421 0, /* tp_init */
1422 0, /* tp_alloc */
1423 starmap_new, /* tp_new */
1424 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001425};
1426
1427
1428/* imap object ************************************************************/
1429
1430typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001431 PyObject_HEAD
1432 PyObject *iters;
1433 PyObject *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001434} imapobject;
1435
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001436static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001437
1438static PyObject *
1439imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1440{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001441 PyObject *it, *iters, *func;
1442 imapobject *lz;
1443 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001444
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001445 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
1446 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001447
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001448 numargs = PyTuple_Size(args);
1449 if (numargs < 2) {
1450 PyErr_SetString(PyExc_TypeError,
1451 "imap() must have at least two arguments.");
1452 return NULL;
1453 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001454
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001455 iters = PyTuple_New(numargs-1);
1456 if (iters == NULL)
1457 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001458
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001459 for (i=1 ; i<numargs ; i++) {
1460 /* Get iterator. */
1461 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1462 if (it == NULL) {
1463 Py_DECREF(iters);
1464 return NULL;
1465 }
1466 PyTuple_SET_ITEM(iters, i-1, it);
1467 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001468
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001469 /* create imapobject structure */
1470 lz = (imapobject *)type->tp_alloc(type, 0);
1471 if (lz == NULL) {
1472 Py_DECREF(iters);
1473 return NULL;
1474 }
1475 lz->iters = iters;
1476 func = PyTuple_GET_ITEM(args, 0);
1477 Py_INCREF(func);
1478 lz->func = func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001479
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001480 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001481}
1482
1483static void
1484imap_dealloc(imapobject *lz)
1485{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001486 PyObject_GC_UnTrack(lz);
1487 Py_XDECREF(lz->iters);
1488 Py_XDECREF(lz->func);
1489 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001490}
1491
1492static int
1493imap_traverse(imapobject *lz, visitproc visit, void *arg)
1494{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001495 Py_VISIT(lz->iters);
1496 Py_VISIT(lz->func);
1497 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001498}
1499
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001500/*
Raymond Hettinger2012f172003-02-07 05:32:58 +00001501imap() is an iterator version of __builtins__.map() except that it does
1502not have the None fill-in feature. That was intentionally left out for
1503the following reasons:
1504
1505 1) Itertools are designed to be easily combined and chained together.
1506 Having all tools stop with the shortest input is a unifying principle
1507 that makes it easier to combine finite iterators (supplying data) with
1508 infinite iterators like count() and repeat() (for supplying sequential
1509 or constant arguments to a function).
1510
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001511 2) In typical use cases for combining itertools, having one finite data
Raymond Hettinger2012f172003-02-07 05:32:58 +00001512 supplier run out before another is likely to be an error condition which
1513 should not pass silently by automatically supplying None.
1514
1515 3) The use cases for automatic None fill-in are rare -- not many functions
1516 do something useful when a parameter suddenly switches type and becomes
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001517 None.
Raymond Hettinger2012f172003-02-07 05:32:58 +00001518
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001519 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001520 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001521
1522 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1523*/
1524
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001525static PyObject *
1526imap_next(imapobject *lz)
1527{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001528 PyObject *val;
1529 PyObject *argtuple;
1530 PyObject *result;
1531 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001532
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001533 numargs = PyTuple_Size(lz->iters);
1534 argtuple = PyTuple_New(numargs);
1535 if (argtuple == NULL)
1536 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001537
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001538 for (i=0 ; i<numargs ; i++) {
1539 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1540 if (val == NULL) {
1541 Py_DECREF(argtuple);
1542 return NULL;
1543 }
1544 PyTuple_SET_ITEM(argtuple, i, val);
1545 }
1546 if (lz->func == Py_None)
1547 return argtuple;
1548 result = PyObject_Call(lz->func, argtuple, NULL);
1549 Py_DECREF(argtuple);
1550 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001551}
1552
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001553PyDoc_STRVAR(imap_doc,
1554"imap(func, *iterables) --> imap object\n\
1555\n\
1556Make an iterator that computes the function using arguments from\n\
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001557each of the iterables. Like map() except that it returns\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001558an iterator instead of a list and that it stops when the shortest\n\
1559iterable is exhausted instead of filling in None for shorter\n\
1560iterables.");
1561
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001562static PyTypeObject imap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001563 PyVarObject_HEAD_INIT(NULL, 0)
1564 "itertools.imap", /* tp_name */
1565 sizeof(imapobject), /* tp_basicsize */
1566 0, /* tp_itemsize */
1567 /* methods */
1568 (destructor)imap_dealloc, /* tp_dealloc */
1569 0, /* tp_print */
1570 0, /* tp_getattr */
1571 0, /* tp_setattr */
1572 0, /* tp_compare */
1573 0, /* tp_repr */
1574 0, /* tp_as_number */
1575 0, /* tp_as_sequence */
1576 0, /* tp_as_mapping */
1577 0, /* tp_hash */
1578 0, /* tp_call */
1579 0, /* tp_str */
1580 PyObject_GenericGetAttr, /* tp_getattro */
1581 0, /* tp_setattro */
1582 0, /* tp_as_buffer */
1583 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1584 Py_TPFLAGS_BASETYPE, /* tp_flags */
1585 imap_doc, /* tp_doc */
1586 (traverseproc)imap_traverse, /* tp_traverse */
1587 0, /* tp_clear */
1588 0, /* tp_richcompare */
1589 0, /* tp_weaklistoffset */
1590 PyObject_SelfIter, /* tp_iter */
1591 (iternextfunc)imap_next, /* tp_iternext */
1592 0, /* tp_methods */
1593 0, /* tp_members */
1594 0, /* tp_getset */
1595 0, /* tp_base */
1596 0, /* tp_dict */
1597 0, /* tp_descr_get */
1598 0, /* tp_descr_set */
1599 0, /* tp_dictoffset */
1600 0, /* tp_init */
1601 0, /* tp_alloc */
1602 imap_new, /* tp_new */
1603 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001604};
1605
1606
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001607/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001608
1609typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001610 PyObject_HEAD
1611 PyObject *source; /* Iterator over input iterables */
1612 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001613} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001614
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001615static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001616
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001617static PyObject *
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001618chain_new_internal(PyTypeObject *type, PyObject *source)
1619{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001620 chainobject *lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001621
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001622 lz = (chainobject *)type->tp_alloc(type, 0);
1623 if (lz == NULL) {
1624 Py_DECREF(source);
1625 return NULL;
1626 }
1627
1628 lz->source = source;
1629 lz->active = NULL;
1630 return (PyObject *)lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001631}
1632
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001633static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001634chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001635{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001636 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001637
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001638 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1639 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001640
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001641 source = PyObject_GetIter(args);
1642 if (source == NULL)
1643 return NULL;
1644
1645 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001646}
1647
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001648static PyObject *
1649chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1650{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001651 PyObject *source;
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001652
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001653 source = PyObject_GetIter(arg);
1654 if (source == NULL)
1655 return NULL;
1656
1657 return chain_new_internal(type, source);
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001658}
1659
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001660static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001661chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001662{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001663 PyObject_GC_UnTrack(lz);
1664 Py_XDECREF(lz->active);
1665 Py_XDECREF(lz->source);
1666 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001667}
1668
Raymond Hettinger2012f172003-02-07 05:32:58 +00001669static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001670chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001671{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001672 Py_VISIT(lz->source);
1673 Py_VISIT(lz->active);
1674 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001675}
1676
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001677static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001678chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001679{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001680 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001681
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001682 if (lz->source == NULL)
1683 return NULL; /* already stopped */
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001684
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001685 if (lz->active == NULL) {
1686 PyObject *iterable = PyIter_Next(lz->source);
1687 if (iterable == NULL) {
1688 Py_CLEAR(lz->source);
1689 return NULL; /* no more input sources */
1690 }
1691 lz->active = PyObject_GetIter(iterable);
1692 Py_DECREF(iterable);
1693 if (lz->active == NULL) {
1694 Py_CLEAR(lz->source);
1695 return NULL; /* input not iterable */
1696 }
1697 }
1698 item = PyIter_Next(lz->active);
1699 if (item != NULL)
1700 return item;
1701 if (PyErr_Occurred()) {
1702 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1703 PyErr_Clear();
1704 else
1705 return NULL; /* input raised an exception */
1706 }
1707 Py_CLEAR(lz->active);
1708 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001709}
1710
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001711PyDoc_STRVAR(chain_doc,
1712"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001713\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001714Return a chain object whose .next() method returns elements from the\n\
1715first iterable until it is exhausted, then elements from the next\n\
1716iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001717
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001718PyDoc_STRVAR(chain_from_iterable_doc,
1719"chain.from_iterable(iterable) --> chain object\n\
1720\n\
1721Alternate chain() contructor taking a single iterable argument\n\
1722that evaluates lazily.");
1723
1724static PyMethodDef chain_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001725 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1726 chain_from_iterable_doc},
1727 {NULL, NULL} /* sentinel */
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001728};
1729
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001730static PyTypeObject chain_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001731 PyVarObject_HEAD_INIT(NULL, 0)
1732 "itertools.chain", /* tp_name */
1733 sizeof(chainobject), /* tp_basicsize */
1734 0, /* tp_itemsize */
1735 /* methods */
1736 (destructor)chain_dealloc, /* tp_dealloc */
1737 0, /* tp_print */
1738 0, /* tp_getattr */
1739 0, /* tp_setattr */
1740 0, /* tp_compare */
1741 0, /* tp_repr */
1742 0, /* tp_as_number */
1743 0, /* tp_as_sequence */
1744 0, /* tp_as_mapping */
1745 0, /* tp_hash */
1746 0, /* tp_call */
1747 0, /* tp_str */
1748 PyObject_GenericGetAttr, /* tp_getattro */
1749 0, /* tp_setattro */
1750 0, /* tp_as_buffer */
1751 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1752 Py_TPFLAGS_BASETYPE, /* tp_flags */
1753 chain_doc, /* tp_doc */
1754 (traverseproc)chain_traverse, /* tp_traverse */
1755 0, /* tp_clear */
1756 0, /* tp_richcompare */
1757 0, /* tp_weaklistoffset */
1758 PyObject_SelfIter, /* tp_iter */
1759 (iternextfunc)chain_next, /* tp_iternext */
1760 chain_methods, /* tp_methods */
1761 0, /* tp_members */
1762 0, /* tp_getset */
1763 0, /* tp_base */
1764 0, /* tp_dict */
1765 0, /* tp_descr_get */
1766 0, /* tp_descr_set */
1767 0, /* tp_dictoffset */
1768 0, /* tp_init */
1769 0, /* tp_alloc */
1770 chain_new, /* tp_new */
1771 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001772};
1773
1774
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001775/* product object ************************************************************/
1776
1777typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001778 PyObject_HEAD
1779 PyObject *pools; /* tuple of pool tuples */
1780 Py_ssize_t *indices; /* one index per pool */
1781 PyObject *result; /* most recently returned result tuple */
1782 int stopped; /* set to 1 when the product iterator is exhausted */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001783} productobject;
1784
1785static PyTypeObject product_type;
1786
1787static PyObject *
1788product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1789{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001790 productobject *lz;
1791 Py_ssize_t nargs, npools, repeat=1;
1792 PyObject *pools = NULL;
1793 Py_ssize_t *indices = NULL;
1794 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001795
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001796 if (kwds != NULL) {
1797 char *kwlist[] = {"repeat", 0};
1798 PyObject *tmpargs = PyTuple_New(0);
1799 if (tmpargs == NULL)
1800 return NULL;
1801 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1802 Py_DECREF(tmpargs);
1803 return NULL;
1804 }
1805 Py_DECREF(tmpargs);
1806 if (repeat < 0) {
1807 PyErr_SetString(PyExc_ValueError,
1808 "repeat argument cannot be negative");
1809 return NULL;
1810 }
1811 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001812
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001813 assert(PyTuple_Check(args));
1814 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1815 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001816
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001817 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1818 if (indices == NULL) {
1819 PyErr_NoMemory();
1820 goto error;
1821 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001822
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001823 pools = PyTuple_New(npools);
1824 if (pools == NULL)
1825 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001826
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001827 for (i=0; i < nargs ; ++i) {
1828 PyObject *item = PyTuple_GET_ITEM(args, i);
1829 PyObject *pool = PySequence_Tuple(item);
1830 if (pool == NULL)
1831 goto error;
1832 PyTuple_SET_ITEM(pools, i, pool);
1833 indices[i] = 0;
1834 }
1835 for ( ; i < npools; ++i) {
1836 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1837 Py_INCREF(pool);
1838 PyTuple_SET_ITEM(pools, i, pool);
1839 indices[i] = 0;
1840 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001841
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001842 /* create productobject structure */
1843 lz = (productobject *)type->tp_alloc(type, 0);
1844 if (lz == NULL)
1845 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001846
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001847 lz->pools = pools;
1848 lz->indices = indices;
1849 lz->result = NULL;
1850 lz->stopped = 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001851
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001852 return (PyObject *)lz;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001853
1854error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001855 if (indices != NULL)
1856 PyMem_Free(indices);
1857 Py_XDECREF(pools);
1858 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001859}
1860
1861static void
1862product_dealloc(productobject *lz)
1863{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001864 PyObject_GC_UnTrack(lz);
1865 Py_XDECREF(lz->pools);
1866 Py_XDECREF(lz->result);
1867 if (lz->indices != NULL)
1868 PyMem_Free(lz->indices);
1869 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001870}
1871
1872static int
1873product_traverse(productobject *lz, visitproc visit, void *arg)
1874{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001875 Py_VISIT(lz->pools);
1876 Py_VISIT(lz->result);
1877 return 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001878}
1879
1880static PyObject *
1881product_next(productobject *lz)
1882{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001883 PyObject *pool;
1884 PyObject *elem;
1885 PyObject *oldelem;
1886 PyObject *pools = lz->pools;
1887 PyObject *result = lz->result;
1888 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1889 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001890
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001891 if (lz->stopped)
1892 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001893
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001894 if (result == NULL) {
1895 /* On the first pass, return an initial tuple filled with the
1896 first element from each pool. */
1897 result = PyTuple_New(npools);
1898 if (result == NULL)
1899 goto empty;
1900 lz->result = result;
1901 for (i=0; i < npools; i++) {
1902 pool = PyTuple_GET_ITEM(pools, i);
1903 if (PyTuple_GET_SIZE(pool) == 0)
1904 goto empty;
1905 elem = PyTuple_GET_ITEM(pool, 0);
1906 Py_INCREF(elem);
1907 PyTuple_SET_ITEM(result, i, elem);
1908 }
1909 } else {
1910 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001911
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001912 /* Copy the previous result tuple or re-use it if available */
1913 if (Py_REFCNT(result) > 1) {
1914 PyObject *old_result = result;
1915 result = PyTuple_New(npools);
1916 if (result == NULL)
1917 goto empty;
1918 lz->result = result;
1919 for (i=0; i < npools; i++) {
1920 elem = PyTuple_GET_ITEM(old_result, i);
1921 Py_INCREF(elem);
1922 PyTuple_SET_ITEM(result, i, elem);
1923 }
1924 Py_DECREF(old_result);
1925 }
1926 /* Now, we've got the only copy so we can update it in-place */
1927 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001928
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001929 /* Update the pool indices right-to-left. Only advance to the
1930 next pool when the previous one rolls-over */
1931 for (i=npools-1 ; i >= 0 ; i--) {
1932 pool = PyTuple_GET_ITEM(pools, i);
1933 indices[i]++;
1934 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1935 /* Roll-over and advance to next pool */
1936 indices[i] = 0;
1937 elem = PyTuple_GET_ITEM(pool, 0);
1938 Py_INCREF(elem);
1939 oldelem = PyTuple_GET_ITEM(result, i);
1940 PyTuple_SET_ITEM(result, i, elem);
1941 Py_DECREF(oldelem);
1942 } else {
1943 /* No rollover. Just increment and stop here. */
1944 elem = PyTuple_GET_ITEM(pool, indices[i]);
1945 Py_INCREF(elem);
1946 oldelem = PyTuple_GET_ITEM(result, i);
1947 PyTuple_SET_ITEM(result, i, elem);
1948 Py_DECREF(oldelem);
1949 break;
1950 }
1951 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00001952
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001953 /* If i is negative, then the indices have all rolled-over
1954 and we're done. */
1955 if (i < 0)
1956 goto empty;
1957 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001958
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001959 Py_INCREF(result);
1960 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001961
1962empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001963 lz->stopped = 1;
1964 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001965}
1966
1967PyDoc_STRVAR(product_doc,
1968"product(*iterables) --> product object\n\
1969\n\
Raymond Hettinger73d79632008-02-23 02:20:41 +00001970Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001971For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1972The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1973cycle in a manner similar to an odometer (with the rightmost element changing\n\
1974on every iteration).\n\n\
Raymond Hettinger94034ea2009-02-09 18:39:41 +00001975To compute the product of an iterable with itself, specify the number\n\
1976of repetitions with the optional repeat keyword argument. For example,\n\
1977product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001978product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1979product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1980
1981static PyTypeObject product_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001982 PyVarObject_HEAD_INIT(NULL, 0)
1983 "itertools.product", /* tp_name */
1984 sizeof(productobject), /* tp_basicsize */
1985 0, /* tp_itemsize */
1986 /* methods */
1987 (destructor)product_dealloc, /* tp_dealloc */
1988 0, /* tp_print */
1989 0, /* tp_getattr */
1990 0, /* tp_setattr */
1991 0, /* tp_compare */
1992 0, /* tp_repr */
1993 0, /* tp_as_number */
1994 0, /* tp_as_sequence */
1995 0, /* tp_as_mapping */
1996 0, /* tp_hash */
1997 0, /* tp_call */
1998 0, /* tp_str */
1999 PyObject_GenericGetAttr, /* tp_getattro */
2000 0, /* tp_setattro */
2001 0, /* tp_as_buffer */
2002 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2003 Py_TPFLAGS_BASETYPE, /* tp_flags */
2004 product_doc, /* tp_doc */
2005 (traverseproc)product_traverse, /* tp_traverse */
2006 0, /* tp_clear */
2007 0, /* tp_richcompare */
2008 0, /* tp_weaklistoffset */
2009 PyObject_SelfIter, /* tp_iter */
2010 (iternextfunc)product_next, /* tp_iternext */
2011 0, /* tp_methods */
2012 0, /* tp_members */
2013 0, /* tp_getset */
2014 0, /* tp_base */
2015 0, /* tp_dict */
2016 0, /* tp_descr_get */
2017 0, /* tp_descr_set */
2018 0, /* tp_dictoffset */
2019 0, /* tp_init */
2020 0, /* tp_alloc */
2021 product_new, /* tp_new */
2022 PyObject_GC_Del, /* tp_free */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002023};
2024
2025
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002026/* combinations object ************************************************************/
2027
2028typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002029 PyObject_HEAD
2030 PyObject *pool; /* input converted to a tuple */
2031 Py_ssize_t *indices; /* one index per result element */
2032 PyObject *result; /* most recently returned result tuple */
2033 Py_ssize_t r; /* size of result tuple */
2034 int stopped; /* set to 1 when the combinations iterator is exhausted */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002035} combinationsobject;
2036
2037static PyTypeObject combinations_type;
2038
2039static PyObject *
2040combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2041{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002042 combinationsobject *co;
2043 Py_ssize_t n;
2044 Py_ssize_t r;
2045 PyObject *pool = NULL;
2046 PyObject *iterable = NULL;
2047 Py_ssize_t *indices = NULL;
2048 Py_ssize_t i;
2049 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002050
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002051 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2052 &iterable, &r))
2053 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002054
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002055 pool = PySequence_Tuple(iterable);
2056 if (pool == NULL)
2057 goto error;
2058 n = PyTuple_GET_SIZE(pool);
2059 if (r < 0) {
2060 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2061 goto error;
2062 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002063
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002064 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2065 if (indices == NULL) {
2066 PyErr_NoMemory();
2067 goto error;
2068 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002069
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002070 for (i=0 ; i<r ; i++)
2071 indices[i] = i;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002072
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002073 /* create combinationsobject structure */
2074 co = (combinationsobject *)type->tp_alloc(type, 0);
2075 if (co == NULL)
2076 goto error;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002077
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002078 co->pool = pool;
2079 co->indices = indices;
2080 co->result = NULL;
2081 co->r = r;
2082 co->stopped = r > n ? 1 : 0;
2083
2084 return (PyObject *)co;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002085
2086error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002087 if (indices != NULL)
2088 PyMem_Free(indices);
2089 Py_XDECREF(pool);
2090 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002091}
2092
2093static void
2094combinations_dealloc(combinationsobject *co)
2095{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002096 PyObject_GC_UnTrack(co);
2097 Py_XDECREF(co->pool);
2098 Py_XDECREF(co->result);
2099 if (co->indices != NULL)
2100 PyMem_Free(co->indices);
2101 Py_TYPE(co)->tp_free(co);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002102}
2103
2104static int
2105combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2106{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002107 Py_VISIT(co->pool);
2108 Py_VISIT(co->result);
2109 return 0;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002110}
2111
2112static PyObject *
2113combinations_next(combinationsobject *co)
2114{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002115 PyObject *elem;
2116 PyObject *oldelem;
2117 PyObject *pool = co->pool;
2118 Py_ssize_t *indices = co->indices;
2119 PyObject *result = co->result;
2120 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2121 Py_ssize_t r = co->r;
2122 Py_ssize_t i, j, index;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002123
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002124 if (co->stopped)
2125 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002126
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002127 if (result == NULL) {
2128 /* On the first pass, initialize result tuple using the indices */
2129 result = PyTuple_New(r);
2130 if (result == NULL)
2131 goto empty;
2132 co->result = result;
2133 for (i=0; i<r ; i++) {
2134 index = indices[i];
2135 elem = PyTuple_GET_ITEM(pool, index);
2136 Py_INCREF(elem);
2137 PyTuple_SET_ITEM(result, i, elem);
2138 }
2139 } else {
2140 /* Copy the previous result tuple or re-use it if available */
2141 if (Py_REFCNT(result) > 1) {
2142 PyObject *old_result = result;
2143 result = PyTuple_New(r);
2144 if (result == NULL)
2145 goto empty;
2146 co->result = result;
2147 for (i=0; i<r ; i++) {
2148 elem = PyTuple_GET_ITEM(old_result, i);
2149 Py_INCREF(elem);
2150 PyTuple_SET_ITEM(result, i, elem);
2151 }
2152 Py_DECREF(old_result);
2153 }
2154 /* Now, we've got the only copy so we can update it in-place
2155 * CPython's empty tuple is a singleton and cached in
2156 * PyTuple's freelist.
2157 */
2158 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002159
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002160 /* Scan indices right-to-left until finding one that is not
2161 at its maximum (i + n - r). */
2162 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2163 ;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002164
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002165 /* If i is negative, then the indices are all at
2166 their maximum value and we're done. */
2167 if (i < 0)
2168 goto empty;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002169
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002170 /* Increment the current index which we know is not at its
2171 maximum. Then move back to the right setting each index
2172 to its lowest possible value (one higher than the index
2173 to its left -- this maintains the sort order invariant). */
2174 indices[i]++;
2175 for (j=i+1 ; j<r ; j++)
2176 indices[j] = indices[j-1] + 1;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002177
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002178 /* Update the result tuple for the new indices
2179 starting with i, the leftmost index that changed */
2180 for ( ; i<r ; i++) {
2181 index = indices[i];
2182 elem = PyTuple_GET_ITEM(pool, index);
2183 Py_INCREF(elem);
2184 oldelem = PyTuple_GET_ITEM(result, i);
2185 PyTuple_SET_ITEM(result, i, elem);
2186 Py_DECREF(oldelem);
2187 }
2188 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002189
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002190 Py_INCREF(result);
2191 return result;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002192
2193empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002194 co->stopped = 1;
2195 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002196}
2197
2198PyDoc_STRVAR(combinations_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002199"combinations(iterable, r) --> combinations object\n\
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002200\n\
2201Return successive r-length combinations of elements in the iterable.\n\n\
2202combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2203
2204static PyTypeObject combinations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002205 PyVarObject_HEAD_INIT(NULL, 0)
2206 "itertools.combinations", /* tp_name */
2207 sizeof(combinationsobject), /* tp_basicsize */
2208 0, /* tp_itemsize */
2209 /* methods */
2210 (destructor)combinations_dealloc, /* tp_dealloc */
2211 0, /* tp_print */
2212 0, /* tp_getattr */
2213 0, /* tp_setattr */
2214 0, /* tp_compare */
2215 0, /* tp_repr */
2216 0, /* tp_as_number */
2217 0, /* tp_as_sequence */
2218 0, /* tp_as_mapping */
2219 0, /* tp_hash */
2220 0, /* tp_call */
2221 0, /* tp_str */
2222 PyObject_GenericGetAttr, /* tp_getattro */
2223 0, /* tp_setattro */
2224 0, /* tp_as_buffer */
2225 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2226 Py_TPFLAGS_BASETYPE, /* tp_flags */
2227 combinations_doc, /* tp_doc */
2228 (traverseproc)combinations_traverse, /* tp_traverse */
2229 0, /* tp_clear */
2230 0, /* tp_richcompare */
2231 0, /* tp_weaklistoffset */
2232 PyObject_SelfIter, /* tp_iter */
2233 (iternextfunc)combinations_next, /* tp_iternext */
2234 0, /* tp_methods */
2235 0, /* tp_members */
2236 0, /* tp_getset */
2237 0, /* tp_base */
2238 0, /* tp_dict */
2239 0, /* tp_descr_get */
2240 0, /* tp_descr_set */
2241 0, /* tp_dictoffset */
2242 0, /* tp_init */
2243 0, /* tp_alloc */
2244 combinations_new, /* tp_new */
2245 PyObject_GC_Del, /* tp_free */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002246};
2247
2248
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002249/* combinations with replacement object *******************************************/
2250
2251/* Equivalent to:
2252
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002253 def combinations_with_replacement(iterable, r):
2254 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2255 # number items returned: (n+r-1)! / r! / (n-1)!
2256 pool = tuple(iterable)
2257 n = len(pool)
2258 indices = [0] * r
2259 yield tuple(pool[i] for i in indices)
2260 while 1:
2261 for i in reversed(range(r)):
2262 if indices[i] != n - 1:
2263 break
2264 else:
2265 return
2266 indices[i:] = [indices[i] + 1] * (r - i)
2267 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002268
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002269 def combinations_with_replacement2(iterable, r):
2270 'Alternate version that filters from product()'
2271 pool = tuple(iterable)
2272 n = len(pool)
2273 for indices in product(range(n), repeat=r):
2274 if sorted(indices) == list(indices):
2275 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002276*/
2277typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002278 PyObject_HEAD
2279 PyObject *pool; /* input converted to a tuple */
2280 Py_ssize_t *indices; /* one index per result element */
2281 PyObject *result; /* most recently returned result tuple */
2282 Py_ssize_t r; /* size of result tuple */
2283 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002284} cwrobject;
2285
2286static PyTypeObject cwr_type;
2287
2288static PyObject *
2289cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2290{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002291 cwrobject *co;
2292 Py_ssize_t n;
2293 Py_ssize_t r;
2294 PyObject *pool = NULL;
2295 PyObject *iterable = NULL;
2296 Py_ssize_t *indices = NULL;
2297 Py_ssize_t i;
2298 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002299
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002300 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2301 &iterable, &r))
2302 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002303
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002304 pool = PySequence_Tuple(iterable);
2305 if (pool == NULL)
2306 goto error;
2307 n = PyTuple_GET_SIZE(pool);
2308 if (r < 0) {
2309 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2310 goto error;
2311 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002312
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002313 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2314 if (indices == NULL) {
2315 PyErr_NoMemory();
2316 goto error;
2317 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002318
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002319 for (i=0 ; i<r ; i++)
2320 indices[i] = 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002321
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002322 /* create cwrobject structure */
2323 co = (cwrobject *)type->tp_alloc(type, 0);
2324 if (co == NULL)
2325 goto error;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002326
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002327 co->pool = pool;
2328 co->indices = indices;
2329 co->result = NULL;
2330 co->r = r;
2331 co->stopped = !n && r;
2332
2333 return (PyObject *)co;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002334
2335error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002336 if (indices != NULL)
2337 PyMem_Free(indices);
2338 Py_XDECREF(pool);
2339 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002340}
2341
2342static void
2343cwr_dealloc(cwrobject *co)
2344{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002345 PyObject_GC_UnTrack(co);
2346 Py_XDECREF(co->pool);
2347 Py_XDECREF(co->result);
2348 if (co->indices != NULL)
2349 PyMem_Free(co->indices);
2350 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002351}
2352
2353static int
2354cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2355{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002356 Py_VISIT(co->pool);
2357 Py_VISIT(co->result);
2358 return 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002359}
2360
2361static PyObject *
2362cwr_next(cwrobject *co)
2363{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002364 PyObject *elem;
2365 PyObject *oldelem;
2366 PyObject *pool = co->pool;
2367 Py_ssize_t *indices = co->indices;
2368 PyObject *result = co->result;
2369 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2370 Py_ssize_t r = co->r;
2371 Py_ssize_t i, j, index;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002372
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002373 if (co->stopped)
2374 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002375
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002376 if (result == NULL) {
2377 /* On the first pass, initialize result tuple using the indices */
2378 result = PyTuple_New(r);
2379 if (result == NULL)
2380 goto empty;
2381 co->result = result;
2382 for (i=0; i<r ; i++) {
2383 index = indices[i];
2384 elem = PyTuple_GET_ITEM(pool, index);
2385 Py_INCREF(elem);
2386 PyTuple_SET_ITEM(result, i, elem);
2387 }
2388 } else {
2389 /* Copy the previous result tuple or re-use it if available */
2390 if (Py_REFCNT(result) > 1) {
2391 PyObject *old_result = result;
2392 result = PyTuple_New(r);
2393 if (result == NULL)
2394 goto empty;
2395 co->result = result;
2396 for (i=0; i<r ; i++) {
2397 elem = PyTuple_GET_ITEM(old_result, i);
2398 Py_INCREF(elem);
2399 PyTuple_SET_ITEM(result, i, elem);
2400 }
2401 Py_DECREF(old_result);
2402 }
2403 /* Now, we've got the only copy so we can update it in-place CPython's
2404 empty tuple is a singleton and cached in PyTuple's freelist. */
2405 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002406
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002407 /* Scan indices right-to-left until finding one that is not
2408 * at its maximum (n-1). */
2409 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2410 ;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002411
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002412 /* If i is negative, then the indices are all at
2413 their maximum value and we're done. */
2414 if (i < 0)
2415 goto empty;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002416
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002417 /* Increment the current index which we know is not at its
2418 maximum. Then set all to the right to the same value. */
2419 indices[i]++;
2420 for (j=i+1 ; j<r ; j++)
2421 indices[j] = indices[j-1];
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002422
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002423 /* Update the result tuple for the new indices
2424 starting with i, the leftmost index that changed */
2425 for ( ; i<r ; i++) {
2426 index = indices[i];
2427 elem = PyTuple_GET_ITEM(pool, index);
2428 Py_INCREF(elem);
2429 oldelem = PyTuple_GET_ITEM(result, i);
2430 PyTuple_SET_ITEM(result, i, elem);
2431 Py_DECREF(oldelem);
2432 }
2433 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002434
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002435 Py_INCREF(result);
2436 return result;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002437
2438empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002439 co->stopped = 1;
2440 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002441}
2442
2443PyDoc_STRVAR(cwr_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002444"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002445\n\
2446Return successive r-length combinations of elements in the iterable\n\
2447allowing individual elements to have successive repeats.\n\
2448combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2449
2450static PyTypeObject cwr_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002451 PyVarObject_HEAD_INIT(NULL, 0)
2452 "itertools.combinations_with_replacement", /* tp_name */
2453 sizeof(cwrobject), /* tp_basicsize */
2454 0, /* tp_itemsize */
2455 /* methods */
2456 (destructor)cwr_dealloc, /* tp_dealloc */
2457 0, /* tp_print */
2458 0, /* tp_getattr */
2459 0, /* tp_setattr */
2460 0, /* tp_compare */
2461 0, /* tp_repr */
2462 0, /* tp_as_number */
2463 0, /* tp_as_sequence */
2464 0, /* tp_as_mapping */
2465 0, /* tp_hash */
2466 0, /* tp_call */
2467 0, /* tp_str */
2468 PyObject_GenericGetAttr, /* tp_getattro */
2469 0, /* tp_setattro */
2470 0, /* tp_as_buffer */
2471 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2472 Py_TPFLAGS_BASETYPE, /* tp_flags */
2473 cwr_doc, /* tp_doc */
2474 (traverseproc)cwr_traverse, /* tp_traverse */
2475 0, /* tp_clear */
2476 0, /* tp_richcompare */
2477 0, /* tp_weaklistoffset */
2478 PyObject_SelfIter, /* tp_iter */
2479 (iternextfunc)cwr_next, /* tp_iternext */
2480 0, /* tp_methods */
2481 0, /* tp_members */
2482 0, /* tp_getset */
2483 0, /* tp_base */
2484 0, /* tp_dict */
2485 0, /* tp_descr_get */
2486 0, /* tp_descr_set */
2487 0, /* tp_dictoffset */
2488 0, /* tp_init */
2489 0, /* tp_alloc */
2490 cwr_new, /* tp_new */
2491 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002492};
2493
2494
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002495/* permutations object ************************************************************
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002496
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002497def permutations(iterable, r=None):
2498 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2499 pool = tuple(iterable)
2500 n = len(pool)
2501 r = n if r is None else r
2502 indices = range(n)
2503 cycles = range(n-r+1, n+1)[::-1]
2504 yield tuple(pool[i] for i in indices[:r])
2505 while n:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002506 for i in reversed(range(r)):
2507 cycles[i] -= 1
2508 if cycles[i] == 0:
2509 indices[i:] = indices[i+1:] + indices[i:i+1]
2510 cycles[i] = n - i
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002511 else:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002512 j = cycles[i]
2513 indices[i], indices[-j] = indices[-j], indices[i]
2514 yield tuple(pool[i] for i in indices[:r])
2515 break
2516 else:
2517 return
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002518*/
2519
2520typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002521 PyObject_HEAD
2522 PyObject *pool; /* input converted to a tuple */
2523 Py_ssize_t *indices; /* one index per element in the pool */
2524 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2525 PyObject *result; /* most recently returned result tuple */
2526 Py_ssize_t r; /* size of result tuple */
2527 int stopped; /* set to 1 when the permutations iterator is exhausted */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002528} permutationsobject;
2529
2530static PyTypeObject permutations_type;
2531
2532static PyObject *
2533permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2534{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002535 permutationsobject *po;
2536 Py_ssize_t n;
2537 Py_ssize_t r;
2538 PyObject *robj = Py_None;
2539 PyObject *pool = NULL;
2540 PyObject *iterable = NULL;
2541 Py_ssize_t *indices = NULL;
2542 Py_ssize_t *cycles = NULL;
2543 Py_ssize_t i;
2544 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002545
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002546 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2547 &iterable, &robj))
2548 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002549
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002550 pool = PySequence_Tuple(iterable);
2551 if (pool == NULL)
2552 goto error;
2553 n = PyTuple_GET_SIZE(pool);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002554
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002555 r = n;
2556 if (robj != Py_None) {
2557 r = PyInt_AsSsize_t(robj);
2558 if (r == -1 && PyErr_Occurred())
2559 goto error;
2560 }
2561 if (r < 0) {
2562 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2563 goto error;
2564 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002565
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002566 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2567 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2568 if (indices == NULL || cycles == NULL) {
2569 PyErr_NoMemory();
2570 goto error;
2571 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002572
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002573 for (i=0 ; i<n ; i++)
2574 indices[i] = i;
2575 for (i=0 ; i<r ; i++)
2576 cycles[i] = n - i;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002577
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002578 /* create permutationsobject structure */
2579 po = (permutationsobject *)type->tp_alloc(type, 0);
2580 if (po == NULL)
2581 goto error;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002582
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002583 po->pool = pool;
2584 po->indices = indices;
2585 po->cycles = cycles;
2586 po->result = NULL;
2587 po->r = r;
2588 po->stopped = r > n ? 1 : 0;
2589
2590 return (PyObject *)po;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002591
2592error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002593 if (indices != NULL)
2594 PyMem_Free(indices);
2595 if (cycles != NULL)
2596 PyMem_Free(cycles);
2597 Py_XDECREF(pool);
2598 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002599}
2600
2601static void
2602permutations_dealloc(permutationsobject *po)
2603{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002604 PyObject_GC_UnTrack(po);
2605 Py_XDECREF(po->pool);
2606 Py_XDECREF(po->result);
2607 PyMem_Free(po->indices);
2608 PyMem_Free(po->cycles);
2609 Py_TYPE(po)->tp_free(po);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002610}
2611
2612static int
2613permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2614{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002615 Py_VISIT(po->pool);
2616 Py_VISIT(po->result);
2617 return 0;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002618}
2619
2620static PyObject *
2621permutations_next(permutationsobject *po)
2622{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002623 PyObject *elem;
2624 PyObject *oldelem;
2625 PyObject *pool = po->pool;
2626 Py_ssize_t *indices = po->indices;
2627 Py_ssize_t *cycles = po->cycles;
2628 PyObject *result = po->result;
2629 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2630 Py_ssize_t r = po->r;
2631 Py_ssize_t i, j, k, index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002632
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002633 if (po->stopped)
2634 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002635
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002636 if (result == NULL) {
2637 /* On the first pass, initialize result tuple using the indices */
2638 result = PyTuple_New(r);
2639 if (result == NULL)
2640 goto empty;
2641 po->result = result;
2642 for (i=0; i<r ; i++) {
2643 index = indices[i];
2644 elem = PyTuple_GET_ITEM(pool, index);
2645 Py_INCREF(elem);
2646 PyTuple_SET_ITEM(result, i, elem);
2647 }
2648 } else {
2649 if (n == 0)
2650 goto empty;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002651
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002652 /* Copy the previous result tuple or re-use it if available */
2653 if (Py_REFCNT(result) > 1) {
2654 PyObject *old_result = result;
2655 result = PyTuple_New(r);
2656 if (result == NULL)
2657 goto empty;
2658 po->result = result;
2659 for (i=0; i<r ; i++) {
2660 elem = PyTuple_GET_ITEM(old_result, i);
2661 Py_INCREF(elem);
2662 PyTuple_SET_ITEM(result, i, elem);
2663 }
2664 Py_DECREF(old_result);
2665 }
2666 /* Now, we've got the only copy so we can update it in-place */
2667 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002668
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002669 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2670 for (i=r-1 ; i>=0 ; i--) {
2671 cycles[i] -= 1;
2672 if (cycles[i] == 0) {
2673 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2674 index = indices[i];
2675 for (j=i ; j<n-1 ; j++)
2676 indices[j] = indices[j+1];
2677 indices[n-1] = index;
2678 cycles[i] = n - i;
2679 } else {
2680 j = cycles[i];
2681 index = indices[i];
2682 indices[i] = indices[n-j];
2683 indices[n-j] = index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002684
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002685 for (k=i; k<r ; k++) {
2686 /* start with i, the leftmost element that changed */
2687 /* yield tuple(pool[k] for k in indices[:r]) */
2688 index = indices[k];
2689 elem = PyTuple_GET_ITEM(pool, index);
2690 Py_INCREF(elem);
2691 oldelem = PyTuple_GET_ITEM(result, k);
2692 PyTuple_SET_ITEM(result, k, elem);
2693 Py_DECREF(oldelem);
2694 }
2695 break;
2696 }
2697 }
2698 /* If i is negative, then the cycles have all
2699 rolled-over and we're done. */
2700 if (i < 0)
2701 goto empty;
2702 }
2703 Py_INCREF(result);
2704 return result;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002705
2706empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002707 po->stopped = 1;
2708 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002709}
2710
2711PyDoc_STRVAR(permutations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +00002712"permutations(iterable[, r]) --> permutations object\n\
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002713\n\
2714Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl8550d692008-06-10 12:46:39 +00002715permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002716
2717static PyTypeObject permutations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002718 PyVarObject_HEAD_INIT(NULL, 0)
2719 "itertools.permutations", /* tp_name */
2720 sizeof(permutationsobject), /* tp_basicsize */
2721 0, /* tp_itemsize */
2722 /* methods */
2723 (destructor)permutations_dealloc, /* tp_dealloc */
2724 0, /* tp_print */
2725 0, /* tp_getattr */
2726 0, /* tp_setattr */
2727 0, /* tp_compare */
2728 0, /* tp_repr */
2729 0, /* tp_as_number */
2730 0, /* tp_as_sequence */
2731 0, /* tp_as_mapping */
2732 0, /* tp_hash */
2733 0, /* tp_call */
2734 0, /* tp_str */
2735 PyObject_GenericGetAttr, /* tp_getattro */
2736 0, /* tp_setattro */
2737 0, /* tp_as_buffer */
2738 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2739 Py_TPFLAGS_BASETYPE, /* tp_flags */
2740 permutations_doc, /* tp_doc */
2741 (traverseproc)permutations_traverse, /* tp_traverse */
2742 0, /* tp_clear */
2743 0, /* tp_richcompare */
2744 0, /* tp_weaklistoffset */
2745 PyObject_SelfIter, /* tp_iter */
2746 (iternextfunc)permutations_next, /* tp_iternext */
2747 0, /* tp_methods */
2748 0, /* tp_members */
2749 0, /* tp_getset */
2750 0, /* tp_base */
2751 0, /* tp_dict */
2752 0, /* tp_descr_get */
2753 0, /* tp_descr_set */
2754 0, /* tp_dictoffset */
2755 0, /* tp_init */
2756 0, /* tp_alloc */
2757 permutations_new, /* tp_new */
2758 PyObject_GC_Del, /* tp_free */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002759};
2760
2761
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002762/* compress object ************************************************************/
2763
2764/* Equivalent to:
2765
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002766 def compress(data, selectors):
2767 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2768 return (d for d, s in izip(data, selectors) if s)
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002769*/
2770
2771typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002772 PyObject_HEAD
2773 PyObject *data;
2774 PyObject *selectors;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002775} compressobject;
2776
2777static PyTypeObject compress_type;
2778
2779static PyObject *
2780compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2781{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002782 PyObject *seq1, *seq2;
2783 PyObject *data=NULL, *selectors=NULL;
2784 compressobject *lz;
2785 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002786
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002787 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2788 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002789
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002790 data = PyObject_GetIter(seq1);
2791 if (data == NULL)
2792 goto fail;
2793 selectors = PyObject_GetIter(seq2);
2794 if (selectors == NULL)
2795 goto fail;
2796
2797 /* create compressobject structure */
2798 lz = (compressobject *)type->tp_alloc(type, 0);
2799 if (lz == NULL)
2800 goto fail;
2801 lz->data = data;
2802 lz->selectors = selectors;
2803 return (PyObject *)lz;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002804
2805fail:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002806 Py_XDECREF(data);
2807 Py_XDECREF(selectors);
2808 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002809}
2810
2811static void
2812compress_dealloc(compressobject *lz)
2813{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002814 PyObject_GC_UnTrack(lz);
2815 Py_XDECREF(lz->data);
2816 Py_XDECREF(lz->selectors);
2817 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002818}
2819
2820static int
2821compress_traverse(compressobject *lz, visitproc visit, void *arg)
2822{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002823 Py_VISIT(lz->data);
2824 Py_VISIT(lz->selectors);
2825 return 0;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002826}
2827
2828static PyObject *
2829compress_next(compressobject *lz)
2830{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002831 PyObject *data = lz->data, *selectors = lz->selectors;
2832 PyObject *datum, *selector;
2833 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2834 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2835 int ok;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002836
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002837 while (1) {
2838 /* Steps: get datum, get selector, evaluate selector.
2839 Order is important (to match the pure python version
2840 in terms of which input gets a chance to raise an
2841 exception first).
2842 */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002843
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002844 datum = datanext(data);
2845 if (datum == NULL)
2846 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002847
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002848 selector = selectornext(selectors);
2849 if (selector == NULL) {
2850 Py_DECREF(datum);
2851 return NULL;
2852 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002853
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002854 ok = PyObject_IsTrue(selector);
2855 Py_DECREF(selector);
2856 if (ok == 1)
2857 return datum;
2858 Py_DECREF(datum);
2859 if (ok == -1)
2860 return NULL;
2861 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002862}
2863
2864PyDoc_STRVAR(compress_doc,
Raymond Hettinger2e2909f2009-02-19 02:15:14 +00002865"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002866\n\
2867Return data elements corresponding to true selector elements.\n\
2868Forms a shorter iterator from selected data elements using the\n\
2869selectors to choose the data elements.");
2870
2871static PyTypeObject compress_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002872 PyVarObject_HEAD_INIT(NULL, 0)
2873 "itertools.compress", /* tp_name */
2874 sizeof(compressobject), /* tp_basicsize */
2875 0, /* tp_itemsize */
2876 /* methods */
2877 (destructor)compress_dealloc, /* tp_dealloc */
2878 0, /* tp_print */
2879 0, /* tp_getattr */
2880 0, /* tp_setattr */
2881 0, /* tp_compare */
2882 0, /* tp_repr */
2883 0, /* tp_as_number */
2884 0, /* tp_as_sequence */
2885 0, /* tp_as_mapping */
2886 0, /* tp_hash */
2887 0, /* tp_call */
2888 0, /* tp_str */
2889 PyObject_GenericGetAttr, /* tp_getattro */
2890 0, /* tp_setattro */
2891 0, /* tp_as_buffer */
2892 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2893 Py_TPFLAGS_BASETYPE, /* tp_flags */
2894 compress_doc, /* tp_doc */
2895 (traverseproc)compress_traverse, /* tp_traverse */
2896 0, /* tp_clear */
2897 0, /* tp_richcompare */
2898 0, /* tp_weaklistoffset */
2899 PyObject_SelfIter, /* tp_iter */
2900 (iternextfunc)compress_next, /* tp_iternext */
2901 0, /* tp_methods */
2902 0, /* tp_members */
2903 0, /* tp_getset */
2904 0, /* tp_base */
2905 0, /* tp_dict */
2906 0, /* tp_descr_get */
2907 0, /* tp_descr_set */
2908 0, /* tp_dictoffset */
2909 0, /* tp_init */
2910 0, /* tp_alloc */
2911 compress_new, /* tp_new */
2912 PyObject_GC_Del, /* tp_free */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002913};
2914
2915
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002916/* ifilter object ************************************************************/
2917
2918typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002919 PyObject_HEAD
2920 PyObject *func;
2921 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002922} ifilterobject;
2923
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002924static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002925
2926static PyObject *
2927ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2928{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002929 PyObject *func, *seq;
2930 PyObject *it;
2931 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002932
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002933 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
2934 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002935
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002936 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
2937 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002938
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002939 /* Get iterator. */
2940 it = PyObject_GetIter(seq);
2941 if (it == NULL)
2942 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002943
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002944 /* create ifilterobject structure */
2945 lz = (ifilterobject *)type->tp_alloc(type, 0);
2946 if (lz == NULL) {
2947 Py_DECREF(it);
2948 return NULL;
2949 }
2950 Py_INCREF(func);
2951 lz->func = func;
2952 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002953
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002954 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002955}
2956
2957static void
2958ifilter_dealloc(ifilterobject *lz)
2959{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002960 PyObject_GC_UnTrack(lz);
2961 Py_XDECREF(lz->func);
2962 Py_XDECREF(lz->it);
2963 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002964}
2965
2966static int
2967ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
2968{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002969 Py_VISIT(lz->it);
2970 Py_VISIT(lz->func);
2971 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002972}
2973
2974static PyObject *
2975ifilter_next(ifilterobject *lz)
2976{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002977 PyObject *item;
2978 PyObject *it = lz->it;
2979 long ok;
2980 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002981
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002982 iternext = *Py_TYPE(it)->tp_iternext;
2983 for (;;) {
2984 item = iternext(it);
2985 if (item == NULL)
2986 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002987
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002988 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2989 ok = PyObject_IsTrue(item);
2990 } else {
2991 PyObject *good;
2992 good = PyObject_CallFunctionObjArgs(lz->func,
2993 item, NULL);
2994 if (good == NULL) {
2995 Py_DECREF(item);
2996 return NULL;
2997 }
2998 ok = PyObject_IsTrue(good);
2999 Py_DECREF(good);
3000 }
3001 if (ok)
3002 return item;
3003 Py_DECREF(item);
3004 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003005}
3006
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003007PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003008"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003009\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003010Return those items of sequence for which function(item) is true.\n\
3011If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003012
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003013static PyTypeObject ifilter_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003014 PyVarObject_HEAD_INIT(NULL, 0)
3015 "itertools.ifilter", /* tp_name */
3016 sizeof(ifilterobject), /* tp_basicsize */
3017 0, /* tp_itemsize */
3018 /* methods */
3019 (destructor)ifilter_dealloc, /* tp_dealloc */
3020 0, /* tp_print */
3021 0, /* tp_getattr */
3022 0, /* tp_setattr */
3023 0, /* tp_compare */
3024 0, /* tp_repr */
3025 0, /* tp_as_number */
3026 0, /* tp_as_sequence */
3027 0, /* tp_as_mapping */
3028 0, /* tp_hash */
3029 0, /* tp_call */
3030 0, /* tp_str */
3031 PyObject_GenericGetAttr, /* tp_getattro */
3032 0, /* tp_setattro */
3033 0, /* tp_as_buffer */
3034 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3035 Py_TPFLAGS_BASETYPE, /* tp_flags */
3036 ifilter_doc, /* tp_doc */
3037 (traverseproc)ifilter_traverse, /* tp_traverse */
3038 0, /* tp_clear */
3039 0, /* tp_richcompare */
3040 0, /* tp_weaklistoffset */
3041 PyObject_SelfIter, /* tp_iter */
3042 (iternextfunc)ifilter_next, /* tp_iternext */
3043 0, /* tp_methods */
3044 0, /* tp_members */
3045 0, /* tp_getset */
3046 0, /* tp_base */
3047 0, /* tp_dict */
3048 0, /* tp_descr_get */
3049 0, /* tp_descr_set */
3050 0, /* tp_dictoffset */
3051 0, /* tp_init */
3052 0, /* tp_alloc */
3053 ifilter_new, /* tp_new */
3054 PyObject_GC_Del, /* tp_free */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003055};
3056
3057
3058/* ifilterfalse object ************************************************************/
3059
3060typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003061 PyObject_HEAD
3062 PyObject *func;
3063 PyObject *it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003064} ifilterfalseobject;
3065
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003066static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003067
3068static PyObject *
3069ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3070{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003071 PyObject *func, *seq;
3072 PyObject *it;
3073 ifilterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003074
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003075 if (type == &ifilterfalse_type &&
3076 !_PyArg_NoKeywords("ifilterfalse()", kwds))
3077 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003078
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003079 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
3080 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003081
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003082 /* Get iterator. */
3083 it = PyObject_GetIter(seq);
3084 if (it == NULL)
3085 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003086
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003087 /* create ifilterfalseobject structure */
3088 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
3089 if (lz == NULL) {
3090 Py_DECREF(it);
3091 return NULL;
3092 }
3093 Py_INCREF(func);
3094 lz->func = func;
3095 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003096
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003097 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003098}
3099
3100static void
3101ifilterfalse_dealloc(ifilterfalseobject *lz)
3102{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003103 PyObject_GC_UnTrack(lz);
3104 Py_XDECREF(lz->func);
3105 Py_XDECREF(lz->it);
3106 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003107}
3108
3109static int
3110ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
3111{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003112 Py_VISIT(lz->it);
3113 Py_VISIT(lz->func);
3114 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003115}
3116
3117static PyObject *
3118ifilterfalse_next(ifilterfalseobject *lz)
3119{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003120 PyObject *item;
3121 PyObject *it = lz->it;
3122 long ok;
3123 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003124
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003125 iternext = *Py_TYPE(it)->tp_iternext;
3126 for (;;) {
3127 item = iternext(it);
3128 if (item == NULL)
3129 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003130
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003131 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3132 ok = PyObject_IsTrue(item);
3133 } else {
3134 PyObject *good;
3135 good = PyObject_CallFunctionObjArgs(lz->func,
3136 item, NULL);
3137 if (good == NULL) {
3138 Py_DECREF(item);
3139 return NULL;
3140 }
3141 ok = PyObject_IsTrue(good);
3142 Py_DECREF(good);
3143 }
3144 if (!ok)
3145 return item;
3146 Py_DECREF(item);
3147 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003148}
3149
Raymond Hettinger60eca932003-02-09 06:40:58 +00003150PyDoc_STRVAR(ifilterfalse_doc,
3151"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
3152\n\
3153Return those items of sequence for which function(item) is false.\n\
3154If function is None, return the items that are false.");
3155
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003156static PyTypeObject ifilterfalse_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003157 PyVarObject_HEAD_INIT(NULL, 0)
3158 "itertools.ifilterfalse", /* tp_name */
3159 sizeof(ifilterfalseobject), /* tp_basicsize */
3160 0, /* tp_itemsize */
3161 /* methods */
3162 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
3163 0, /* tp_print */
3164 0, /* tp_getattr */
3165 0, /* tp_setattr */
3166 0, /* tp_compare */
3167 0, /* tp_repr */
3168 0, /* tp_as_number */
3169 0, /* tp_as_sequence */
3170 0, /* tp_as_mapping */
3171 0, /* tp_hash */
3172 0, /* tp_call */
3173 0, /* tp_str */
3174 PyObject_GenericGetAttr, /* tp_getattro */
3175 0, /* tp_setattro */
3176 0, /* tp_as_buffer */
3177 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3178 Py_TPFLAGS_BASETYPE, /* tp_flags */
3179 ifilterfalse_doc, /* tp_doc */
3180 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
3181 0, /* tp_clear */
3182 0, /* tp_richcompare */
3183 0, /* tp_weaklistoffset */
3184 PyObject_SelfIter, /* tp_iter */
3185 (iternextfunc)ifilterfalse_next, /* tp_iternext */
3186 0, /* tp_methods */
3187 0, /* tp_members */
3188 0, /* tp_getset */
3189 0, /* tp_base */
3190 0, /* tp_dict */
3191 0, /* tp_descr_get */
3192 0, /* tp_descr_set */
3193 0, /* tp_dictoffset */
3194 0, /* tp_init */
3195 0, /* tp_alloc */
3196 ifilterfalse_new, /* tp_new */
3197 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003198};
3199
3200
3201/* count object ************************************************************/
3202
3203typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003204 PyObject_HEAD
3205 Py_ssize_t cnt;
3206 PyObject *long_cnt;
3207 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003208} countobject;
3209
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003210/* Counting logic and invariants:
3211
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003212fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003213
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003214 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3215 Advances with: cnt += 1
3216 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003217
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003218slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003219
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003220 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3221 All counting is done with python objects (no overflows or underflows).
3222 Advances with: long_cnt += long_step
3223 Step may be zero -- effectively a slow version of repeat(cnt).
3224 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003225*/
3226
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003227static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003228
3229static PyObject *
3230count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3231{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003232 countobject *lz;
3233 int slow_mode = 0;
3234 Py_ssize_t cnt = 0;
3235 PyObject *long_cnt = NULL;
3236 PyObject *long_step = NULL;
3237 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003238
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003239 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3240 kwlist, &long_cnt, &long_step))
3241 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003242
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003243 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3244 (long_step != NULL && !PyNumber_Check(long_step))) {
3245 PyErr_SetString(PyExc_TypeError, "a number is required");
3246 return NULL;
3247 }
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003248
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003249 if (long_cnt != NULL) {
3250 cnt = PyInt_AsSsize_t(long_cnt);
3251 if ((cnt == -1 && PyErr_Occurred()) || !PyInt_Check(long_cnt)) {
3252 PyErr_Clear();
3253 slow_mode = 1;
3254 }
3255 Py_INCREF(long_cnt);
3256 } else {
3257 cnt = 0;
3258 long_cnt = PyInt_FromLong(0);
3259 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003260
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003261 /* If not specified, step defaults to 1 */
3262 if (long_step == NULL) {
3263 long_step = PyInt_FromLong(1);
3264 if (long_step == NULL) {
3265 Py_DECREF(long_cnt);
3266 return NULL;
3267 }
3268 } else
3269 Py_INCREF(long_step);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003270
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003271 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003272
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003273 /* Fast mode only works when the step is 1 */
3274 if (!PyInt_Check(long_step) ||
3275 PyInt_AS_LONG(long_step) != 1) {
3276 slow_mode = 1;
3277 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003278
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003279 if (slow_mode)
3280 cnt = PY_SSIZE_T_MAX;
3281 else
3282 Py_CLEAR(long_cnt);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003283
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003284 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3285 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3286 assert(slow_mode ||
3287 (PyInt_Check(long_step) && PyInt_AS_LONG(long_step) == 1));
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003288
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003289 /* create countobject structure */
3290 lz = (countobject *)type->tp_alloc(type, 0);
3291 if (lz == NULL) {
3292 Py_XDECREF(long_cnt);
3293 return NULL;
3294 }
3295 lz->cnt = cnt;
3296 lz->long_cnt = long_cnt;
3297 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003298
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003299 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003300}
3301
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003302static void
3303count_dealloc(countobject *lz)
3304{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003305 PyObject_GC_UnTrack(lz);
3306 Py_XDECREF(lz->long_cnt);
3307 Py_XDECREF(lz->long_step);
3308 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003309}
3310
Benjamin Peterson062a7c32009-02-16 21:07:52 +00003311static int
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003312count_traverse(countobject *lz, visitproc visit, void *arg)
3313{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003314 Py_VISIT(lz->long_cnt);
3315 Py_VISIT(lz->long_step);
3316 return 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003317}
3318
3319static PyObject *
3320count_nextlong(countobject *lz)
3321{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003322 PyObject *long_cnt;
3323 PyObject *stepped_up;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003324
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003325 long_cnt = lz->long_cnt;
3326 if (long_cnt == NULL) {
3327 /* Switch to slow_mode */
3328 long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
3329 if (long_cnt == NULL)
3330 return NULL;
3331 }
3332 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003333
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003334 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3335 if (stepped_up == NULL)
3336 return NULL;
3337 lz->long_cnt = stepped_up;
3338 return long_cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003339}
3340
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003341static PyObject *
3342count_next(countobject *lz)
3343{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003344 if (lz->cnt == PY_SSIZE_T_MAX)
3345 return count_nextlong(lz);
3346 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003347}
3348
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003349static PyObject *
3350count_repr(countobject *lz)
3351{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003352 PyObject *cnt_repr, *step_repr = NULL;
3353 PyObject *result = NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003354
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003355 if (lz->cnt != PY_SSIZE_T_MAX)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003356 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003357
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003358 cnt_repr = PyObject_Repr(lz->long_cnt);
3359 if (cnt_repr == NULL)
3360 return NULL;
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003361
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003362 if (PyInt_Check(lz->long_step) && PyInt_AS_LONG(lz->long_step) == 1) {
3363 /* Don't display step when it is an integer equal to 1 */
3364 result = PyString_FromFormat("count(%s)",
3365 PyString_AS_STRING(cnt_repr));
3366 } else {
3367 step_repr = PyObject_Repr(lz->long_step);
3368 if (step_repr != NULL)
3369 result = PyString_FromFormat("count(%s, %s)",
3370 PyString_AS_STRING(cnt_repr),
3371 PyString_AS_STRING(step_repr));
3372 }
3373 Py_DECREF(cnt_repr);
3374 Py_XDECREF(step_repr);
3375 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003376}
3377
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003378static PyObject *
3379count_reduce(countobject *lz)
3380{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003381 if (lz->cnt == PY_SSIZE_T_MAX)
3382 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3383 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003384}
3385
3386PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3387
3388static PyMethodDef count_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003389 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3390 count_reduce_doc},
3391 {NULL, NULL} /* sentinel */
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003392};
3393
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003394PyDoc_STRVAR(count_doc,
Ezio Melottifdc1e0d2010-06-11 02:21:25 +00003395 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003396\n\
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003397Return a count object whose .next() method returns consecutive values.\n\
3398Equivalent to:\n\n\
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003399 def count(firstval=0, step=1):\n\
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003400 x = firstval\n\
3401 while 1:\n\
3402 yield x\n\
3403 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003404
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003405static PyTypeObject count_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003406 PyVarObject_HEAD_INIT(NULL, 0)
3407 "itertools.count", /* tp_name */
3408 sizeof(countobject), /* tp_basicsize */
3409 0, /* tp_itemsize */
3410 /* methods */
3411 (destructor)count_dealloc, /* tp_dealloc */
3412 0, /* tp_print */
3413 0, /* tp_getattr */
3414 0, /* tp_setattr */
3415 0, /* tp_compare */
3416 (reprfunc)count_repr, /* tp_repr */
3417 0, /* tp_as_number */
3418 0, /* tp_as_sequence */
3419 0, /* tp_as_mapping */
3420 0, /* tp_hash */
3421 0, /* tp_call */
3422 0, /* tp_str */
3423 PyObject_GenericGetAttr, /* tp_getattro */
3424 0, /* tp_setattro */
3425 0, /* tp_as_buffer */
3426 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3427 Py_TPFLAGS_BASETYPE, /* tp_flags */
3428 count_doc, /* tp_doc */
3429 (traverseproc)count_traverse, /* tp_traverse */
3430 0, /* tp_clear */
3431 0, /* tp_richcompare */
3432 0, /* tp_weaklistoffset */
3433 PyObject_SelfIter, /* tp_iter */
3434 (iternextfunc)count_next, /* tp_iternext */
3435 count_methods, /* tp_methods */
3436 0, /* tp_members */
3437 0, /* tp_getset */
3438 0, /* tp_base */
3439 0, /* tp_dict */
3440 0, /* tp_descr_get */
3441 0, /* tp_descr_set */
3442 0, /* tp_dictoffset */
3443 0, /* tp_init */
3444 0, /* tp_alloc */
3445 count_new, /* tp_new */
3446 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003447};
3448
3449
3450/* izip object ************************************************************/
3451
3452#include "Python.h"
3453
3454typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003455 PyObject_HEAD
3456 Py_ssize_t tuplesize;
3457 PyObject *ittuple; /* tuple of iterators */
3458 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003459} izipobject;
3460
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003461static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003462
3463static PyObject *
3464izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3465{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003466 izipobject *lz;
3467 Py_ssize_t i;
3468 PyObject *ittuple; /* tuple of iterators */
3469 PyObject *result;
3470 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003471
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003472 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
3473 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003474
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003475 /* args must be a tuple */
3476 assert(PyTuple_Check(args));
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003477
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003478 /* obtain iterators */
3479 ittuple = PyTuple_New(tuplesize);
3480 if (ittuple == NULL)
3481 return NULL;
3482 for (i=0; i < tuplesize; ++i) {
3483 PyObject *item = PyTuple_GET_ITEM(args, i);
3484 PyObject *it = PyObject_GetIter(item);
3485 if (it == NULL) {
3486 if (PyErr_ExceptionMatches(PyExc_TypeError))
3487 PyErr_Format(PyExc_TypeError,
3488 "izip argument #%zd must support iteration",
3489 i+1);
3490 Py_DECREF(ittuple);
3491 return NULL;
3492 }
3493 PyTuple_SET_ITEM(ittuple, i, it);
3494 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003495
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003496 /* create a result holder */
3497 result = PyTuple_New(tuplesize);
3498 if (result == NULL) {
3499 Py_DECREF(ittuple);
3500 return NULL;
3501 }
3502 for (i=0 ; i < tuplesize ; i++) {
3503 Py_INCREF(Py_None);
3504 PyTuple_SET_ITEM(result, i, Py_None);
3505 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003506
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003507 /* create izipobject structure */
3508 lz = (izipobject *)type->tp_alloc(type, 0);
3509 if (lz == NULL) {
3510 Py_DECREF(ittuple);
3511 Py_DECREF(result);
3512 return NULL;
3513 }
3514 lz->ittuple = ittuple;
3515 lz->tuplesize = tuplesize;
3516 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003517
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003518 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003519}
3520
3521static void
3522izip_dealloc(izipobject *lz)
3523{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003524 PyObject_GC_UnTrack(lz);
3525 Py_XDECREF(lz->ittuple);
3526 Py_XDECREF(lz->result);
3527 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003528}
3529
3530static int
3531izip_traverse(izipobject *lz, visitproc visit, void *arg)
3532{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003533 Py_VISIT(lz->ittuple);
3534 Py_VISIT(lz->result);
3535 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003536}
3537
3538static PyObject *
3539izip_next(izipobject *lz)
3540{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003541 Py_ssize_t i;
3542 Py_ssize_t tuplesize = lz->tuplesize;
3543 PyObject *result = lz->result;
3544 PyObject *it;
3545 PyObject *item;
3546 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003547
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003548 if (tuplesize == 0)
3549 return NULL;
3550 if (Py_REFCNT(result) == 1) {
3551 Py_INCREF(result);
3552 for (i=0 ; i < tuplesize ; i++) {
3553 it = PyTuple_GET_ITEM(lz->ittuple, i);
3554 item = (*Py_TYPE(it)->tp_iternext)(it);
3555 if (item == NULL) {
3556 Py_DECREF(result);
3557 return NULL;
3558 }
3559 olditem = PyTuple_GET_ITEM(result, i);
3560 PyTuple_SET_ITEM(result, i, item);
3561 Py_DECREF(olditem);
3562 }
3563 } else {
3564 result = PyTuple_New(tuplesize);
3565 if (result == NULL)
3566 return NULL;
3567 for (i=0 ; i < tuplesize ; i++) {
3568 it = PyTuple_GET_ITEM(lz->ittuple, i);
3569 item = (*Py_TYPE(it)->tp_iternext)(it);
3570 if (item == NULL) {
3571 Py_DECREF(result);
3572 return NULL;
3573 }
3574 PyTuple_SET_ITEM(result, i, item);
3575 }
3576 }
3577 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003578}
3579
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003580PyDoc_STRVAR(izip_doc,
3581"izip(iter1 [,iter2 [...]]) --> izip object\n\
3582\n\
3583Return a izip object whose .next() method returns a tuple where\n\
3584the i-th element comes from the i-th iterable argument. The .next()\n\
3585method continues until the shortest iterable in the argument sequence\n\
3586is exhausted and then it raises StopIteration. Works like the zip()\n\
3587function but consumes less memory by returning an iterator instead of\n\
3588a list.");
3589
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003590static PyTypeObject izip_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003591 PyVarObject_HEAD_INIT(NULL, 0)
3592 "itertools.izip", /* tp_name */
3593 sizeof(izipobject), /* tp_basicsize */
3594 0, /* tp_itemsize */
3595 /* methods */
3596 (destructor)izip_dealloc, /* tp_dealloc */
3597 0, /* tp_print */
3598 0, /* tp_getattr */
3599 0, /* tp_setattr */
3600 0, /* tp_compare */
3601 0, /* tp_repr */
3602 0, /* tp_as_number */
3603 0, /* tp_as_sequence */
3604 0, /* tp_as_mapping */
3605 0, /* tp_hash */
3606 0, /* tp_call */
3607 0, /* tp_str */
3608 PyObject_GenericGetAttr, /* tp_getattro */
3609 0, /* tp_setattro */
3610 0, /* tp_as_buffer */
3611 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3612 Py_TPFLAGS_BASETYPE, /* tp_flags */
3613 izip_doc, /* tp_doc */
3614 (traverseproc)izip_traverse, /* tp_traverse */
3615 0, /* tp_clear */
3616 0, /* tp_richcompare */
3617 0, /* tp_weaklistoffset */
3618 PyObject_SelfIter, /* tp_iter */
3619 (iternextfunc)izip_next, /* tp_iternext */
3620 0, /* tp_methods */
3621 0, /* tp_members */
3622 0, /* tp_getset */
3623 0, /* tp_base */
3624 0, /* tp_dict */
3625 0, /* tp_descr_get */
3626 0, /* tp_descr_set */
3627 0, /* tp_dictoffset */
3628 0, /* tp_init */
3629 0, /* tp_alloc */
3630 izip_new, /* tp_new */
3631 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003632};
3633
3634
3635/* repeat object ************************************************************/
3636
3637typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003638 PyObject_HEAD
3639 PyObject *element;
3640 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003641} repeatobject;
3642
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003643static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003644
3645static PyObject *
3646repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3647{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003648 repeatobject *ro;
3649 PyObject *element;
3650 Py_ssize_t cnt = -1;
3651 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003652
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003653 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3654 &element, &cnt))
3655 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003656
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003657 if (PyTuple_Size(args) == 2 && cnt < 0)
3658 cnt = 0;
3659
3660 ro = (repeatobject *)type->tp_alloc(type, 0);
3661 if (ro == NULL)
3662 return NULL;
3663 Py_INCREF(element);
3664 ro->element = element;
3665 ro->cnt = cnt;
3666 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003667}
3668
3669static void
3670repeat_dealloc(repeatobject *ro)
3671{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003672 PyObject_GC_UnTrack(ro);
3673 Py_XDECREF(ro->element);
3674 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003675}
3676
3677static int
3678repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3679{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003680 Py_VISIT(ro->element);
3681 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003682}
3683
3684static PyObject *
3685repeat_next(repeatobject *ro)
3686{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003687 if (ro->cnt == 0)
3688 return NULL;
3689 if (ro->cnt > 0)
3690 ro->cnt--;
3691 Py_INCREF(ro->element);
3692 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003693}
3694
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003695static PyObject *
3696repeat_repr(repeatobject *ro)
3697{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003698 PyObject *result, *objrepr;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003699
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003700 objrepr = PyObject_Repr(ro->element);
3701 if (objrepr == NULL)
3702 return NULL;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003703
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003704 if (ro->cnt == -1)
3705 result = PyString_FromFormat("repeat(%s)",
3706 PyString_AS_STRING(objrepr));
3707 else
3708 result = PyString_FromFormat("repeat(%s, %zd)",
3709 PyString_AS_STRING(objrepr), ro->cnt);
3710 Py_DECREF(objrepr);
3711 return result;
3712}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003713
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003714static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003715repeat_len(repeatobject *ro)
3716{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003717 if (ro->cnt == -1) {
3718 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3719 return NULL;
3720 }
3721 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003722}
3723
Armin Rigof5b3e362006-02-11 21:32:43 +00003724PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003725
3726static PyMethodDef repeat_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003727 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3728 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003729};
3730
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003731PyDoc_STRVAR(repeat_doc,
Raymond Hettinger182edae2009-02-19 02:38:25 +00003732"repeat(object [,times]) -> create an iterator which returns the object\n\
3733for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003734endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003735
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003736static PyTypeObject repeat_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003737 PyVarObject_HEAD_INIT(NULL, 0)
3738 "itertools.repeat", /* tp_name */
3739 sizeof(repeatobject), /* tp_basicsize */
3740 0, /* tp_itemsize */
3741 /* methods */
3742 (destructor)repeat_dealloc, /* tp_dealloc */
3743 0, /* tp_print */
3744 0, /* tp_getattr */
3745 0, /* tp_setattr */
3746 0, /* tp_compare */
3747 (reprfunc)repeat_repr, /* tp_repr */
3748 0, /* tp_as_number */
3749 0, /* tp_as_sequence */
3750 0, /* tp_as_mapping */
3751 0, /* tp_hash */
3752 0, /* tp_call */
3753 0, /* tp_str */
3754 PyObject_GenericGetAttr, /* tp_getattro */
3755 0, /* tp_setattro */
3756 0, /* tp_as_buffer */
3757 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3758 Py_TPFLAGS_BASETYPE, /* tp_flags */
3759 repeat_doc, /* tp_doc */
3760 (traverseproc)repeat_traverse, /* tp_traverse */
3761 0, /* tp_clear */
3762 0, /* tp_richcompare */
3763 0, /* tp_weaklistoffset */
3764 PyObject_SelfIter, /* tp_iter */
3765 (iternextfunc)repeat_next, /* tp_iternext */
3766 repeat_methods, /* tp_methods */
3767 0, /* tp_members */
3768 0, /* tp_getset */
3769 0, /* tp_base */
3770 0, /* tp_dict */
3771 0, /* tp_descr_get */
3772 0, /* tp_descr_set */
3773 0, /* tp_dictoffset */
3774 0, /* tp_init */
3775 0, /* tp_alloc */
3776 repeat_new, /* tp_new */
3777 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003778};
3779
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003780/* iziplongest object ************************************************************/
3781
3782#include "Python.h"
3783
3784typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003785 PyObject_HEAD
3786 Py_ssize_t tuplesize;
3787 Py_ssize_t numactive;
3788 PyObject *ittuple; /* tuple of iterators */
3789 PyObject *result;
3790 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003791} iziplongestobject;
3792
3793static PyTypeObject iziplongest_type;
3794
3795static PyObject *
3796izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3797{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003798 iziplongestobject *lz;
3799 Py_ssize_t i;
3800 PyObject *ittuple; /* tuple of iterators */
3801 PyObject *result;
3802 PyObject *fillvalue = Py_None;
3803 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003804
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003805 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3806 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3807 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3808 PyErr_SetString(PyExc_TypeError,
3809 "izip_longest() got an unexpected keyword argument");
3810 return NULL;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003811 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003812 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003813
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003814 /* args must be a tuple */
3815 assert(PyTuple_Check(args));
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003816
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003817 /* obtain iterators */
3818 ittuple = PyTuple_New(tuplesize);
3819 if (ittuple == NULL)
3820 return NULL;
3821 for (i=0; i < tuplesize; ++i) {
3822 PyObject *item = PyTuple_GET_ITEM(args, i);
3823 PyObject *it = PyObject_GetIter(item);
3824 if (it == NULL) {
3825 if (PyErr_ExceptionMatches(PyExc_TypeError))
3826 PyErr_Format(PyExc_TypeError,
3827 "izip_longest argument #%zd must support iteration",
3828 i+1);
3829 Py_DECREF(ittuple);
3830 return NULL;
3831 }
3832 PyTuple_SET_ITEM(ittuple, i, it);
3833 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003834
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003835 /* create a result holder */
3836 result = PyTuple_New(tuplesize);
3837 if (result == NULL) {
3838 Py_DECREF(ittuple);
3839 return NULL;
3840 }
3841 for (i=0 ; i < tuplesize ; i++) {
3842 Py_INCREF(Py_None);
3843 PyTuple_SET_ITEM(result, i, Py_None);
3844 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003845
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003846 /* create iziplongestobject structure */
3847 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3848 if (lz == NULL) {
3849 Py_DECREF(ittuple);
3850 Py_DECREF(result);
3851 return NULL;
3852 }
3853 lz->ittuple = ittuple;
3854 lz->tuplesize = tuplesize;
3855 lz->numactive = tuplesize;
3856 lz->result = result;
3857 Py_INCREF(fillvalue);
3858 lz->fillvalue = fillvalue;
3859 return (PyObject *)lz;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003860}
3861
3862static void
3863izip_longest_dealloc(iziplongestobject *lz)
3864{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003865 PyObject_GC_UnTrack(lz);
3866 Py_XDECREF(lz->ittuple);
3867 Py_XDECREF(lz->result);
3868 Py_XDECREF(lz->fillvalue);
3869 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003870}
3871
3872static int
3873izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3874{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003875 Py_VISIT(lz->ittuple);
3876 Py_VISIT(lz->result);
3877 Py_VISIT(lz->fillvalue);
3878 return 0;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003879}
3880
3881static PyObject *
3882izip_longest_next(iziplongestobject *lz)
3883{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003884 Py_ssize_t i;
3885 Py_ssize_t tuplesize = lz->tuplesize;
3886 PyObject *result = lz->result;
3887 PyObject *it;
3888 PyObject *item;
3889 PyObject *olditem;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003890
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003891 if (tuplesize == 0)
3892 return NULL;
3893 if (lz->numactive == 0)
3894 return NULL;
3895 if (Py_REFCNT(result) == 1) {
3896 Py_INCREF(result);
3897 for (i=0 ; i < tuplesize ; i++) {
3898 it = PyTuple_GET_ITEM(lz->ittuple, i);
3899 if (it == NULL) {
3900 Py_INCREF(lz->fillvalue);
3901 item = lz->fillvalue;
3902 } else {
3903 item = PyIter_Next(it);
3904 if (item == NULL) {
3905 lz->numactive -= 1;
3906 if (lz->numactive == 0 || PyErr_Occurred()) {
3907 lz->numactive = 0;
3908 Py_DECREF(result);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003909 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003910 } else {
3911 Py_INCREF(lz->fillvalue);
3912 item = lz->fillvalue;
3913 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3914 Py_DECREF(it);
3915 }
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003916 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003917 }
3918 olditem = PyTuple_GET_ITEM(result, i);
3919 PyTuple_SET_ITEM(result, i, item);
3920 Py_DECREF(olditem);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003921 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003922 } else {
3923 result = PyTuple_New(tuplesize);
3924 if (result == NULL)
3925 return NULL;
3926 for (i=0 ; i < tuplesize ; i++) {
3927 it = PyTuple_GET_ITEM(lz->ittuple, i);
3928 if (it == NULL) {
3929 Py_INCREF(lz->fillvalue);
3930 item = lz->fillvalue;
3931 } else {
3932 item = PyIter_Next(it);
3933 if (item == NULL) {
3934 lz->numactive -= 1;
3935 if (lz->numactive == 0 || PyErr_Occurred()) {
3936 lz->numactive = 0;
3937 Py_DECREF(result);
3938 return NULL;
3939 } else {
3940 Py_INCREF(lz->fillvalue);
3941 item = lz->fillvalue;
3942 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3943 Py_DECREF(it);
3944 }
3945 }
3946 }
3947 PyTuple_SET_ITEM(result, i, item);
3948 }
3949 }
3950 return result;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003951}
3952
3953PyDoc_STRVAR(izip_longest_doc,
3954"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3955\n\
3956Return an izip_longest object whose .next() method returns a tuple where\n\
3957the i-th element comes from the i-th iterable argument. The .next()\n\
3958method continues until the longest iterable in the argument sequence\n\
3959is exhausted and then it raises StopIteration. When the shorter iterables\n\
3960are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3961defaults to None or can be specified by a keyword argument.\n\
3962");
3963
3964static PyTypeObject iziplongest_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003965 PyVarObject_HEAD_INIT(NULL, 0)
3966 "itertools.izip_longest", /* tp_name */
3967 sizeof(iziplongestobject), /* tp_basicsize */
3968 0, /* tp_itemsize */
3969 /* methods */
3970 (destructor)izip_longest_dealloc, /* tp_dealloc */
3971 0, /* tp_print */
3972 0, /* tp_getattr */
3973 0, /* tp_setattr */
3974 0, /* tp_compare */
3975 0, /* tp_repr */
3976 0, /* tp_as_number */
3977 0, /* tp_as_sequence */
3978 0, /* tp_as_mapping */
3979 0, /* tp_hash */
3980 0, /* tp_call */
3981 0, /* tp_str */
3982 PyObject_GenericGetAttr, /* tp_getattro */
3983 0, /* tp_setattro */
3984 0, /* tp_as_buffer */
3985 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3986 Py_TPFLAGS_BASETYPE, /* tp_flags */
3987 izip_longest_doc, /* tp_doc */
3988 (traverseproc)izip_longest_traverse, /* tp_traverse */
3989 0, /* tp_clear */
3990 0, /* tp_richcompare */
3991 0, /* tp_weaklistoffset */
3992 PyObject_SelfIter, /* tp_iter */
3993 (iternextfunc)izip_longest_next, /* tp_iternext */
3994 0, /* tp_methods */
3995 0, /* tp_members */
3996 0, /* tp_getset */
3997 0, /* tp_base */
3998 0, /* tp_dict */
3999 0, /* tp_descr_get */
4000 0, /* tp_descr_set */
4001 0, /* tp_dictoffset */
4002 0, /* tp_init */
4003 0, /* tp_alloc */
4004 izip_longest_new, /* tp_new */
4005 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd36862c2007-02-21 05:20:38 +00004006};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004007
4008/* module level code ********************************************************/
4009
4010PyDoc_STRVAR(module_doc,
4011"Functional tools for creating and using iterators.\n\
4012\n\
4013Infinite iterators:\n\
4014count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004015cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004016repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004017\n\
4018Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004019chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettingere76816b2009-01-29 03:43:44 +00004020compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004021dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4022groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00004023ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
4024ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004025islice(seq, [start,] stop [, step]) --> elements from\n\
4026 seq[start:stop:step]\n\
4027imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
4028starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004029tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004030takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004031izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4032izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4033\n\
4034Combinatoric generators:\n\
4035product(p, q, ... [repeat=1]) --> cartesian product\n\
4036permutations(p[, r])\n\
Raymond Hettinger9eac1192009-11-19 01:22:04 +00004037combinations(p, r)\n\
4038combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004039");
4040
4041
Raymond Hettingerad983e72003-11-12 14:32:26 +00004042static PyMethodDef module_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004043 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4044 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004045};
4046
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004047PyMODINIT_FUNC
4048inititertools(void)
4049{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004050 int i;
4051 PyObject *m;
4052 char *name;
4053 PyTypeObject *typelist[] = {
4054 &combinations_type,
4055 &cwr_type,
4056 &cycle_type,
4057 &dropwhile_type,
4058 &takewhile_type,
4059 &islice_type,
4060 &starmap_type,
4061 &imap_type,
4062 &chain_type,
4063 &compress_type,
4064 &ifilter_type,
4065 &ifilterfalse_type,
4066 &count_type,
4067 &izip_type,
4068 &iziplongest_type,
4069 &permutations_type,
4070 &product_type,
4071 &repeat_type,
4072 &groupby_type,
4073 NULL
4074 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004075
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004076 Py_TYPE(&teedataobject_type) = &PyType_Type;
4077 m = Py_InitModule3("itertools", module_methods, module_doc);
4078 if (m == NULL)
4079 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004080
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004081 for (i=0 ; typelist[i] != NULL ; i++) {
4082 if (PyType_Ready(typelist[i]) < 0)
4083 return;
4084 name = strchr(typelist[i]->tp_name, '.');
4085 assert (name != NULL);
4086 Py_INCREF(typelist[i]);
4087 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4088 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004089
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004090 if (PyType_Ready(&teedataobject_type) < 0)
4091 return;
4092 if (PyType_Ready(&tee_type) < 0)
4093 return;
4094 if (PyType_Ready(&_grouper_type) < 0)
4095 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004096}