blob: ff25335c04bc320023040dc38ba59abd9d483208 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
2#include "Python.h"
Fred Drake08ebfec2004-10-17 19:36:57 +00003#include "structmember.h"
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004
Antoine Pitrouc83ea132010-05-09 14:46:46 +00005/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00006 by Raymond D. Hettinger <python@rcn.com>
7 Copyright (c) 2003 Python Software Foundation.
8 All rights reserved.
9*/
10
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000011
12/* groupby object ***********************************************************/
13
14typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +000015 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000021} groupbyobject;
22
23static PyTypeObject groupby_type;
24static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26static PyObject *
27groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28{
Antoine Pitrouc83ea132010-05-09 14:46:46 +000029 static char *kwargs[] = {"iterable", "key", NULL};
30 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000032
Antoine Pitrouc83ea132010-05-09 14:46:46 +000033 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
34 &it, &keyfunc))
35 return NULL;
36
37 gbo = (groupbyobject *)type->tp_alloc(type, 0);
38 if (gbo == NULL)
39 return NULL;
40 gbo->tgtkey = NULL;
41 gbo->currkey = NULL;
42 gbo->currvalue = NULL;
43 gbo->keyfunc = keyfunc;
44 Py_INCREF(keyfunc);
45 gbo->it = PyObject_GetIter(it);
46 if (gbo->it == NULL) {
47 Py_DECREF(gbo);
48 return NULL;
49 }
50 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000051}
52
53static void
54groupby_dealloc(groupbyobject *gbo)
55{
Antoine Pitrouc83ea132010-05-09 14:46:46 +000056 PyObject_GC_UnTrack(gbo);
57 Py_XDECREF(gbo->it);
58 Py_XDECREF(gbo->keyfunc);
59 Py_XDECREF(gbo->tgtkey);
60 Py_XDECREF(gbo->currkey);
61 Py_XDECREF(gbo->currvalue);
62 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000063}
64
65static int
66groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
67{
Antoine Pitrouc83ea132010-05-09 14:46:46 +000068 Py_VISIT(gbo->it);
69 Py_VISIT(gbo->keyfunc);
70 Py_VISIT(gbo->tgtkey);
71 Py_VISIT(gbo->currkey);
72 Py_VISIT(gbo->currvalue);
73 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000074}
75
76static PyObject *
77groupby_next(groupbyobject *gbo)
78{
Antoine Pitrouc83ea132010-05-09 14:46:46 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
Antoine Pitrouc83ea132010-05-09 14:46:46 +000081 /* skip to next iteration group */
82 for (;;) {
83 if (gbo->currkey == NULL)
84 /* pass */;
85 else if (gbo->tgtkey == NULL)
86 break;
87 else {
88 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000089
Antoine Pitrouc83ea132010-05-09 14:46:46 +000090 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
91 gbo->currkey, Py_EQ);
92 if (rcmp == -1)
93 return NULL;
94 else if (rcmp == 0)
95 break;
96 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000097
Antoine Pitrouc83ea132010-05-09 14:46:46 +000098 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000101
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000102 if (gbo->keyfunc == Py_None) {
103 newkey = newvalue;
104 Py_INCREF(newvalue);
105 } else {
106 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
107 newvalue, NULL);
108 if (newkey == NULL) {
109 Py_DECREF(newvalue);
110 return NULL;
111 }
112 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000113
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000114 tmp = gbo->currkey;
115 gbo->currkey = newkey;
116 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000117
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000118 tmp = gbo->currvalue;
119 gbo->currvalue = newvalue;
120 Py_XDECREF(tmp);
121 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000123 Py_INCREF(gbo->currkey);
124 tmp = gbo->tgtkey;
125 gbo->tgtkey = gbo->currkey;
126 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000127
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000131
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000135}
136
137PyDoc_STRVAR(groupby_doc,
138"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
139(key, sub-iterator) grouped by each value of key(value).\n");
140
141static PyTypeObject groupby_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000142 PyVarObject_HEAD_INIT(NULL, 0)
143 "itertools.groupby", /* tp_name */
144 sizeof(groupbyobject), /* tp_basicsize */
145 0, /* tp_itemsize */
146 /* methods */
147 (destructor)groupby_dealloc, /* tp_dealloc */
148 0, /* tp_print */
149 0, /* tp_getattr */
150 0, /* tp_setattr */
151 0, /* tp_compare */
152 0, /* tp_repr */
153 0, /* tp_as_number */
154 0, /* tp_as_sequence */
155 0, /* tp_as_mapping */
156 0, /* tp_hash */
157 0, /* tp_call */
158 0, /* tp_str */
159 PyObject_GenericGetAttr, /* tp_getattro */
160 0, /* tp_setattro */
161 0, /* tp_as_buffer */
162 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
163 Py_TPFLAGS_BASETYPE, /* tp_flags */
164 groupby_doc, /* tp_doc */
165 (traverseproc)groupby_traverse, /* tp_traverse */
166 0, /* tp_clear */
167 0, /* tp_richcompare */
168 0, /* tp_weaklistoffset */
169 PyObject_SelfIter, /* tp_iter */
170 (iternextfunc)groupby_next, /* tp_iternext */
171 0, /* tp_methods */
172 0, /* tp_members */
173 0, /* tp_getset */
174 0, /* tp_base */
175 0, /* tp_dict */
176 0, /* tp_descr_get */
177 0, /* tp_descr_set */
178 0, /* tp_dictoffset */
179 0, /* tp_init */
180 0, /* tp_alloc */
181 groupby_new, /* tp_new */
182 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000183};
184
185
186/* _grouper object (internal) ************************************************/
187
188typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000189 PyObject_HEAD
190 PyObject *parent;
191 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000192} _grouperobject;
193
194static PyTypeObject _grouper_type;
195
196static PyObject *
197_grouper_create(groupbyobject *parent, PyObject *tgtkey)
198{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000199 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000200
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000201 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
202 if (igo == NULL)
203 return NULL;
204 igo->parent = (PyObject *)parent;
205 Py_INCREF(parent);
206 igo->tgtkey = tgtkey;
207 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000208
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000209 PyObject_GC_Track(igo);
210 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000211}
212
213static void
214_grouper_dealloc(_grouperobject *igo)
215{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000216 PyObject_GC_UnTrack(igo);
217 Py_DECREF(igo->parent);
218 Py_DECREF(igo->tgtkey);
219 PyObject_GC_Del(igo);
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000220}
221
222static int
223_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
224{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000225 Py_VISIT(igo->parent);
226 Py_VISIT(igo->tgtkey);
227 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000228}
229
230static PyObject *
231_grouper_next(_grouperobject *igo)
232{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000233 groupbyobject *gbo = (groupbyobject *)igo->parent;
234 PyObject *newvalue, *newkey, *r;
235 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000236
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000237 if (gbo->currvalue == NULL) {
238 newvalue = PyIter_Next(gbo->it);
239 if (newvalue == NULL)
240 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000241
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000242 if (gbo->keyfunc == Py_None) {
243 newkey = newvalue;
244 Py_INCREF(newvalue);
245 } else {
246 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
247 newvalue, NULL);
248 if (newkey == NULL) {
249 Py_DECREF(newvalue);
250 return NULL;
251 }
252 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000253
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000254 assert(gbo->currkey == NULL);
255 gbo->currkey = newkey;
256 gbo->currvalue = newvalue;
257 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000259 assert(gbo->currkey != NULL);
260 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
261 if (rcmp <= 0)
262 /* got any error or current group is end */
263 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000264
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000265 r = gbo->currvalue;
266 gbo->currvalue = NULL;
267 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000268
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000269 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000270}
271
272static PyTypeObject _grouper_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000273 PyVarObject_HEAD_INIT(NULL, 0)
274 "itertools._grouper", /* tp_name */
275 sizeof(_grouperobject), /* tp_basicsize */
276 0, /* tp_itemsize */
277 /* methods */
278 (destructor)_grouper_dealloc, /* tp_dealloc */
279 0, /* tp_print */
280 0, /* tp_getattr */
281 0, /* tp_setattr */
282 0, /* tp_compare */
283 0, /* tp_repr */
284 0, /* tp_as_number */
285 0, /* tp_as_sequence */
286 0, /* tp_as_mapping */
287 0, /* tp_hash */
288 0, /* tp_call */
289 0, /* tp_str */
290 PyObject_GenericGetAttr, /* tp_getattro */
291 0, /* tp_setattro */
292 0, /* tp_as_buffer */
293 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
294 0, /* tp_doc */
295 (traverseproc)_grouper_traverse,/* tp_traverse */
296 0, /* tp_clear */
297 0, /* tp_richcompare */
298 0, /* tp_weaklistoffset */
299 PyObject_SelfIter, /* tp_iter */
300 (iternextfunc)_grouper_next, /* tp_iternext */
301 0, /* tp_methods */
302 0, /* tp_members */
303 0, /* tp_getset */
304 0, /* tp_base */
305 0, /* tp_dict */
306 0, /* tp_descr_get */
307 0, /* tp_descr_set */
308 0, /* tp_dictoffset */
309 0, /* tp_init */
310 0, /* tp_alloc */
311 0, /* tp_new */
312 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000313};
314
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000315
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000316
Raymond Hettingerad983e72003-11-12 14:32:26 +0000317/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000318
Raymond Hettingerad983e72003-11-12 14:32:26 +0000319/* The teedataobject pre-allocates space for LINKCELLS number of objects.
320 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000321 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000322 the other structure members including PyHEAD overhead. The larger the
323 value, the less memory overhead per object and the less time spent
324 allocating/deallocating new links. The smaller the number, the less
325 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000326*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000327#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000328
329typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000330 PyObject_HEAD
331 PyObject *it;
332 int numread;
333 PyObject *nextlink;
334 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000335} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000336
337typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000338 PyObject_HEAD
339 teedataobject *dataobj;
340 int index;
341 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000342} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000343
Raymond Hettingerad983e72003-11-12 14:32:26 +0000344static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000345
346static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000347teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000348{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000349 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000350
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000351 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
352 if (tdo == NULL)
353 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000354
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000355 tdo->numread = 0;
356 tdo->nextlink = NULL;
357 Py_INCREF(it);
358 tdo->it = it;
359 PyObject_GC_Track(tdo);
360 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000361}
362
363static PyObject *
364teedataobject_jumplink(teedataobject *tdo)
365{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000366 if (tdo->nextlink == NULL)
367 tdo->nextlink = teedataobject_new(tdo->it);
368 Py_XINCREF(tdo->nextlink);
369 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000370}
371
372static PyObject *
373teedataobject_getitem(teedataobject *tdo, int i)
374{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000375 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000376
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000377 assert(i < LINKCELLS);
378 if (i < tdo->numread)
379 value = tdo->values[i];
380 else {
381 /* this is the lead iterator, so fetch more data */
382 assert(i == tdo->numread);
383 value = PyIter_Next(tdo->it);
384 if (value == NULL)
385 return NULL;
386 tdo->numread++;
387 tdo->values[i] = value;
388 }
389 Py_INCREF(value);
390 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000391}
392
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000393static int
394teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
395{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000396 int i;
397 Py_VISIT(tdo->it);
398 for (i = 0; i < tdo->numread; i++)
399 Py_VISIT(tdo->values[i]);
400 Py_VISIT(tdo->nextlink);
401 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000402}
403
404static int
405teedataobject_clear(teedataobject *tdo)
406{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000407 int i;
408 Py_CLEAR(tdo->it);
409 for (i=0 ; i<tdo->numread ; i++)
410 Py_CLEAR(tdo->values[i]);
411 Py_CLEAR(tdo->nextlink);
412 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000413}
414
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000416teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000418 PyObject_GC_UnTrack(tdo);
419 teedataobject_clear(tdo);
420 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000421}
422
Raymond Hettingerad983e72003-11-12 14:32:26 +0000423PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000424
Raymond Hettingerad983e72003-11-12 14:32:26 +0000425static PyTypeObject teedataobject_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000426 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
427 "itertools.tee_dataobject", /* tp_name */
428 sizeof(teedataobject), /* tp_basicsize */
429 0, /* tp_itemsize */
430 /* methods */
431 (destructor)teedataobject_dealloc, /* tp_dealloc */
432 0, /* tp_print */
433 0, /* tp_getattr */
434 0, /* tp_setattr */
435 0, /* tp_compare */
436 0, /* tp_repr */
437 0, /* tp_as_number */
438 0, /* tp_as_sequence */
439 0, /* tp_as_mapping */
440 0, /* tp_hash */
441 0, /* tp_call */
442 0, /* tp_str */
443 PyObject_GenericGetAttr, /* tp_getattro */
444 0, /* tp_setattro */
445 0, /* tp_as_buffer */
446 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
447 teedataobject_doc, /* tp_doc */
448 (traverseproc)teedataobject_traverse, /* tp_traverse */
449 (inquiry)teedataobject_clear, /* tp_clear */
450 0, /* tp_richcompare */
451 0, /* tp_weaklistoffset */
452 0, /* tp_iter */
453 0, /* tp_iternext */
454 0, /* tp_methods */
455 0, /* tp_members */
456 0, /* tp_getset */
457 0, /* tp_base */
458 0, /* tp_dict */
459 0, /* tp_descr_get */
460 0, /* tp_descr_set */
461 0, /* tp_dictoffset */
462 0, /* tp_init */
463 0, /* tp_alloc */
464 0, /* tp_new */
465 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000466};
467
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000468
469static PyTypeObject tee_type;
470
471static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000472tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000473{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000474 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000475
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000476 if (to->index >= LINKCELLS) {
477 link = teedataobject_jumplink(to->dataobj);
478 Py_DECREF(to->dataobj);
479 to->dataobj = (teedataobject *)link;
480 to->index = 0;
481 }
482 value = teedataobject_getitem(to->dataobj, to->index);
483 if (value == NULL)
484 return NULL;
485 to->index++;
486 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000487}
488
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000489static int
490tee_traverse(teeobject *to, visitproc visit, void *arg)
491{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000492 Py_VISIT((PyObject *)to->dataobj);
493 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000494}
495
Raymond Hettingerad983e72003-11-12 14:32:26 +0000496static PyObject *
497tee_copy(teeobject *to)
498{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000499 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000500
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000501 newto = PyObject_GC_New(teeobject, &tee_type);
502 if (newto == NULL)
503 return NULL;
504 Py_INCREF(to->dataobj);
505 newto->dataobj = to->dataobj;
506 newto->index = to->index;
507 newto->weakreflist = NULL;
508 PyObject_GC_Track(newto);
509 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000510}
511
512PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
513
514static PyObject *
515tee_fromiterable(PyObject *iterable)
516{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000517 teeobject *to;
518 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000519
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000520 it = PyObject_GetIter(iterable);
521 if (it == NULL)
522 return NULL;
523 if (PyObject_TypeCheck(it, &tee_type)) {
524 to = (teeobject *)tee_copy((teeobject *)it);
525 goto done;
526 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000527
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000528 to = PyObject_GC_New(teeobject, &tee_type);
529 if (to == NULL)
530 goto done;
531 to->dataobj = (teedataobject *)teedataobject_new(it);
532 if (!to->dataobj) {
533 PyObject_GC_Del(to);
534 to = NULL;
535 goto done;
536 }
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000537
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000538 to->index = 0;
539 to->weakreflist = NULL;
540 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000541done:
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000542 Py_XDECREF(it);
543 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000544}
545
546static PyObject *
547tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
548{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000549 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000550
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000551 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
552 return NULL;
553 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000554}
555
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000556static int
557tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000558{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000559 if (to->weakreflist != NULL)
560 PyObject_ClearWeakRefs((PyObject *) to);
561 Py_CLEAR(to->dataobj);
562 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000563}
564
565static void
566tee_dealloc(teeobject *to)
567{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000568 PyObject_GC_UnTrack(to);
569 tee_clear(to);
570 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000571}
572
Raymond Hettingerad983e72003-11-12 14:32:26 +0000573PyDoc_STRVAR(teeobject_doc,
574"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000575
Raymond Hettingerad983e72003-11-12 14:32:26 +0000576static PyMethodDef tee_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000577 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
578 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000579};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000580
581static PyTypeObject tee_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000582 PyVarObject_HEAD_INIT(NULL, 0)
583 "itertools.tee", /* tp_name */
584 sizeof(teeobject), /* tp_basicsize */
585 0, /* tp_itemsize */
586 /* methods */
587 (destructor)tee_dealloc, /* tp_dealloc */
588 0, /* tp_print */
589 0, /* tp_getattr */
590 0, /* tp_setattr */
591 0, /* tp_compare */
592 0, /* tp_repr */
593 0, /* tp_as_number */
594 0, /* tp_as_sequence */
595 0, /* tp_as_mapping */
596 0, /* tp_hash */
597 0, /* tp_call */
598 0, /* tp_str */
599 0, /* tp_getattro */
600 0, /* tp_setattro */
601 0, /* tp_as_buffer */
602 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
603 teeobject_doc, /* tp_doc */
604 (traverseproc)tee_traverse, /* tp_traverse */
605 (inquiry)tee_clear, /* tp_clear */
606 0, /* tp_richcompare */
607 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
608 PyObject_SelfIter, /* tp_iter */
609 (iternextfunc)tee_next, /* tp_iternext */
610 tee_methods, /* tp_methods */
611 0, /* tp_members */
612 0, /* tp_getset */
613 0, /* tp_base */
614 0, /* tp_dict */
615 0, /* tp_descr_get */
616 0, /* tp_descr_set */
617 0, /* tp_dictoffset */
618 0, /* tp_init */
619 0, /* tp_alloc */
620 tee_new, /* tp_new */
621 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000622};
623
Raymond Hettingerad983e72003-11-12 14:32:26 +0000624static PyObject *
625tee(PyObject *self, PyObject *args)
626{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000627 Py_ssize_t i, n=2;
628 PyObject *it, *iterable, *copyable, *result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000629
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000630 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
631 return NULL;
632 if (n < 0) {
633 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
634 return NULL;
635 }
636 result = PyTuple_New(n);
637 if (result == NULL)
638 return NULL;
639 if (n == 0)
640 return result;
641 it = PyObject_GetIter(iterable);
642 if (it == NULL) {
643 Py_DECREF(result);
644 return NULL;
645 }
646 if (!PyObject_HasAttrString(it, "__copy__")) {
647 copyable = tee_fromiterable(it);
648 Py_DECREF(it);
649 if (copyable == NULL) {
650 Py_DECREF(result);
651 return NULL;
652 }
653 } else
654 copyable = it;
655 PyTuple_SET_ITEM(result, 0, copyable);
656 for (i=1 ; i<n ; i++) {
657 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
658 if (copyable == NULL) {
659 Py_DECREF(result);
660 return NULL;
661 }
662 PyTuple_SET_ITEM(result, i, copyable);
663 }
664 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000665}
666
667PyDoc_STRVAR(tee_doc,
668"tee(iterable, n=2) --> tuple of n independent iterators.");
669
670
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000671/* cycle object **********************************************************/
672
673typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000674 PyObject_HEAD
675 PyObject *it;
676 PyObject *saved;
677 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000678} cycleobject;
679
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000680static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000681
682static PyObject *
683cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
684{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000685 PyObject *it;
686 PyObject *iterable;
687 PyObject *saved;
688 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000689
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000690 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
691 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000692
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000693 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
694 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000695
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000696 /* Get iterator. */
697 it = PyObject_GetIter(iterable);
698 if (it == NULL)
699 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000700
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000701 saved = PyList_New(0);
702 if (saved == NULL) {
703 Py_DECREF(it);
704 return NULL;
705 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000706
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000707 /* create cycleobject structure */
708 lz = (cycleobject *)type->tp_alloc(type, 0);
709 if (lz == NULL) {
710 Py_DECREF(it);
711 Py_DECREF(saved);
712 return NULL;
713 }
714 lz->it = it;
715 lz->saved = saved;
716 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000717
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000718 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000719}
720
721static void
722cycle_dealloc(cycleobject *lz)
723{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000724 PyObject_GC_UnTrack(lz);
725 Py_XDECREF(lz->saved);
726 Py_XDECREF(lz->it);
727 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000728}
729
730static int
731cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
732{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000733 Py_VISIT(lz->it);
734 Py_VISIT(lz->saved);
735 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000736}
737
738static PyObject *
739cycle_next(cycleobject *lz)
740{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000741 PyObject *item;
742 PyObject *it;
743 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000744
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000745 while (1) {
746 item = PyIter_Next(lz->it);
747 if (item != NULL) {
748 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
749 Py_DECREF(item);
750 return NULL;
751 }
752 return item;
753 }
754 if (PyErr_Occurred()) {
755 if (PyErr_ExceptionMatches(PyExc_StopIteration))
756 PyErr_Clear();
757 else
758 return NULL;
759 }
760 if (PyList_Size(lz->saved) == 0)
761 return NULL;
762 it = PyObject_GetIter(lz->saved);
763 if (it == NULL)
764 return NULL;
765 tmp = lz->it;
766 lz->it = it;
767 lz->firstpass = 1;
768 Py_DECREF(tmp);
769 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000770}
771
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000772PyDoc_STRVAR(cycle_doc,
773"cycle(iterable) --> cycle object\n\
774\n\
775Return elements from the iterable until it is exhausted.\n\
776Then repeat the sequence indefinitely.");
777
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000778static PyTypeObject cycle_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000779 PyVarObject_HEAD_INIT(NULL, 0)
780 "itertools.cycle", /* tp_name */
781 sizeof(cycleobject), /* tp_basicsize */
782 0, /* tp_itemsize */
783 /* methods */
784 (destructor)cycle_dealloc, /* tp_dealloc */
785 0, /* tp_print */
786 0, /* tp_getattr */
787 0, /* tp_setattr */
788 0, /* tp_compare */
789 0, /* tp_repr */
790 0, /* tp_as_number */
791 0, /* tp_as_sequence */
792 0, /* tp_as_mapping */
793 0, /* tp_hash */
794 0, /* tp_call */
795 0, /* tp_str */
796 PyObject_GenericGetAttr, /* tp_getattro */
797 0, /* tp_setattro */
798 0, /* tp_as_buffer */
799 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
800 Py_TPFLAGS_BASETYPE, /* tp_flags */
801 cycle_doc, /* tp_doc */
802 (traverseproc)cycle_traverse, /* tp_traverse */
803 0, /* tp_clear */
804 0, /* tp_richcompare */
805 0, /* tp_weaklistoffset */
806 PyObject_SelfIter, /* tp_iter */
807 (iternextfunc)cycle_next, /* tp_iternext */
808 0, /* tp_methods */
809 0, /* tp_members */
810 0, /* tp_getset */
811 0, /* tp_base */
812 0, /* tp_dict */
813 0, /* tp_descr_get */
814 0, /* tp_descr_set */
815 0, /* tp_dictoffset */
816 0, /* tp_init */
817 0, /* tp_alloc */
818 cycle_new, /* tp_new */
819 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000820};
821
822
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000823/* dropwhile object **********************************************************/
824
825typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000826 PyObject_HEAD
827 PyObject *func;
828 PyObject *it;
829 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000830} dropwhileobject;
831
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000832static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000833
834static PyObject *
835dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
836{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000837 PyObject *func, *seq;
838 PyObject *it;
839 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000840
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000841 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
842 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000843
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000844 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
845 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000846
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000847 /* Get iterator. */
848 it = PyObject_GetIter(seq);
849 if (it == NULL)
850 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000851
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000852 /* create dropwhileobject structure */
853 lz = (dropwhileobject *)type->tp_alloc(type, 0);
854 if (lz == NULL) {
855 Py_DECREF(it);
856 return NULL;
857 }
858 Py_INCREF(func);
859 lz->func = func;
860 lz->it = it;
861 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000862
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000863 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000864}
865
866static void
867dropwhile_dealloc(dropwhileobject *lz)
868{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000869 PyObject_GC_UnTrack(lz);
870 Py_XDECREF(lz->func);
871 Py_XDECREF(lz->it);
872 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000873}
874
875static int
876dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
877{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000878 Py_VISIT(lz->it);
879 Py_VISIT(lz->func);
880 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000881}
882
883static PyObject *
884dropwhile_next(dropwhileobject *lz)
885{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000886 PyObject *item, *good;
887 PyObject *it = lz->it;
888 long ok;
889 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000890
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000891 iternext = *Py_TYPE(it)->tp_iternext;
892 for (;;) {
893 item = iternext(it);
894 if (item == NULL)
895 return NULL;
896 if (lz->start == 1)
897 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000898
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000899 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
900 if (good == NULL) {
901 Py_DECREF(item);
902 return NULL;
903 }
904 ok = PyObject_IsTrue(good);
905 Py_DECREF(good);
906 if (!ok) {
907 lz->start = 1;
908 return item;
909 }
910 Py_DECREF(item);
911 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000912}
913
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000914PyDoc_STRVAR(dropwhile_doc,
915"dropwhile(predicate, iterable) --> dropwhile object\n\
916\n\
917Drop items from the iterable while predicate(item) is true.\n\
918Afterwards, return every element until the iterable is exhausted.");
919
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000920static PyTypeObject dropwhile_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000921 PyVarObject_HEAD_INIT(NULL, 0)
922 "itertools.dropwhile", /* tp_name */
923 sizeof(dropwhileobject), /* tp_basicsize */
924 0, /* tp_itemsize */
925 /* methods */
926 (destructor)dropwhile_dealloc, /* tp_dealloc */
927 0, /* tp_print */
928 0, /* tp_getattr */
929 0, /* tp_setattr */
930 0, /* tp_compare */
931 0, /* tp_repr */
932 0, /* tp_as_number */
933 0, /* tp_as_sequence */
934 0, /* tp_as_mapping */
935 0, /* tp_hash */
936 0, /* tp_call */
937 0, /* tp_str */
938 PyObject_GenericGetAttr, /* tp_getattro */
939 0, /* tp_setattro */
940 0, /* tp_as_buffer */
941 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
942 Py_TPFLAGS_BASETYPE, /* tp_flags */
943 dropwhile_doc, /* tp_doc */
944 (traverseproc)dropwhile_traverse, /* tp_traverse */
945 0, /* tp_clear */
946 0, /* tp_richcompare */
947 0, /* tp_weaklistoffset */
948 PyObject_SelfIter, /* tp_iter */
949 (iternextfunc)dropwhile_next, /* tp_iternext */
950 0, /* tp_methods */
951 0, /* tp_members */
952 0, /* tp_getset */
953 0, /* tp_base */
954 0, /* tp_dict */
955 0, /* tp_descr_get */
956 0, /* tp_descr_set */
957 0, /* tp_dictoffset */
958 0, /* tp_init */
959 0, /* tp_alloc */
960 dropwhile_new, /* tp_new */
961 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000962};
963
964
965/* takewhile object **********************************************************/
966
967typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000968 PyObject_HEAD
969 PyObject *func;
970 PyObject *it;
971 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000972} takewhileobject;
973
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000974static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000975
976static PyObject *
977takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
978{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000979 PyObject *func, *seq;
980 PyObject *it;
981 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000982
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000983 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
984 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000985
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000986 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
987 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000988
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000989 /* Get iterator. */
990 it = PyObject_GetIter(seq);
991 if (it == NULL)
992 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000993
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000994 /* create takewhileobject structure */
995 lz = (takewhileobject *)type->tp_alloc(type, 0);
996 if (lz == NULL) {
997 Py_DECREF(it);
998 return NULL;
999 }
1000 Py_INCREF(func);
1001 lz->func = func;
1002 lz->it = it;
1003 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001004
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001005 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001006}
1007
1008static void
1009takewhile_dealloc(takewhileobject *lz)
1010{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001011 PyObject_GC_UnTrack(lz);
1012 Py_XDECREF(lz->func);
1013 Py_XDECREF(lz->it);
1014 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001015}
1016
1017static int
1018takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1019{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001020 Py_VISIT(lz->it);
1021 Py_VISIT(lz->func);
1022 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001023}
1024
1025static PyObject *
1026takewhile_next(takewhileobject *lz)
1027{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001028 PyObject *item, *good;
1029 PyObject *it = lz->it;
1030 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001031
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001032 if (lz->stop == 1)
1033 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001034
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001035 item = (*Py_TYPE(it)->tp_iternext)(it);
1036 if (item == NULL)
1037 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001038
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001039 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1040 if (good == NULL) {
1041 Py_DECREF(item);
1042 return NULL;
1043 }
1044 ok = PyObject_IsTrue(good);
1045 Py_DECREF(good);
1046 if (ok)
1047 return item;
1048 Py_DECREF(item);
1049 lz->stop = 1;
1050 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001051}
1052
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053PyDoc_STRVAR(takewhile_doc,
1054"takewhile(predicate, iterable) --> takewhile object\n\
1055\n\
1056Return successive entries from an iterable as long as the \n\
1057predicate evaluates to true for each entry.");
1058
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001059static PyTypeObject takewhile_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001060 PyVarObject_HEAD_INIT(NULL, 0)
1061 "itertools.takewhile", /* tp_name */
1062 sizeof(takewhileobject), /* tp_basicsize */
1063 0, /* tp_itemsize */
1064 /* methods */
1065 (destructor)takewhile_dealloc, /* tp_dealloc */
1066 0, /* tp_print */
1067 0, /* tp_getattr */
1068 0, /* tp_setattr */
1069 0, /* tp_compare */
1070 0, /* tp_repr */
1071 0, /* tp_as_number */
1072 0, /* tp_as_sequence */
1073 0, /* tp_as_mapping */
1074 0, /* tp_hash */
1075 0, /* tp_call */
1076 0, /* tp_str */
1077 PyObject_GenericGetAttr, /* tp_getattro */
1078 0, /* tp_setattro */
1079 0, /* tp_as_buffer */
1080 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1081 Py_TPFLAGS_BASETYPE, /* tp_flags */
1082 takewhile_doc, /* tp_doc */
1083 (traverseproc)takewhile_traverse, /* tp_traverse */
1084 0, /* tp_clear */
1085 0, /* tp_richcompare */
1086 0, /* tp_weaklistoffset */
1087 PyObject_SelfIter, /* tp_iter */
1088 (iternextfunc)takewhile_next, /* tp_iternext */
1089 0, /* tp_methods */
1090 0, /* tp_members */
1091 0, /* tp_getset */
1092 0, /* tp_base */
1093 0, /* tp_dict */
1094 0, /* tp_descr_get */
1095 0, /* tp_descr_set */
1096 0, /* tp_dictoffset */
1097 0, /* tp_init */
1098 0, /* tp_alloc */
1099 takewhile_new, /* tp_new */
1100 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001101};
1102
1103
1104/* islice object ************************************************************/
1105
1106typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001107 PyObject_HEAD
1108 PyObject *it;
1109 Py_ssize_t next;
1110 Py_ssize_t stop;
1111 Py_ssize_t step;
1112 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001113} isliceobject;
1114
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001115static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116
1117static PyObject *
1118islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1119{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001120 PyObject *seq;
1121 Py_ssize_t start=0, stop=-1, step=1;
1122 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1123 Py_ssize_t numargs;
1124 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001125
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001126 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1127 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001128
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001129 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1130 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001131
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001132 numargs = PyTuple_Size(args);
1133 if (numargs == 2) {
1134 if (a1 != Py_None) {
1135 stop = PyInt_AsSsize_t(a1);
1136 if (stop == -1) {
1137 if (PyErr_Occurred())
1138 PyErr_Clear();
1139 PyErr_SetString(PyExc_ValueError,
1140 "Stop argument for islice() must be None or an integer: 0 <= x <= maxint.");
1141 return NULL;
1142 }
1143 }
1144 } else {
1145 if (a1 != Py_None)
1146 start = PyInt_AsSsize_t(a1);
1147 if (start == -1 && PyErr_Occurred())
1148 PyErr_Clear();
1149 if (a2 != Py_None) {
1150 stop = PyInt_AsSsize_t(a2);
1151 if (stop == -1) {
1152 if (PyErr_Occurred())
1153 PyErr_Clear();
1154 PyErr_SetString(PyExc_ValueError,
1155 "Stop argument for islice() must be None or an integer: 0 <= x <= maxint.");
1156 return NULL;
1157 }
1158 }
1159 }
1160 if (start<0 || stop<-1) {
1161 PyErr_SetString(PyExc_ValueError,
1162 "Indices for islice() must be None or an integer: 0 <= x <= maxint.");
1163 return NULL;
1164 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001165
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001166 if (a3 != NULL) {
1167 if (a3 != Py_None)
1168 step = PyInt_AsSsize_t(a3);
1169 if (step == -1 && PyErr_Occurred())
1170 PyErr_Clear();
1171 }
1172 if (step<1) {
1173 PyErr_SetString(PyExc_ValueError,
1174 "Step for islice() must be a positive integer or None.");
1175 return NULL;
1176 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001177
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001178 /* Get iterator. */
1179 it = PyObject_GetIter(seq);
1180 if (it == NULL)
1181 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001182
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001183 /* create isliceobject structure */
1184 lz = (isliceobject *)type->tp_alloc(type, 0);
1185 if (lz == NULL) {
1186 Py_DECREF(it);
1187 return NULL;
1188 }
1189 lz->it = it;
1190 lz->next = start;
1191 lz->stop = stop;
1192 lz->step = step;
1193 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001194
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001195 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196}
1197
1198static void
1199islice_dealloc(isliceobject *lz)
1200{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001201 PyObject_GC_UnTrack(lz);
1202 Py_XDECREF(lz->it);
1203 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204}
1205
1206static int
1207islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1208{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001209 Py_VISIT(lz->it);
1210 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001211}
1212
1213static PyObject *
1214islice_next(isliceobject *lz)
1215{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001216 PyObject *item;
1217 PyObject *it = lz->it;
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001218 Py_ssize_t stop = lz->stop;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001219 Py_ssize_t oldnext;
1220 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001221
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001222 iternext = *Py_TYPE(it)->tp_iternext;
1223 while (lz->cnt < lz->next) {
1224 item = iternext(it);
1225 if (item == NULL)
1226 return NULL;
1227 Py_DECREF(item);
1228 lz->cnt++;
1229 }
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001230 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001231 return NULL;
1232 item = iternext(it);
1233 if (item == NULL)
1234 return NULL;
1235 lz->cnt++;
1236 oldnext = lz->next;
1237 lz->next += lz->step;
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001238 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1239 lz->next = stop;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001240 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001241}
1242
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001243PyDoc_STRVAR(islice_doc,
1244"islice(iterable, [start,] stop [, step]) --> islice object\n\
1245\n\
1246Return an iterator whose next() method returns selected values from an\n\
1247iterable. If start is specified, will skip all preceding elements;\n\
1248otherwise, start defaults to zero. Step defaults to one. If\n\
1249specified as another value, step determines how many values are \n\
1250skipped between successive calls. Works like a slice() on a list\n\
1251but returns an iterator.");
1252
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001253static PyTypeObject islice_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001254 PyVarObject_HEAD_INIT(NULL, 0)
1255 "itertools.islice", /* tp_name */
1256 sizeof(isliceobject), /* tp_basicsize */
1257 0, /* tp_itemsize */
1258 /* methods */
1259 (destructor)islice_dealloc, /* tp_dealloc */
1260 0, /* tp_print */
1261 0, /* tp_getattr */
1262 0, /* tp_setattr */
1263 0, /* tp_compare */
1264 0, /* tp_repr */
1265 0, /* tp_as_number */
1266 0, /* tp_as_sequence */
1267 0, /* tp_as_mapping */
1268 0, /* tp_hash */
1269 0, /* tp_call */
1270 0, /* tp_str */
1271 PyObject_GenericGetAttr, /* tp_getattro */
1272 0, /* tp_setattro */
1273 0, /* tp_as_buffer */
1274 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1275 Py_TPFLAGS_BASETYPE, /* tp_flags */
1276 islice_doc, /* tp_doc */
1277 (traverseproc)islice_traverse, /* tp_traverse */
1278 0, /* tp_clear */
1279 0, /* tp_richcompare */
1280 0, /* tp_weaklistoffset */
1281 PyObject_SelfIter, /* tp_iter */
1282 (iternextfunc)islice_next, /* tp_iternext */
1283 0, /* tp_methods */
1284 0, /* tp_members */
1285 0, /* tp_getset */
1286 0, /* tp_base */
1287 0, /* tp_dict */
1288 0, /* tp_descr_get */
1289 0, /* tp_descr_set */
1290 0, /* tp_dictoffset */
1291 0, /* tp_init */
1292 0, /* tp_alloc */
1293 islice_new, /* tp_new */
1294 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001295};
1296
1297
1298/* starmap object ************************************************************/
1299
1300typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001301 PyObject_HEAD
1302 PyObject *func;
1303 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001304} starmapobject;
1305
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001306static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001307
1308static PyObject *
1309starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1310{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001311 PyObject *func, *seq;
1312 PyObject *it;
1313 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001314
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001315 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1316 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001317
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001318 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1319 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001320
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001321 /* Get iterator. */
1322 it = PyObject_GetIter(seq);
1323 if (it == NULL)
1324 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001325
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001326 /* create starmapobject structure */
1327 lz = (starmapobject *)type->tp_alloc(type, 0);
1328 if (lz == NULL) {
1329 Py_DECREF(it);
1330 return NULL;
1331 }
1332 Py_INCREF(func);
1333 lz->func = func;
1334 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001335
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001336 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001337}
1338
1339static void
1340starmap_dealloc(starmapobject *lz)
1341{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001342 PyObject_GC_UnTrack(lz);
1343 Py_XDECREF(lz->func);
1344 Py_XDECREF(lz->it);
1345 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001346}
1347
1348static int
1349starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1350{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001351 Py_VISIT(lz->it);
1352 Py_VISIT(lz->func);
1353 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001354}
1355
1356static PyObject *
1357starmap_next(starmapobject *lz)
1358{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001359 PyObject *args;
1360 PyObject *result;
1361 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001362
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001363 args = (*Py_TYPE(it)->tp_iternext)(it);
1364 if (args == NULL)
1365 return NULL;
1366 if (!PyTuple_CheckExact(args)) {
1367 PyObject *newargs = PySequence_Tuple(args);
1368 Py_DECREF(args);
1369 if (newargs == NULL)
1370 return NULL;
1371 args = newargs;
1372 }
1373 result = PyObject_Call(lz->func, args, NULL);
1374 Py_DECREF(args);
1375 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001376}
1377
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001378PyDoc_STRVAR(starmap_doc,
1379"starmap(function, sequence) --> starmap object\n\
1380\n\
1381Return an iterator whose values are returned from the function evaluated\n\
1382with a argument tuple taken from the given sequence.");
1383
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001384static PyTypeObject starmap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001385 PyVarObject_HEAD_INIT(NULL, 0)
1386 "itertools.starmap", /* tp_name */
1387 sizeof(starmapobject), /* tp_basicsize */
1388 0, /* tp_itemsize */
1389 /* methods */
1390 (destructor)starmap_dealloc, /* tp_dealloc */
1391 0, /* tp_print */
1392 0, /* tp_getattr */
1393 0, /* tp_setattr */
1394 0, /* tp_compare */
1395 0, /* tp_repr */
1396 0, /* tp_as_number */
1397 0, /* tp_as_sequence */
1398 0, /* tp_as_mapping */
1399 0, /* tp_hash */
1400 0, /* tp_call */
1401 0, /* tp_str */
1402 PyObject_GenericGetAttr, /* tp_getattro */
1403 0, /* tp_setattro */
1404 0, /* tp_as_buffer */
1405 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1406 Py_TPFLAGS_BASETYPE, /* tp_flags */
1407 starmap_doc, /* tp_doc */
1408 (traverseproc)starmap_traverse, /* tp_traverse */
1409 0, /* tp_clear */
1410 0, /* tp_richcompare */
1411 0, /* tp_weaklistoffset */
1412 PyObject_SelfIter, /* tp_iter */
1413 (iternextfunc)starmap_next, /* tp_iternext */
1414 0, /* tp_methods */
1415 0, /* tp_members */
1416 0, /* tp_getset */
1417 0, /* tp_base */
1418 0, /* tp_dict */
1419 0, /* tp_descr_get */
1420 0, /* tp_descr_set */
1421 0, /* tp_dictoffset */
1422 0, /* tp_init */
1423 0, /* tp_alloc */
1424 starmap_new, /* tp_new */
1425 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001426};
1427
1428
1429/* imap object ************************************************************/
1430
1431typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001432 PyObject_HEAD
1433 PyObject *iters;
1434 PyObject *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001435} imapobject;
1436
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001437static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001438
1439static PyObject *
1440imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1441{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001442 PyObject *it, *iters, *func;
1443 imapobject *lz;
1444 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001445
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001446 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
1447 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001448
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001449 numargs = PyTuple_Size(args);
1450 if (numargs < 2) {
1451 PyErr_SetString(PyExc_TypeError,
1452 "imap() must have at least two arguments.");
1453 return NULL;
1454 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001455
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001456 iters = PyTuple_New(numargs-1);
1457 if (iters == NULL)
1458 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001460 for (i=1 ; i<numargs ; i++) {
1461 /* Get iterator. */
1462 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1463 if (it == NULL) {
1464 Py_DECREF(iters);
1465 return NULL;
1466 }
1467 PyTuple_SET_ITEM(iters, i-1, it);
1468 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001470 /* create imapobject structure */
1471 lz = (imapobject *)type->tp_alloc(type, 0);
1472 if (lz == NULL) {
1473 Py_DECREF(iters);
1474 return NULL;
1475 }
1476 lz->iters = iters;
1477 func = PyTuple_GET_ITEM(args, 0);
1478 Py_INCREF(func);
1479 lz->func = func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001480
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001481 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001482}
1483
1484static void
1485imap_dealloc(imapobject *lz)
1486{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001487 PyObject_GC_UnTrack(lz);
1488 Py_XDECREF(lz->iters);
1489 Py_XDECREF(lz->func);
1490 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001491}
1492
1493static int
1494imap_traverse(imapobject *lz, visitproc visit, void *arg)
1495{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001496 Py_VISIT(lz->iters);
1497 Py_VISIT(lz->func);
1498 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001499}
1500
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001501/*
Raymond Hettinger2012f172003-02-07 05:32:58 +00001502imap() is an iterator version of __builtins__.map() except that it does
1503not have the None fill-in feature. That was intentionally left out for
1504the following reasons:
1505
1506 1) Itertools are designed to be easily combined and chained together.
1507 Having all tools stop with the shortest input is a unifying principle
1508 that makes it easier to combine finite iterators (supplying data) with
1509 infinite iterators like count() and repeat() (for supplying sequential
1510 or constant arguments to a function).
1511
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001512 2) In typical use cases for combining itertools, having one finite data
Raymond Hettinger2012f172003-02-07 05:32:58 +00001513 supplier run out before another is likely to be an error condition which
1514 should not pass silently by automatically supplying None.
1515
1516 3) The use cases for automatic None fill-in are rare -- not many functions
1517 do something useful when a parameter suddenly switches type and becomes
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001518 None.
Raymond Hettinger2012f172003-02-07 05:32:58 +00001519
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001520 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001521 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001522
1523 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1524*/
1525
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001526static PyObject *
1527imap_next(imapobject *lz)
1528{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001529 PyObject *val;
1530 PyObject *argtuple;
1531 PyObject *result;
1532 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001533
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001534 numargs = PyTuple_Size(lz->iters);
1535 argtuple = PyTuple_New(numargs);
1536 if (argtuple == NULL)
1537 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001538
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001539 for (i=0 ; i<numargs ; i++) {
1540 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1541 if (val == NULL) {
1542 Py_DECREF(argtuple);
1543 return NULL;
1544 }
1545 PyTuple_SET_ITEM(argtuple, i, val);
1546 }
1547 if (lz->func == Py_None)
1548 return argtuple;
1549 result = PyObject_Call(lz->func, argtuple, NULL);
1550 Py_DECREF(argtuple);
1551 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001552}
1553
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001554PyDoc_STRVAR(imap_doc,
1555"imap(func, *iterables) --> imap object\n\
1556\n\
1557Make an iterator that computes the function using arguments from\n\
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001558each of the iterables. Like map() except that it returns\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001559an iterator instead of a list and that it stops when the shortest\n\
1560iterable is exhausted instead of filling in None for shorter\n\
1561iterables.");
1562
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001563static PyTypeObject imap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001564 PyVarObject_HEAD_INIT(NULL, 0)
1565 "itertools.imap", /* tp_name */
1566 sizeof(imapobject), /* tp_basicsize */
1567 0, /* tp_itemsize */
1568 /* methods */
1569 (destructor)imap_dealloc, /* tp_dealloc */
1570 0, /* tp_print */
1571 0, /* tp_getattr */
1572 0, /* tp_setattr */
1573 0, /* tp_compare */
1574 0, /* tp_repr */
1575 0, /* tp_as_number */
1576 0, /* tp_as_sequence */
1577 0, /* tp_as_mapping */
1578 0, /* tp_hash */
1579 0, /* tp_call */
1580 0, /* tp_str */
1581 PyObject_GenericGetAttr, /* tp_getattro */
1582 0, /* tp_setattro */
1583 0, /* tp_as_buffer */
1584 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1585 Py_TPFLAGS_BASETYPE, /* tp_flags */
1586 imap_doc, /* tp_doc */
1587 (traverseproc)imap_traverse, /* tp_traverse */
1588 0, /* tp_clear */
1589 0, /* tp_richcompare */
1590 0, /* tp_weaklistoffset */
1591 PyObject_SelfIter, /* tp_iter */
1592 (iternextfunc)imap_next, /* tp_iternext */
1593 0, /* tp_methods */
1594 0, /* tp_members */
1595 0, /* tp_getset */
1596 0, /* tp_base */
1597 0, /* tp_dict */
1598 0, /* tp_descr_get */
1599 0, /* tp_descr_set */
1600 0, /* tp_dictoffset */
1601 0, /* tp_init */
1602 0, /* tp_alloc */
1603 imap_new, /* tp_new */
1604 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001605};
1606
1607
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001608/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001609
1610typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001611 PyObject_HEAD
1612 PyObject *source; /* Iterator over input iterables */
1613 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001614} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001615
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001616static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001617
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001618static PyObject *
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001619chain_new_internal(PyTypeObject *type, PyObject *source)
1620{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001621 chainobject *lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001622
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001623 lz = (chainobject *)type->tp_alloc(type, 0);
1624 if (lz == NULL) {
1625 Py_DECREF(source);
1626 return NULL;
1627 }
1628
1629 lz->source = source;
1630 lz->active = NULL;
1631 return (PyObject *)lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001632}
1633
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001634static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001635chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001636{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001637 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001638
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001639 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1640 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001641
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001642 source = PyObject_GetIter(args);
1643 if (source == NULL)
1644 return NULL;
1645
1646 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001647}
1648
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001649static PyObject *
1650chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1651{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001652 PyObject *source;
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001653
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001654 source = PyObject_GetIter(arg);
1655 if (source == NULL)
1656 return NULL;
1657
1658 return chain_new_internal(type, source);
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001659}
1660
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001661static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001662chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001663{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001664 PyObject_GC_UnTrack(lz);
1665 Py_XDECREF(lz->active);
1666 Py_XDECREF(lz->source);
1667 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001668}
1669
Raymond Hettinger2012f172003-02-07 05:32:58 +00001670static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001671chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001672{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001673 Py_VISIT(lz->source);
1674 Py_VISIT(lz->active);
1675 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001676}
1677
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001678static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001679chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001680{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001681 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001682
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001683 if (lz->source == NULL)
1684 return NULL; /* already stopped */
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001685
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001686 if (lz->active == NULL) {
1687 PyObject *iterable = PyIter_Next(lz->source);
1688 if (iterable == NULL) {
1689 Py_CLEAR(lz->source);
1690 return NULL; /* no more input sources */
1691 }
1692 lz->active = PyObject_GetIter(iterable);
1693 Py_DECREF(iterable);
1694 if (lz->active == NULL) {
1695 Py_CLEAR(lz->source);
1696 return NULL; /* input not iterable */
1697 }
1698 }
1699 item = PyIter_Next(lz->active);
1700 if (item != NULL)
1701 return item;
1702 if (PyErr_Occurred()) {
1703 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1704 PyErr_Clear();
1705 else
1706 return NULL; /* input raised an exception */
1707 }
1708 Py_CLEAR(lz->active);
1709 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001710}
1711
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001712PyDoc_STRVAR(chain_doc,
1713"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001714\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001715Return a chain object whose .next() method returns elements from the\n\
1716first iterable until it is exhausted, then elements from the next\n\
1717iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001718
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001719PyDoc_STRVAR(chain_from_iterable_doc,
1720"chain.from_iterable(iterable) --> chain object\n\
1721\n\
1722Alternate chain() contructor taking a single iterable argument\n\
1723that evaluates lazily.");
1724
1725static PyMethodDef chain_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001726 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1727 chain_from_iterable_doc},
1728 {NULL, NULL} /* sentinel */
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001729};
1730
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001731static PyTypeObject chain_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001732 PyVarObject_HEAD_INIT(NULL, 0)
1733 "itertools.chain", /* tp_name */
1734 sizeof(chainobject), /* tp_basicsize */
1735 0, /* tp_itemsize */
1736 /* methods */
1737 (destructor)chain_dealloc, /* tp_dealloc */
1738 0, /* tp_print */
1739 0, /* tp_getattr */
1740 0, /* tp_setattr */
1741 0, /* tp_compare */
1742 0, /* tp_repr */
1743 0, /* tp_as_number */
1744 0, /* tp_as_sequence */
1745 0, /* tp_as_mapping */
1746 0, /* tp_hash */
1747 0, /* tp_call */
1748 0, /* tp_str */
1749 PyObject_GenericGetAttr, /* tp_getattro */
1750 0, /* tp_setattro */
1751 0, /* tp_as_buffer */
1752 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1753 Py_TPFLAGS_BASETYPE, /* tp_flags */
1754 chain_doc, /* tp_doc */
1755 (traverseproc)chain_traverse, /* tp_traverse */
1756 0, /* tp_clear */
1757 0, /* tp_richcompare */
1758 0, /* tp_weaklistoffset */
1759 PyObject_SelfIter, /* tp_iter */
1760 (iternextfunc)chain_next, /* tp_iternext */
1761 chain_methods, /* tp_methods */
1762 0, /* tp_members */
1763 0, /* tp_getset */
1764 0, /* tp_base */
1765 0, /* tp_dict */
1766 0, /* tp_descr_get */
1767 0, /* tp_descr_set */
1768 0, /* tp_dictoffset */
1769 0, /* tp_init */
1770 0, /* tp_alloc */
1771 chain_new, /* tp_new */
1772 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001773};
1774
1775
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001776/* product object ************************************************************/
1777
1778typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001779 PyObject_HEAD
1780 PyObject *pools; /* tuple of pool tuples */
1781 Py_ssize_t *indices; /* one index per pool */
1782 PyObject *result; /* most recently returned result tuple */
1783 int stopped; /* set to 1 when the product iterator is exhausted */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001784} productobject;
1785
1786static PyTypeObject product_type;
1787
1788static PyObject *
1789product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1790{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001791 productobject *lz;
1792 Py_ssize_t nargs, npools, repeat=1;
1793 PyObject *pools = NULL;
1794 Py_ssize_t *indices = NULL;
1795 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001796
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001797 if (kwds != NULL) {
1798 char *kwlist[] = {"repeat", 0};
1799 PyObject *tmpargs = PyTuple_New(0);
1800 if (tmpargs == NULL)
1801 return NULL;
1802 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1803 Py_DECREF(tmpargs);
1804 return NULL;
1805 }
1806 Py_DECREF(tmpargs);
1807 if (repeat < 0) {
1808 PyErr_SetString(PyExc_ValueError,
1809 "repeat argument cannot be negative");
1810 return NULL;
1811 }
1812 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001813
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001814 assert(PyTuple_Check(args));
1815 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1816 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001817
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001818 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1819 if (indices == NULL) {
1820 PyErr_NoMemory();
1821 goto error;
1822 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001823
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001824 pools = PyTuple_New(npools);
1825 if (pools == NULL)
1826 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001827
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001828 for (i=0; i < nargs ; ++i) {
1829 PyObject *item = PyTuple_GET_ITEM(args, i);
1830 PyObject *pool = PySequence_Tuple(item);
1831 if (pool == NULL)
1832 goto error;
1833 PyTuple_SET_ITEM(pools, i, pool);
1834 indices[i] = 0;
1835 }
1836 for ( ; i < npools; ++i) {
1837 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1838 Py_INCREF(pool);
1839 PyTuple_SET_ITEM(pools, i, pool);
1840 indices[i] = 0;
1841 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001842
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001843 /* create productobject structure */
1844 lz = (productobject *)type->tp_alloc(type, 0);
1845 if (lz == NULL)
1846 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001847
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001848 lz->pools = pools;
1849 lz->indices = indices;
1850 lz->result = NULL;
1851 lz->stopped = 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001852
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001853 return (PyObject *)lz;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001854
1855error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001856 if (indices != NULL)
1857 PyMem_Free(indices);
1858 Py_XDECREF(pools);
1859 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001860}
1861
1862static void
1863product_dealloc(productobject *lz)
1864{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001865 PyObject_GC_UnTrack(lz);
1866 Py_XDECREF(lz->pools);
1867 Py_XDECREF(lz->result);
1868 if (lz->indices != NULL)
1869 PyMem_Free(lz->indices);
1870 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001871}
1872
1873static int
1874product_traverse(productobject *lz, visitproc visit, void *arg)
1875{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001876 Py_VISIT(lz->pools);
1877 Py_VISIT(lz->result);
1878 return 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001879}
1880
1881static PyObject *
1882product_next(productobject *lz)
1883{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001884 PyObject *pool;
1885 PyObject *elem;
1886 PyObject *oldelem;
1887 PyObject *pools = lz->pools;
1888 PyObject *result = lz->result;
1889 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1890 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001891
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001892 if (lz->stopped)
1893 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001894
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001895 if (result == NULL) {
1896 /* On the first pass, return an initial tuple filled with the
1897 first element from each pool. */
1898 result = PyTuple_New(npools);
1899 if (result == NULL)
1900 goto empty;
1901 lz->result = result;
1902 for (i=0; i < npools; i++) {
1903 pool = PyTuple_GET_ITEM(pools, i);
1904 if (PyTuple_GET_SIZE(pool) == 0)
1905 goto empty;
1906 elem = PyTuple_GET_ITEM(pool, 0);
1907 Py_INCREF(elem);
1908 PyTuple_SET_ITEM(result, i, elem);
1909 }
1910 } else {
1911 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001912
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001913 /* Copy the previous result tuple or re-use it if available */
1914 if (Py_REFCNT(result) > 1) {
1915 PyObject *old_result = result;
1916 result = PyTuple_New(npools);
1917 if (result == NULL)
1918 goto empty;
1919 lz->result = result;
1920 for (i=0; i < npools; i++) {
1921 elem = PyTuple_GET_ITEM(old_result, i);
1922 Py_INCREF(elem);
1923 PyTuple_SET_ITEM(result, i, elem);
1924 }
1925 Py_DECREF(old_result);
1926 }
1927 /* Now, we've got the only copy so we can update it in-place */
1928 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001929
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001930 /* Update the pool indices right-to-left. Only advance to the
1931 next pool when the previous one rolls-over */
1932 for (i=npools-1 ; i >= 0 ; i--) {
1933 pool = PyTuple_GET_ITEM(pools, i);
1934 indices[i]++;
1935 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1936 /* Roll-over and advance to next pool */
1937 indices[i] = 0;
1938 elem = PyTuple_GET_ITEM(pool, 0);
1939 Py_INCREF(elem);
1940 oldelem = PyTuple_GET_ITEM(result, i);
1941 PyTuple_SET_ITEM(result, i, elem);
1942 Py_DECREF(oldelem);
1943 } else {
1944 /* No rollover. Just increment and stop here. */
1945 elem = PyTuple_GET_ITEM(pool, indices[i]);
1946 Py_INCREF(elem);
1947 oldelem = PyTuple_GET_ITEM(result, i);
1948 PyTuple_SET_ITEM(result, i, elem);
1949 Py_DECREF(oldelem);
1950 break;
1951 }
1952 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00001953
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001954 /* If i is negative, then the indices have all rolled-over
1955 and we're done. */
1956 if (i < 0)
1957 goto empty;
1958 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001959
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001960 Py_INCREF(result);
1961 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001962
1963empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001964 lz->stopped = 1;
1965 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001966}
1967
1968PyDoc_STRVAR(product_doc,
1969"product(*iterables) --> product object\n\
1970\n\
Raymond Hettinger73d79632008-02-23 02:20:41 +00001971Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001972For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1973The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1974cycle in a manner similar to an odometer (with the rightmost element changing\n\
1975on every iteration).\n\n\
Raymond Hettinger94034ea2009-02-09 18:39:41 +00001976To compute the product of an iterable with itself, specify the number\n\
1977of repetitions with the optional repeat keyword argument. For example,\n\
1978product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001979product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1980product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1981
1982static PyTypeObject product_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001983 PyVarObject_HEAD_INIT(NULL, 0)
1984 "itertools.product", /* tp_name */
1985 sizeof(productobject), /* tp_basicsize */
1986 0, /* tp_itemsize */
1987 /* methods */
1988 (destructor)product_dealloc, /* tp_dealloc */
1989 0, /* tp_print */
1990 0, /* tp_getattr */
1991 0, /* tp_setattr */
1992 0, /* tp_compare */
1993 0, /* tp_repr */
1994 0, /* tp_as_number */
1995 0, /* tp_as_sequence */
1996 0, /* tp_as_mapping */
1997 0, /* tp_hash */
1998 0, /* tp_call */
1999 0, /* tp_str */
2000 PyObject_GenericGetAttr, /* tp_getattro */
2001 0, /* tp_setattro */
2002 0, /* tp_as_buffer */
2003 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2004 Py_TPFLAGS_BASETYPE, /* tp_flags */
2005 product_doc, /* tp_doc */
2006 (traverseproc)product_traverse, /* tp_traverse */
2007 0, /* tp_clear */
2008 0, /* tp_richcompare */
2009 0, /* tp_weaklistoffset */
2010 PyObject_SelfIter, /* tp_iter */
2011 (iternextfunc)product_next, /* tp_iternext */
2012 0, /* tp_methods */
2013 0, /* tp_members */
2014 0, /* tp_getset */
2015 0, /* tp_base */
2016 0, /* tp_dict */
2017 0, /* tp_descr_get */
2018 0, /* tp_descr_set */
2019 0, /* tp_dictoffset */
2020 0, /* tp_init */
2021 0, /* tp_alloc */
2022 product_new, /* tp_new */
2023 PyObject_GC_Del, /* tp_free */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002024};
2025
2026
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002027/* combinations object ************************************************************/
2028
2029typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002030 PyObject_HEAD
2031 PyObject *pool; /* input converted to a tuple */
2032 Py_ssize_t *indices; /* one index per result element */
2033 PyObject *result; /* most recently returned result tuple */
2034 Py_ssize_t r; /* size of result tuple */
2035 int stopped; /* set to 1 when the combinations iterator is exhausted */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002036} combinationsobject;
2037
2038static PyTypeObject combinations_type;
2039
2040static PyObject *
2041combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2042{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002043 combinationsobject *co;
2044 Py_ssize_t n;
2045 Py_ssize_t r;
2046 PyObject *pool = NULL;
2047 PyObject *iterable = NULL;
2048 Py_ssize_t *indices = NULL;
2049 Py_ssize_t i;
2050 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002051
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002052 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2053 &iterable, &r))
2054 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002055
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002056 pool = PySequence_Tuple(iterable);
2057 if (pool == NULL)
2058 goto error;
2059 n = PyTuple_GET_SIZE(pool);
2060 if (r < 0) {
2061 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2062 goto error;
2063 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002064
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002065 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2066 if (indices == NULL) {
2067 PyErr_NoMemory();
2068 goto error;
2069 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002070
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002071 for (i=0 ; i<r ; i++)
2072 indices[i] = i;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002073
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002074 /* create combinationsobject structure */
2075 co = (combinationsobject *)type->tp_alloc(type, 0);
2076 if (co == NULL)
2077 goto error;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002078
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002079 co->pool = pool;
2080 co->indices = indices;
2081 co->result = NULL;
2082 co->r = r;
2083 co->stopped = r > n ? 1 : 0;
2084
2085 return (PyObject *)co;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002086
2087error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002088 if (indices != NULL)
2089 PyMem_Free(indices);
2090 Py_XDECREF(pool);
2091 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002092}
2093
2094static void
2095combinations_dealloc(combinationsobject *co)
2096{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002097 PyObject_GC_UnTrack(co);
2098 Py_XDECREF(co->pool);
2099 Py_XDECREF(co->result);
2100 if (co->indices != NULL)
2101 PyMem_Free(co->indices);
2102 Py_TYPE(co)->tp_free(co);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002103}
2104
2105static int
2106combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2107{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002108 Py_VISIT(co->pool);
2109 Py_VISIT(co->result);
2110 return 0;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002111}
2112
2113static PyObject *
2114combinations_next(combinationsobject *co)
2115{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002116 PyObject *elem;
2117 PyObject *oldelem;
2118 PyObject *pool = co->pool;
2119 Py_ssize_t *indices = co->indices;
2120 PyObject *result = co->result;
2121 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2122 Py_ssize_t r = co->r;
2123 Py_ssize_t i, j, index;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002124
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002125 if (co->stopped)
2126 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002127
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002128 if (result == NULL) {
2129 /* On the first pass, initialize result tuple using the indices */
2130 result = PyTuple_New(r);
2131 if (result == NULL)
2132 goto empty;
2133 co->result = result;
2134 for (i=0; i<r ; i++) {
2135 index = indices[i];
2136 elem = PyTuple_GET_ITEM(pool, index);
2137 Py_INCREF(elem);
2138 PyTuple_SET_ITEM(result, i, elem);
2139 }
2140 } else {
2141 /* Copy the previous result tuple or re-use it if available */
2142 if (Py_REFCNT(result) > 1) {
2143 PyObject *old_result = result;
2144 result = PyTuple_New(r);
2145 if (result == NULL)
2146 goto empty;
2147 co->result = result;
2148 for (i=0; i<r ; i++) {
2149 elem = PyTuple_GET_ITEM(old_result, i);
2150 Py_INCREF(elem);
2151 PyTuple_SET_ITEM(result, i, elem);
2152 }
2153 Py_DECREF(old_result);
2154 }
2155 /* Now, we've got the only copy so we can update it in-place
2156 * CPython's empty tuple is a singleton and cached in
2157 * PyTuple's freelist.
2158 */
2159 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002160
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002161 /* Scan indices right-to-left until finding one that is not
2162 at its maximum (i + n - r). */
2163 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2164 ;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002165
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002166 /* If i is negative, then the indices are all at
2167 their maximum value and we're done. */
2168 if (i < 0)
2169 goto empty;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002170
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002171 /* Increment the current index which we know is not at its
2172 maximum. Then move back to the right setting each index
2173 to its lowest possible value (one higher than the index
2174 to its left -- this maintains the sort order invariant). */
2175 indices[i]++;
2176 for (j=i+1 ; j<r ; j++)
2177 indices[j] = indices[j-1] + 1;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002178
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002179 /* Update the result tuple for the new indices
2180 starting with i, the leftmost index that changed */
2181 for ( ; i<r ; i++) {
2182 index = indices[i];
2183 elem = PyTuple_GET_ITEM(pool, index);
2184 Py_INCREF(elem);
2185 oldelem = PyTuple_GET_ITEM(result, i);
2186 PyTuple_SET_ITEM(result, i, elem);
2187 Py_DECREF(oldelem);
2188 }
2189 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002190
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002191 Py_INCREF(result);
2192 return result;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002193
2194empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002195 co->stopped = 1;
2196 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002197}
2198
2199PyDoc_STRVAR(combinations_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002200"combinations(iterable, r) --> combinations object\n\
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002201\n\
2202Return successive r-length combinations of elements in the iterable.\n\n\
2203combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2204
2205static PyTypeObject combinations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002206 PyVarObject_HEAD_INIT(NULL, 0)
2207 "itertools.combinations", /* tp_name */
2208 sizeof(combinationsobject), /* tp_basicsize */
2209 0, /* tp_itemsize */
2210 /* methods */
2211 (destructor)combinations_dealloc, /* tp_dealloc */
2212 0, /* tp_print */
2213 0, /* tp_getattr */
2214 0, /* tp_setattr */
2215 0, /* tp_compare */
2216 0, /* tp_repr */
2217 0, /* tp_as_number */
2218 0, /* tp_as_sequence */
2219 0, /* tp_as_mapping */
2220 0, /* tp_hash */
2221 0, /* tp_call */
2222 0, /* tp_str */
2223 PyObject_GenericGetAttr, /* tp_getattro */
2224 0, /* tp_setattro */
2225 0, /* tp_as_buffer */
2226 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2227 Py_TPFLAGS_BASETYPE, /* tp_flags */
2228 combinations_doc, /* tp_doc */
2229 (traverseproc)combinations_traverse, /* tp_traverse */
2230 0, /* tp_clear */
2231 0, /* tp_richcompare */
2232 0, /* tp_weaklistoffset */
2233 PyObject_SelfIter, /* tp_iter */
2234 (iternextfunc)combinations_next, /* tp_iternext */
2235 0, /* tp_methods */
2236 0, /* tp_members */
2237 0, /* tp_getset */
2238 0, /* tp_base */
2239 0, /* tp_dict */
2240 0, /* tp_descr_get */
2241 0, /* tp_descr_set */
2242 0, /* tp_dictoffset */
2243 0, /* tp_init */
2244 0, /* tp_alloc */
2245 combinations_new, /* tp_new */
2246 PyObject_GC_Del, /* tp_free */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002247};
2248
2249
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002250/* combinations with replacement object *******************************************/
2251
2252/* Equivalent to:
2253
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002254 def combinations_with_replacement(iterable, r):
2255 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2256 # number items returned: (n+r-1)! / r! / (n-1)!
2257 pool = tuple(iterable)
2258 n = len(pool)
2259 indices = [0] * r
2260 yield tuple(pool[i] for i in indices)
2261 while 1:
2262 for i in reversed(range(r)):
2263 if indices[i] != n - 1:
2264 break
2265 else:
2266 return
2267 indices[i:] = [indices[i] + 1] * (r - i)
2268 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002269
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002270 def combinations_with_replacement2(iterable, r):
2271 'Alternate version that filters from product()'
2272 pool = tuple(iterable)
2273 n = len(pool)
2274 for indices in product(range(n), repeat=r):
2275 if sorted(indices) == list(indices):
2276 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002277*/
2278typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002279 PyObject_HEAD
2280 PyObject *pool; /* input converted to a tuple */
2281 Py_ssize_t *indices; /* one index per result element */
2282 PyObject *result; /* most recently returned result tuple */
2283 Py_ssize_t r; /* size of result tuple */
2284 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002285} cwrobject;
2286
2287static PyTypeObject cwr_type;
2288
2289static PyObject *
2290cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2291{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002292 cwrobject *co;
2293 Py_ssize_t n;
2294 Py_ssize_t r;
2295 PyObject *pool = NULL;
2296 PyObject *iterable = NULL;
2297 Py_ssize_t *indices = NULL;
2298 Py_ssize_t i;
2299 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002300
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002301 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2302 &iterable, &r))
2303 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002304
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002305 pool = PySequence_Tuple(iterable);
2306 if (pool == NULL)
2307 goto error;
2308 n = PyTuple_GET_SIZE(pool);
2309 if (r < 0) {
2310 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2311 goto error;
2312 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002313
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002314 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2315 if (indices == NULL) {
2316 PyErr_NoMemory();
2317 goto error;
2318 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002319
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002320 for (i=0 ; i<r ; i++)
2321 indices[i] = 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002322
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002323 /* create cwrobject structure */
2324 co = (cwrobject *)type->tp_alloc(type, 0);
2325 if (co == NULL)
2326 goto error;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002327
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002328 co->pool = pool;
2329 co->indices = indices;
2330 co->result = NULL;
2331 co->r = r;
2332 co->stopped = !n && r;
2333
2334 return (PyObject *)co;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002335
2336error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002337 if (indices != NULL)
2338 PyMem_Free(indices);
2339 Py_XDECREF(pool);
2340 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002341}
2342
2343static void
2344cwr_dealloc(cwrobject *co)
2345{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002346 PyObject_GC_UnTrack(co);
2347 Py_XDECREF(co->pool);
2348 Py_XDECREF(co->result);
2349 if (co->indices != NULL)
2350 PyMem_Free(co->indices);
2351 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002352}
2353
2354static int
2355cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2356{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002357 Py_VISIT(co->pool);
2358 Py_VISIT(co->result);
2359 return 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002360}
2361
2362static PyObject *
2363cwr_next(cwrobject *co)
2364{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002365 PyObject *elem;
2366 PyObject *oldelem;
2367 PyObject *pool = co->pool;
2368 Py_ssize_t *indices = co->indices;
2369 PyObject *result = co->result;
2370 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2371 Py_ssize_t r = co->r;
2372 Py_ssize_t i, j, index;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002373
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002374 if (co->stopped)
2375 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002376
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002377 if (result == NULL) {
2378 /* On the first pass, initialize result tuple using the indices */
2379 result = PyTuple_New(r);
2380 if (result == NULL)
2381 goto empty;
2382 co->result = result;
2383 for (i=0; i<r ; i++) {
2384 index = indices[i];
2385 elem = PyTuple_GET_ITEM(pool, index);
2386 Py_INCREF(elem);
2387 PyTuple_SET_ITEM(result, i, elem);
2388 }
2389 } else {
2390 /* Copy the previous result tuple or re-use it if available */
2391 if (Py_REFCNT(result) > 1) {
2392 PyObject *old_result = result;
2393 result = PyTuple_New(r);
2394 if (result == NULL)
2395 goto empty;
2396 co->result = result;
2397 for (i=0; i<r ; i++) {
2398 elem = PyTuple_GET_ITEM(old_result, i);
2399 Py_INCREF(elem);
2400 PyTuple_SET_ITEM(result, i, elem);
2401 }
2402 Py_DECREF(old_result);
2403 }
2404 /* Now, we've got the only copy so we can update it in-place CPython's
2405 empty tuple is a singleton and cached in PyTuple's freelist. */
2406 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002407
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002408 /* Scan indices right-to-left until finding one that is not
2409 * at its maximum (n-1). */
2410 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2411 ;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002412
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002413 /* If i is negative, then the indices are all at
2414 their maximum value and we're done. */
2415 if (i < 0)
2416 goto empty;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002417
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002418 /* Increment the current index which we know is not at its
2419 maximum. Then set all to the right to the same value. */
2420 indices[i]++;
2421 for (j=i+1 ; j<r ; j++)
2422 indices[j] = indices[j-1];
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002423
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002424 /* Update the result tuple for the new indices
2425 starting with i, the leftmost index that changed */
2426 for ( ; i<r ; i++) {
2427 index = indices[i];
2428 elem = PyTuple_GET_ITEM(pool, index);
2429 Py_INCREF(elem);
2430 oldelem = PyTuple_GET_ITEM(result, i);
2431 PyTuple_SET_ITEM(result, i, elem);
2432 Py_DECREF(oldelem);
2433 }
2434 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002435
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002436 Py_INCREF(result);
2437 return result;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002438
2439empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002440 co->stopped = 1;
2441 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002442}
2443
2444PyDoc_STRVAR(cwr_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002445"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002446\n\
2447Return successive r-length combinations of elements in the iterable\n\
2448allowing individual elements to have successive repeats.\n\
2449combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2450
2451static PyTypeObject cwr_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002452 PyVarObject_HEAD_INIT(NULL, 0)
2453 "itertools.combinations_with_replacement", /* tp_name */
2454 sizeof(cwrobject), /* tp_basicsize */
2455 0, /* tp_itemsize */
2456 /* methods */
2457 (destructor)cwr_dealloc, /* tp_dealloc */
2458 0, /* tp_print */
2459 0, /* tp_getattr */
2460 0, /* tp_setattr */
2461 0, /* tp_compare */
2462 0, /* tp_repr */
2463 0, /* tp_as_number */
2464 0, /* tp_as_sequence */
2465 0, /* tp_as_mapping */
2466 0, /* tp_hash */
2467 0, /* tp_call */
2468 0, /* tp_str */
2469 PyObject_GenericGetAttr, /* tp_getattro */
2470 0, /* tp_setattro */
2471 0, /* tp_as_buffer */
2472 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2473 Py_TPFLAGS_BASETYPE, /* tp_flags */
2474 cwr_doc, /* tp_doc */
2475 (traverseproc)cwr_traverse, /* tp_traverse */
2476 0, /* tp_clear */
2477 0, /* tp_richcompare */
2478 0, /* tp_weaklistoffset */
2479 PyObject_SelfIter, /* tp_iter */
2480 (iternextfunc)cwr_next, /* tp_iternext */
2481 0, /* tp_methods */
2482 0, /* tp_members */
2483 0, /* tp_getset */
2484 0, /* tp_base */
2485 0, /* tp_dict */
2486 0, /* tp_descr_get */
2487 0, /* tp_descr_set */
2488 0, /* tp_dictoffset */
2489 0, /* tp_init */
2490 0, /* tp_alloc */
2491 cwr_new, /* tp_new */
2492 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002493};
2494
2495
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002496/* permutations object ************************************************************
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002497
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002498def permutations(iterable, r=None):
2499 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2500 pool = tuple(iterable)
2501 n = len(pool)
2502 r = n if r is None else r
2503 indices = range(n)
2504 cycles = range(n-r+1, n+1)[::-1]
2505 yield tuple(pool[i] for i in indices[:r])
2506 while n:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002507 for i in reversed(range(r)):
2508 cycles[i] -= 1
2509 if cycles[i] == 0:
2510 indices[i:] = indices[i+1:] + indices[i:i+1]
2511 cycles[i] = n - i
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002512 else:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002513 j = cycles[i]
2514 indices[i], indices[-j] = indices[-j], indices[i]
2515 yield tuple(pool[i] for i in indices[:r])
2516 break
2517 else:
2518 return
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002519*/
2520
2521typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002522 PyObject_HEAD
2523 PyObject *pool; /* input converted to a tuple */
2524 Py_ssize_t *indices; /* one index per element in the pool */
2525 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2526 PyObject *result; /* most recently returned result tuple */
2527 Py_ssize_t r; /* size of result tuple */
2528 int stopped; /* set to 1 when the permutations iterator is exhausted */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002529} permutationsobject;
2530
2531static PyTypeObject permutations_type;
2532
2533static PyObject *
2534permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2535{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002536 permutationsobject *po;
2537 Py_ssize_t n;
2538 Py_ssize_t r;
2539 PyObject *robj = Py_None;
2540 PyObject *pool = NULL;
2541 PyObject *iterable = NULL;
2542 Py_ssize_t *indices = NULL;
2543 Py_ssize_t *cycles = NULL;
2544 Py_ssize_t i;
2545 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002546
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002547 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2548 &iterable, &robj))
2549 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002550
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002551 pool = PySequence_Tuple(iterable);
2552 if (pool == NULL)
2553 goto error;
2554 n = PyTuple_GET_SIZE(pool);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002555
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002556 r = n;
2557 if (robj != Py_None) {
2558 r = PyInt_AsSsize_t(robj);
2559 if (r == -1 && PyErr_Occurred())
2560 goto error;
2561 }
2562 if (r < 0) {
2563 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2564 goto error;
2565 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002566
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002567 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2568 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2569 if (indices == NULL || cycles == NULL) {
2570 PyErr_NoMemory();
2571 goto error;
2572 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002573
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002574 for (i=0 ; i<n ; i++)
2575 indices[i] = i;
2576 for (i=0 ; i<r ; i++)
2577 cycles[i] = n - i;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002578
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002579 /* create permutationsobject structure */
2580 po = (permutationsobject *)type->tp_alloc(type, 0);
2581 if (po == NULL)
2582 goto error;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002583
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002584 po->pool = pool;
2585 po->indices = indices;
2586 po->cycles = cycles;
2587 po->result = NULL;
2588 po->r = r;
2589 po->stopped = r > n ? 1 : 0;
2590
2591 return (PyObject *)po;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002592
2593error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002594 if (indices != NULL)
2595 PyMem_Free(indices);
2596 if (cycles != NULL)
2597 PyMem_Free(cycles);
2598 Py_XDECREF(pool);
2599 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002600}
2601
2602static void
2603permutations_dealloc(permutationsobject *po)
2604{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002605 PyObject_GC_UnTrack(po);
2606 Py_XDECREF(po->pool);
2607 Py_XDECREF(po->result);
2608 PyMem_Free(po->indices);
2609 PyMem_Free(po->cycles);
2610 Py_TYPE(po)->tp_free(po);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002611}
2612
2613static int
2614permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2615{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002616 Py_VISIT(po->pool);
2617 Py_VISIT(po->result);
2618 return 0;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002619}
2620
2621static PyObject *
2622permutations_next(permutationsobject *po)
2623{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002624 PyObject *elem;
2625 PyObject *oldelem;
2626 PyObject *pool = po->pool;
2627 Py_ssize_t *indices = po->indices;
2628 Py_ssize_t *cycles = po->cycles;
2629 PyObject *result = po->result;
2630 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2631 Py_ssize_t r = po->r;
2632 Py_ssize_t i, j, k, index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002633
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002634 if (po->stopped)
2635 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002636
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002637 if (result == NULL) {
2638 /* On the first pass, initialize result tuple using the indices */
2639 result = PyTuple_New(r);
2640 if (result == NULL)
2641 goto empty;
2642 po->result = result;
2643 for (i=0; i<r ; i++) {
2644 index = indices[i];
2645 elem = PyTuple_GET_ITEM(pool, index);
2646 Py_INCREF(elem);
2647 PyTuple_SET_ITEM(result, i, elem);
2648 }
2649 } else {
2650 if (n == 0)
2651 goto empty;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002652
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002653 /* Copy the previous result tuple or re-use it if available */
2654 if (Py_REFCNT(result) > 1) {
2655 PyObject *old_result = result;
2656 result = PyTuple_New(r);
2657 if (result == NULL)
2658 goto empty;
2659 po->result = result;
2660 for (i=0; i<r ; i++) {
2661 elem = PyTuple_GET_ITEM(old_result, i);
2662 Py_INCREF(elem);
2663 PyTuple_SET_ITEM(result, i, elem);
2664 }
2665 Py_DECREF(old_result);
2666 }
2667 /* Now, we've got the only copy so we can update it in-place */
2668 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002669
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002670 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2671 for (i=r-1 ; i>=0 ; i--) {
2672 cycles[i] -= 1;
2673 if (cycles[i] == 0) {
2674 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2675 index = indices[i];
2676 for (j=i ; j<n-1 ; j++)
2677 indices[j] = indices[j+1];
2678 indices[n-1] = index;
2679 cycles[i] = n - i;
2680 } else {
2681 j = cycles[i];
2682 index = indices[i];
2683 indices[i] = indices[n-j];
2684 indices[n-j] = index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002685
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002686 for (k=i; k<r ; k++) {
2687 /* start with i, the leftmost element that changed */
2688 /* yield tuple(pool[k] for k in indices[:r]) */
2689 index = indices[k];
2690 elem = PyTuple_GET_ITEM(pool, index);
2691 Py_INCREF(elem);
2692 oldelem = PyTuple_GET_ITEM(result, k);
2693 PyTuple_SET_ITEM(result, k, elem);
2694 Py_DECREF(oldelem);
2695 }
2696 break;
2697 }
2698 }
2699 /* If i is negative, then the cycles have all
2700 rolled-over and we're done. */
2701 if (i < 0)
2702 goto empty;
2703 }
2704 Py_INCREF(result);
2705 return result;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002706
2707empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002708 po->stopped = 1;
2709 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002710}
2711
2712PyDoc_STRVAR(permutations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +00002713"permutations(iterable[, r]) --> permutations object\n\
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002714\n\
2715Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl8550d692008-06-10 12:46:39 +00002716permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002717
2718static PyTypeObject permutations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002719 PyVarObject_HEAD_INIT(NULL, 0)
2720 "itertools.permutations", /* tp_name */
2721 sizeof(permutationsobject), /* tp_basicsize */
2722 0, /* tp_itemsize */
2723 /* methods */
2724 (destructor)permutations_dealloc, /* tp_dealloc */
2725 0, /* tp_print */
2726 0, /* tp_getattr */
2727 0, /* tp_setattr */
2728 0, /* tp_compare */
2729 0, /* tp_repr */
2730 0, /* tp_as_number */
2731 0, /* tp_as_sequence */
2732 0, /* tp_as_mapping */
2733 0, /* tp_hash */
2734 0, /* tp_call */
2735 0, /* tp_str */
2736 PyObject_GenericGetAttr, /* tp_getattro */
2737 0, /* tp_setattro */
2738 0, /* tp_as_buffer */
2739 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2740 Py_TPFLAGS_BASETYPE, /* tp_flags */
2741 permutations_doc, /* tp_doc */
2742 (traverseproc)permutations_traverse, /* tp_traverse */
2743 0, /* tp_clear */
2744 0, /* tp_richcompare */
2745 0, /* tp_weaklistoffset */
2746 PyObject_SelfIter, /* tp_iter */
2747 (iternextfunc)permutations_next, /* tp_iternext */
2748 0, /* tp_methods */
2749 0, /* tp_members */
2750 0, /* tp_getset */
2751 0, /* tp_base */
2752 0, /* tp_dict */
2753 0, /* tp_descr_get */
2754 0, /* tp_descr_set */
2755 0, /* tp_dictoffset */
2756 0, /* tp_init */
2757 0, /* tp_alloc */
2758 permutations_new, /* tp_new */
2759 PyObject_GC_Del, /* tp_free */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002760};
2761
2762
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002763/* compress object ************************************************************/
2764
2765/* Equivalent to:
2766
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002767 def compress(data, selectors):
2768 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2769 return (d for d, s in izip(data, selectors) if s)
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002770*/
2771
2772typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002773 PyObject_HEAD
2774 PyObject *data;
2775 PyObject *selectors;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002776} compressobject;
2777
2778static PyTypeObject compress_type;
2779
2780static PyObject *
2781compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2782{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002783 PyObject *seq1, *seq2;
2784 PyObject *data=NULL, *selectors=NULL;
2785 compressobject *lz;
2786 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002787
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002788 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2789 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002790
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002791 data = PyObject_GetIter(seq1);
2792 if (data == NULL)
2793 goto fail;
2794 selectors = PyObject_GetIter(seq2);
2795 if (selectors == NULL)
2796 goto fail;
2797
2798 /* create compressobject structure */
2799 lz = (compressobject *)type->tp_alloc(type, 0);
2800 if (lz == NULL)
2801 goto fail;
2802 lz->data = data;
2803 lz->selectors = selectors;
2804 return (PyObject *)lz;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002805
2806fail:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002807 Py_XDECREF(data);
2808 Py_XDECREF(selectors);
2809 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002810}
2811
2812static void
2813compress_dealloc(compressobject *lz)
2814{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002815 PyObject_GC_UnTrack(lz);
2816 Py_XDECREF(lz->data);
2817 Py_XDECREF(lz->selectors);
2818 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002819}
2820
2821static int
2822compress_traverse(compressobject *lz, visitproc visit, void *arg)
2823{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002824 Py_VISIT(lz->data);
2825 Py_VISIT(lz->selectors);
2826 return 0;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002827}
2828
2829static PyObject *
2830compress_next(compressobject *lz)
2831{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002832 PyObject *data = lz->data, *selectors = lz->selectors;
2833 PyObject *datum, *selector;
2834 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2835 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2836 int ok;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002837
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002838 while (1) {
2839 /* Steps: get datum, get selector, evaluate selector.
2840 Order is important (to match the pure python version
2841 in terms of which input gets a chance to raise an
2842 exception first).
2843 */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002844
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002845 datum = datanext(data);
2846 if (datum == NULL)
2847 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002848
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002849 selector = selectornext(selectors);
2850 if (selector == NULL) {
2851 Py_DECREF(datum);
2852 return NULL;
2853 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002854
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002855 ok = PyObject_IsTrue(selector);
2856 Py_DECREF(selector);
2857 if (ok == 1)
2858 return datum;
2859 Py_DECREF(datum);
2860 if (ok == -1)
2861 return NULL;
2862 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002863}
2864
2865PyDoc_STRVAR(compress_doc,
Raymond Hettinger2e2909f2009-02-19 02:15:14 +00002866"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002867\n\
2868Return data elements corresponding to true selector elements.\n\
2869Forms a shorter iterator from selected data elements using the\n\
2870selectors to choose the data elements.");
2871
2872static PyTypeObject compress_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002873 PyVarObject_HEAD_INIT(NULL, 0)
2874 "itertools.compress", /* tp_name */
2875 sizeof(compressobject), /* tp_basicsize */
2876 0, /* tp_itemsize */
2877 /* methods */
2878 (destructor)compress_dealloc, /* tp_dealloc */
2879 0, /* tp_print */
2880 0, /* tp_getattr */
2881 0, /* tp_setattr */
2882 0, /* tp_compare */
2883 0, /* tp_repr */
2884 0, /* tp_as_number */
2885 0, /* tp_as_sequence */
2886 0, /* tp_as_mapping */
2887 0, /* tp_hash */
2888 0, /* tp_call */
2889 0, /* tp_str */
2890 PyObject_GenericGetAttr, /* tp_getattro */
2891 0, /* tp_setattro */
2892 0, /* tp_as_buffer */
2893 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2894 Py_TPFLAGS_BASETYPE, /* tp_flags */
2895 compress_doc, /* tp_doc */
2896 (traverseproc)compress_traverse, /* tp_traverse */
2897 0, /* tp_clear */
2898 0, /* tp_richcompare */
2899 0, /* tp_weaklistoffset */
2900 PyObject_SelfIter, /* tp_iter */
2901 (iternextfunc)compress_next, /* tp_iternext */
2902 0, /* tp_methods */
2903 0, /* tp_members */
2904 0, /* tp_getset */
2905 0, /* tp_base */
2906 0, /* tp_dict */
2907 0, /* tp_descr_get */
2908 0, /* tp_descr_set */
2909 0, /* tp_dictoffset */
2910 0, /* tp_init */
2911 0, /* tp_alloc */
2912 compress_new, /* tp_new */
2913 PyObject_GC_Del, /* tp_free */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002914};
2915
2916
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002917/* ifilter object ************************************************************/
2918
2919typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002920 PyObject_HEAD
2921 PyObject *func;
2922 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002923} ifilterobject;
2924
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002925static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002926
2927static PyObject *
2928ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2929{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002930 PyObject *func, *seq;
2931 PyObject *it;
2932 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002933
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002934 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
2935 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002936
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002937 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
2938 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002939
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002940 /* Get iterator. */
2941 it = PyObject_GetIter(seq);
2942 if (it == NULL)
2943 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002944
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002945 /* create ifilterobject structure */
2946 lz = (ifilterobject *)type->tp_alloc(type, 0);
2947 if (lz == NULL) {
2948 Py_DECREF(it);
2949 return NULL;
2950 }
2951 Py_INCREF(func);
2952 lz->func = func;
2953 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002954
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002955 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002956}
2957
2958static void
2959ifilter_dealloc(ifilterobject *lz)
2960{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002961 PyObject_GC_UnTrack(lz);
2962 Py_XDECREF(lz->func);
2963 Py_XDECREF(lz->it);
2964 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002965}
2966
2967static int
2968ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
2969{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002970 Py_VISIT(lz->it);
2971 Py_VISIT(lz->func);
2972 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002973}
2974
2975static PyObject *
2976ifilter_next(ifilterobject *lz)
2977{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002978 PyObject *item;
2979 PyObject *it = lz->it;
2980 long ok;
2981 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002982
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002983 iternext = *Py_TYPE(it)->tp_iternext;
2984 for (;;) {
2985 item = iternext(it);
2986 if (item == NULL)
2987 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002988
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002989 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2990 ok = PyObject_IsTrue(item);
2991 } else {
2992 PyObject *good;
2993 good = PyObject_CallFunctionObjArgs(lz->func,
2994 item, NULL);
2995 if (good == NULL) {
2996 Py_DECREF(item);
2997 return NULL;
2998 }
2999 ok = PyObject_IsTrue(good);
3000 Py_DECREF(good);
3001 }
3002 if (ok)
3003 return item;
3004 Py_DECREF(item);
3005 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003006}
3007
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003008PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003009"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003010\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003011Return those items of sequence for which function(item) is true.\n\
3012If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003013
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003014static PyTypeObject ifilter_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003015 PyVarObject_HEAD_INIT(NULL, 0)
3016 "itertools.ifilter", /* tp_name */
3017 sizeof(ifilterobject), /* tp_basicsize */
3018 0, /* tp_itemsize */
3019 /* methods */
3020 (destructor)ifilter_dealloc, /* tp_dealloc */
3021 0, /* tp_print */
3022 0, /* tp_getattr */
3023 0, /* tp_setattr */
3024 0, /* tp_compare */
3025 0, /* tp_repr */
3026 0, /* tp_as_number */
3027 0, /* tp_as_sequence */
3028 0, /* tp_as_mapping */
3029 0, /* tp_hash */
3030 0, /* tp_call */
3031 0, /* tp_str */
3032 PyObject_GenericGetAttr, /* tp_getattro */
3033 0, /* tp_setattro */
3034 0, /* tp_as_buffer */
3035 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3036 Py_TPFLAGS_BASETYPE, /* tp_flags */
3037 ifilter_doc, /* tp_doc */
3038 (traverseproc)ifilter_traverse, /* tp_traverse */
3039 0, /* tp_clear */
3040 0, /* tp_richcompare */
3041 0, /* tp_weaklistoffset */
3042 PyObject_SelfIter, /* tp_iter */
3043 (iternextfunc)ifilter_next, /* tp_iternext */
3044 0, /* tp_methods */
3045 0, /* tp_members */
3046 0, /* tp_getset */
3047 0, /* tp_base */
3048 0, /* tp_dict */
3049 0, /* tp_descr_get */
3050 0, /* tp_descr_set */
3051 0, /* tp_dictoffset */
3052 0, /* tp_init */
3053 0, /* tp_alloc */
3054 ifilter_new, /* tp_new */
3055 PyObject_GC_Del, /* tp_free */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003056};
3057
3058
3059/* ifilterfalse object ************************************************************/
3060
3061typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003062 PyObject_HEAD
3063 PyObject *func;
3064 PyObject *it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003065} ifilterfalseobject;
3066
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003067static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003068
3069static PyObject *
3070ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3071{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003072 PyObject *func, *seq;
3073 PyObject *it;
3074 ifilterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003075
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003076 if (type == &ifilterfalse_type &&
3077 !_PyArg_NoKeywords("ifilterfalse()", kwds))
3078 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003079
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003080 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
3081 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003082
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003083 /* Get iterator. */
3084 it = PyObject_GetIter(seq);
3085 if (it == NULL)
3086 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003087
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003088 /* create ifilterfalseobject structure */
3089 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
3090 if (lz == NULL) {
3091 Py_DECREF(it);
3092 return NULL;
3093 }
3094 Py_INCREF(func);
3095 lz->func = func;
3096 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003097
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003098 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003099}
3100
3101static void
3102ifilterfalse_dealloc(ifilterfalseobject *lz)
3103{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003104 PyObject_GC_UnTrack(lz);
3105 Py_XDECREF(lz->func);
3106 Py_XDECREF(lz->it);
3107 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003108}
3109
3110static int
3111ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
3112{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003113 Py_VISIT(lz->it);
3114 Py_VISIT(lz->func);
3115 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003116}
3117
3118static PyObject *
3119ifilterfalse_next(ifilterfalseobject *lz)
3120{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003121 PyObject *item;
3122 PyObject *it = lz->it;
3123 long ok;
3124 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003125
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003126 iternext = *Py_TYPE(it)->tp_iternext;
3127 for (;;) {
3128 item = iternext(it);
3129 if (item == NULL)
3130 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003131
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003132 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3133 ok = PyObject_IsTrue(item);
3134 } else {
3135 PyObject *good;
3136 good = PyObject_CallFunctionObjArgs(lz->func,
3137 item, NULL);
3138 if (good == NULL) {
3139 Py_DECREF(item);
3140 return NULL;
3141 }
3142 ok = PyObject_IsTrue(good);
3143 Py_DECREF(good);
3144 }
3145 if (!ok)
3146 return item;
3147 Py_DECREF(item);
3148 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003149}
3150
Raymond Hettinger60eca932003-02-09 06:40:58 +00003151PyDoc_STRVAR(ifilterfalse_doc,
3152"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
3153\n\
3154Return those items of sequence for which function(item) is false.\n\
3155If function is None, return the items that are false.");
3156
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003157static PyTypeObject ifilterfalse_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003158 PyVarObject_HEAD_INIT(NULL, 0)
3159 "itertools.ifilterfalse", /* tp_name */
3160 sizeof(ifilterfalseobject), /* tp_basicsize */
3161 0, /* tp_itemsize */
3162 /* methods */
3163 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
3164 0, /* tp_print */
3165 0, /* tp_getattr */
3166 0, /* tp_setattr */
3167 0, /* tp_compare */
3168 0, /* tp_repr */
3169 0, /* tp_as_number */
3170 0, /* tp_as_sequence */
3171 0, /* tp_as_mapping */
3172 0, /* tp_hash */
3173 0, /* tp_call */
3174 0, /* tp_str */
3175 PyObject_GenericGetAttr, /* tp_getattro */
3176 0, /* tp_setattro */
3177 0, /* tp_as_buffer */
3178 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3179 Py_TPFLAGS_BASETYPE, /* tp_flags */
3180 ifilterfalse_doc, /* tp_doc */
3181 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
3182 0, /* tp_clear */
3183 0, /* tp_richcompare */
3184 0, /* tp_weaklistoffset */
3185 PyObject_SelfIter, /* tp_iter */
3186 (iternextfunc)ifilterfalse_next, /* tp_iternext */
3187 0, /* tp_methods */
3188 0, /* tp_members */
3189 0, /* tp_getset */
3190 0, /* tp_base */
3191 0, /* tp_dict */
3192 0, /* tp_descr_get */
3193 0, /* tp_descr_set */
3194 0, /* tp_dictoffset */
3195 0, /* tp_init */
3196 0, /* tp_alloc */
3197 ifilterfalse_new, /* tp_new */
3198 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003199};
3200
3201
3202/* count object ************************************************************/
3203
3204typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003205 PyObject_HEAD
3206 Py_ssize_t cnt;
3207 PyObject *long_cnt;
3208 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003209} countobject;
3210
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003211/* Counting logic and invariants:
3212
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003213fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003214
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003215 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3216 Advances with: cnt += 1
3217 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003218
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003219slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003220
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003221 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3222 All counting is done with python objects (no overflows or underflows).
3223 Advances with: long_cnt += long_step
3224 Step may be zero -- effectively a slow version of repeat(cnt).
3225 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003226*/
3227
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003228static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003229
3230static PyObject *
3231count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3232{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003233 countobject *lz;
3234 int slow_mode = 0;
3235 Py_ssize_t cnt = 0;
3236 PyObject *long_cnt = NULL;
3237 PyObject *long_step = NULL;
3238 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003239
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003240 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3241 kwlist, &long_cnt, &long_step))
3242 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003243
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003244 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3245 (long_step != NULL && !PyNumber_Check(long_step))) {
3246 PyErr_SetString(PyExc_TypeError, "a number is required");
3247 return NULL;
3248 }
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003249
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003250 if (long_cnt != NULL) {
3251 cnt = PyInt_AsSsize_t(long_cnt);
3252 if ((cnt == -1 && PyErr_Occurred()) || !PyInt_Check(long_cnt)) {
3253 PyErr_Clear();
3254 slow_mode = 1;
3255 }
3256 Py_INCREF(long_cnt);
3257 } else {
3258 cnt = 0;
3259 long_cnt = PyInt_FromLong(0);
3260 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003261
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003262 /* If not specified, step defaults to 1 */
3263 if (long_step == NULL) {
3264 long_step = PyInt_FromLong(1);
3265 if (long_step == NULL) {
3266 Py_DECREF(long_cnt);
3267 return NULL;
3268 }
3269 } else
3270 Py_INCREF(long_step);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003271
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003272 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003273
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003274 /* Fast mode only works when the step is 1 */
3275 if (!PyInt_Check(long_step) ||
3276 PyInt_AS_LONG(long_step) != 1) {
3277 slow_mode = 1;
3278 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003279
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003280 if (slow_mode)
3281 cnt = PY_SSIZE_T_MAX;
3282 else
3283 Py_CLEAR(long_cnt);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003284
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003285 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3286 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3287 assert(slow_mode ||
3288 (PyInt_Check(long_step) && PyInt_AS_LONG(long_step) == 1));
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003289
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003290 /* create countobject structure */
3291 lz = (countobject *)type->tp_alloc(type, 0);
3292 if (lz == NULL) {
3293 Py_XDECREF(long_cnt);
3294 return NULL;
3295 }
3296 lz->cnt = cnt;
3297 lz->long_cnt = long_cnt;
3298 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003299
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003300 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003301}
3302
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003303static void
3304count_dealloc(countobject *lz)
3305{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003306 PyObject_GC_UnTrack(lz);
3307 Py_XDECREF(lz->long_cnt);
3308 Py_XDECREF(lz->long_step);
3309 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003310}
3311
Benjamin Peterson062a7c32009-02-16 21:07:52 +00003312static int
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003313count_traverse(countobject *lz, visitproc visit, void *arg)
3314{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003315 Py_VISIT(lz->long_cnt);
3316 Py_VISIT(lz->long_step);
3317 return 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003318}
3319
3320static PyObject *
3321count_nextlong(countobject *lz)
3322{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003323 PyObject *long_cnt;
3324 PyObject *stepped_up;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003325
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003326 long_cnt = lz->long_cnt;
3327 if (long_cnt == NULL) {
3328 /* Switch to slow_mode */
3329 long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
3330 if (long_cnt == NULL)
3331 return NULL;
3332 }
3333 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003334
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003335 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3336 if (stepped_up == NULL)
3337 return NULL;
3338 lz->long_cnt = stepped_up;
3339 return long_cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003340}
3341
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003342static PyObject *
3343count_next(countobject *lz)
3344{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003345 if (lz->cnt == PY_SSIZE_T_MAX)
3346 return count_nextlong(lz);
3347 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003348}
3349
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003350static PyObject *
3351count_repr(countobject *lz)
3352{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003353 PyObject *cnt_repr, *step_repr = NULL;
3354 PyObject *result = NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003355
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003356 if (lz->cnt != PY_SSIZE_T_MAX)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003357 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003358
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003359 cnt_repr = PyObject_Repr(lz->long_cnt);
3360 if (cnt_repr == NULL)
3361 return NULL;
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003362
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003363 if (PyInt_Check(lz->long_step) && PyInt_AS_LONG(lz->long_step) == 1) {
3364 /* Don't display step when it is an integer equal to 1 */
3365 result = PyString_FromFormat("count(%s)",
3366 PyString_AS_STRING(cnt_repr));
3367 } else {
3368 step_repr = PyObject_Repr(lz->long_step);
3369 if (step_repr != NULL)
3370 result = PyString_FromFormat("count(%s, %s)",
3371 PyString_AS_STRING(cnt_repr),
3372 PyString_AS_STRING(step_repr));
3373 }
3374 Py_DECREF(cnt_repr);
3375 Py_XDECREF(step_repr);
3376 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003377}
3378
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003379static PyObject *
3380count_reduce(countobject *lz)
3381{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003382 if (lz->cnt == PY_SSIZE_T_MAX)
3383 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3384 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003385}
3386
3387PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3388
3389static PyMethodDef count_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003390 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3391 count_reduce_doc},
3392 {NULL, NULL} /* sentinel */
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003393};
3394
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003395PyDoc_STRVAR(count_doc,
Ezio Melottifdc1e0d2010-06-11 02:21:25 +00003396 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003397\n\
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003398Return a count object whose .next() method returns consecutive values.\n\
3399Equivalent to:\n\n\
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003400 def count(firstval=0, step=1):\n\
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003401 x = firstval\n\
3402 while 1:\n\
3403 yield x\n\
3404 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003405
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003406static PyTypeObject count_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003407 PyVarObject_HEAD_INIT(NULL, 0)
3408 "itertools.count", /* tp_name */
3409 sizeof(countobject), /* tp_basicsize */
3410 0, /* tp_itemsize */
3411 /* methods */
3412 (destructor)count_dealloc, /* tp_dealloc */
3413 0, /* tp_print */
3414 0, /* tp_getattr */
3415 0, /* tp_setattr */
3416 0, /* tp_compare */
3417 (reprfunc)count_repr, /* tp_repr */
3418 0, /* tp_as_number */
3419 0, /* tp_as_sequence */
3420 0, /* tp_as_mapping */
3421 0, /* tp_hash */
3422 0, /* tp_call */
3423 0, /* tp_str */
3424 PyObject_GenericGetAttr, /* tp_getattro */
3425 0, /* tp_setattro */
3426 0, /* tp_as_buffer */
3427 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3428 Py_TPFLAGS_BASETYPE, /* tp_flags */
3429 count_doc, /* tp_doc */
3430 (traverseproc)count_traverse, /* tp_traverse */
3431 0, /* tp_clear */
3432 0, /* tp_richcompare */
3433 0, /* tp_weaklistoffset */
3434 PyObject_SelfIter, /* tp_iter */
3435 (iternextfunc)count_next, /* tp_iternext */
3436 count_methods, /* tp_methods */
3437 0, /* tp_members */
3438 0, /* tp_getset */
3439 0, /* tp_base */
3440 0, /* tp_dict */
3441 0, /* tp_descr_get */
3442 0, /* tp_descr_set */
3443 0, /* tp_dictoffset */
3444 0, /* tp_init */
3445 0, /* tp_alloc */
3446 count_new, /* tp_new */
3447 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003448};
3449
3450
3451/* izip object ************************************************************/
3452
3453#include "Python.h"
3454
3455typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003456 PyObject_HEAD
3457 Py_ssize_t tuplesize;
3458 PyObject *ittuple; /* tuple of iterators */
3459 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003460} izipobject;
3461
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003462static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003463
3464static PyObject *
3465izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3466{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003467 izipobject *lz;
3468 Py_ssize_t i;
3469 PyObject *ittuple; /* tuple of iterators */
3470 PyObject *result;
3471 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003472
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003473 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
3474 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003475
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003476 /* args must be a tuple */
3477 assert(PyTuple_Check(args));
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003478
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003479 /* obtain iterators */
3480 ittuple = PyTuple_New(tuplesize);
3481 if (ittuple == NULL)
3482 return NULL;
3483 for (i=0; i < tuplesize; ++i) {
3484 PyObject *item = PyTuple_GET_ITEM(args, i);
3485 PyObject *it = PyObject_GetIter(item);
3486 if (it == NULL) {
3487 if (PyErr_ExceptionMatches(PyExc_TypeError))
3488 PyErr_Format(PyExc_TypeError,
3489 "izip argument #%zd must support iteration",
3490 i+1);
3491 Py_DECREF(ittuple);
3492 return NULL;
3493 }
3494 PyTuple_SET_ITEM(ittuple, i, it);
3495 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003496
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003497 /* create a result holder */
3498 result = PyTuple_New(tuplesize);
3499 if (result == NULL) {
3500 Py_DECREF(ittuple);
3501 return NULL;
3502 }
3503 for (i=0 ; i < tuplesize ; i++) {
3504 Py_INCREF(Py_None);
3505 PyTuple_SET_ITEM(result, i, Py_None);
3506 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003507
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003508 /* create izipobject structure */
3509 lz = (izipobject *)type->tp_alloc(type, 0);
3510 if (lz == NULL) {
3511 Py_DECREF(ittuple);
3512 Py_DECREF(result);
3513 return NULL;
3514 }
3515 lz->ittuple = ittuple;
3516 lz->tuplesize = tuplesize;
3517 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003518
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003519 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003520}
3521
3522static void
3523izip_dealloc(izipobject *lz)
3524{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003525 PyObject_GC_UnTrack(lz);
3526 Py_XDECREF(lz->ittuple);
3527 Py_XDECREF(lz->result);
3528 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003529}
3530
3531static int
3532izip_traverse(izipobject *lz, visitproc visit, void *arg)
3533{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003534 Py_VISIT(lz->ittuple);
3535 Py_VISIT(lz->result);
3536 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003537}
3538
3539static PyObject *
3540izip_next(izipobject *lz)
3541{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003542 Py_ssize_t i;
3543 Py_ssize_t tuplesize = lz->tuplesize;
3544 PyObject *result = lz->result;
3545 PyObject *it;
3546 PyObject *item;
3547 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003548
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003549 if (tuplesize == 0)
3550 return NULL;
3551 if (Py_REFCNT(result) == 1) {
3552 Py_INCREF(result);
3553 for (i=0 ; i < tuplesize ; i++) {
3554 it = PyTuple_GET_ITEM(lz->ittuple, i);
3555 item = (*Py_TYPE(it)->tp_iternext)(it);
3556 if (item == NULL) {
3557 Py_DECREF(result);
3558 return NULL;
3559 }
3560 olditem = PyTuple_GET_ITEM(result, i);
3561 PyTuple_SET_ITEM(result, i, item);
3562 Py_DECREF(olditem);
3563 }
3564 } else {
3565 result = PyTuple_New(tuplesize);
3566 if (result == NULL)
3567 return NULL;
3568 for (i=0 ; i < tuplesize ; i++) {
3569 it = PyTuple_GET_ITEM(lz->ittuple, i);
3570 item = (*Py_TYPE(it)->tp_iternext)(it);
3571 if (item == NULL) {
3572 Py_DECREF(result);
3573 return NULL;
3574 }
3575 PyTuple_SET_ITEM(result, i, item);
3576 }
3577 }
3578 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003579}
3580
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003581PyDoc_STRVAR(izip_doc,
3582"izip(iter1 [,iter2 [...]]) --> izip object\n\
3583\n\
3584Return a izip object whose .next() method returns a tuple where\n\
3585the i-th element comes from the i-th iterable argument. The .next()\n\
3586method continues until the shortest iterable in the argument sequence\n\
3587is exhausted and then it raises StopIteration. Works like the zip()\n\
3588function but consumes less memory by returning an iterator instead of\n\
3589a list.");
3590
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003591static PyTypeObject izip_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003592 PyVarObject_HEAD_INIT(NULL, 0)
3593 "itertools.izip", /* tp_name */
3594 sizeof(izipobject), /* tp_basicsize */
3595 0, /* tp_itemsize */
3596 /* methods */
3597 (destructor)izip_dealloc, /* tp_dealloc */
3598 0, /* tp_print */
3599 0, /* tp_getattr */
3600 0, /* tp_setattr */
3601 0, /* tp_compare */
3602 0, /* tp_repr */
3603 0, /* tp_as_number */
3604 0, /* tp_as_sequence */
3605 0, /* tp_as_mapping */
3606 0, /* tp_hash */
3607 0, /* tp_call */
3608 0, /* tp_str */
3609 PyObject_GenericGetAttr, /* tp_getattro */
3610 0, /* tp_setattro */
3611 0, /* tp_as_buffer */
3612 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3613 Py_TPFLAGS_BASETYPE, /* tp_flags */
3614 izip_doc, /* tp_doc */
3615 (traverseproc)izip_traverse, /* tp_traverse */
3616 0, /* tp_clear */
3617 0, /* tp_richcompare */
3618 0, /* tp_weaklistoffset */
3619 PyObject_SelfIter, /* tp_iter */
3620 (iternextfunc)izip_next, /* tp_iternext */
3621 0, /* tp_methods */
3622 0, /* tp_members */
3623 0, /* tp_getset */
3624 0, /* tp_base */
3625 0, /* tp_dict */
3626 0, /* tp_descr_get */
3627 0, /* tp_descr_set */
3628 0, /* tp_dictoffset */
3629 0, /* tp_init */
3630 0, /* tp_alloc */
3631 izip_new, /* tp_new */
3632 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003633};
3634
3635
3636/* repeat object ************************************************************/
3637
3638typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003639 PyObject_HEAD
3640 PyObject *element;
3641 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003642} repeatobject;
3643
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003644static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003645
3646static PyObject *
3647repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3648{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003649 repeatobject *ro;
3650 PyObject *element;
3651 Py_ssize_t cnt = -1;
3652 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003653
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003654 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3655 &element, &cnt))
3656 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003657
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003658 if (PyTuple_Size(args) == 2 && cnt < 0)
3659 cnt = 0;
3660
3661 ro = (repeatobject *)type->tp_alloc(type, 0);
3662 if (ro == NULL)
3663 return NULL;
3664 Py_INCREF(element);
3665 ro->element = element;
3666 ro->cnt = cnt;
3667 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003668}
3669
3670static void
3671repeat_dealloc(repeatobject *ro)
3672{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003673 PyObject_GC_UnTrack(ro);
3674 Py_XDECREF(ro->element);
3675 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003676}
3677
3678static int
3679repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3680{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003681 Py_VISIT(ro->element);
3682 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003683}
3684
3685static PyObject *
3686repeat_next(repeatobject *ro)
3687{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003688 if (ro->cnt == 0)
3689 return NULL;
3690 if (ro->cnt > 0)
3691 ro->cnt--;
3692 Py_INCREF(ro->element);
3693 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003694}
3695
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003696static PyObject *
3697repeat_repr(repeatobject *ro)
3698{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003699 PyObject *result, *objrepr;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003700
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003701 objrepr = PyObject_Repr(ro->element);
3702 if (objrepr == NULL)
3703 return NULL;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003704
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003705 if (ro->cnt == -1)
3706 result = PyString_FromFormat("repeat(%s)",
3707 PyString_AS_STRING(objrepr));
3708 else
3709 result = PyString_FromFormat("repeat(%s, %zd)",
3710 PyString_AS_STRING(objrepr), ro->cnt);
3711 Py_DECREF(objrepr);
3712 return result;
3713}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003714
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003715static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003716repeat_len(repeatobject *ro)
3717{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003718 if (ro->cnt == -1) {
3719 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3720 return NULL;
3721 }
3722 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003723}
3724
Armin Rigof5b3e362006-02-11 21:32:43 +00003725PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003726
3727static PyMethodDef repeat_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003728 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3729 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003730};
3731
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003732PyDoc_STRVAR(repeat_doc,
Raymond Hettinger182edae2009-02-19 02:38:25 +00003733"repeat(object [,times]) -> create an iterator which returns the object\n\
3734for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003735endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003736
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003737static PyTypeObject repeat_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003738 PyVarObject_HEAD_INIT(NULL, 0)
3739 "itertools.repeat", /* tp_name */
3740 sizeof(repeatobject), /* tp_basicsize */
3741 0, /* tp_itemsize */
3742 /* methods */
3743 (destructor)repeat_dealloc, /* tp_dealloc */
3744 0, /* tp_print */
3745 0, /* tp_getattr */
3746 0, /* tp_setattr */
3747 0, /* tp_compare */
3748 (reprfunc)repeat_repr, /* tp_repr */
3749 0, /* tp_as_number */
3750 0, /* tp_as_sequence */
3751 0, /* tp_as_mapping */
3752 0, /* tp_hash */
3753 0, /* tp_call */
3754 0, /* tp_str */
3755 PyObject_GenericGetAttr, /* tp_getattro */
3756 0, /* tp_setattro */
3757 0, /* tp_as_buffer */
3758 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3759 Py_TPFLAGS_BASETYPE, /* tp_flags */
3760 repeat_doc, /* tp_doc */
3761 (traverseproc)repeat_traverse, /* tp_traverse */
3762 0, /* tp_clear */
3763 0, /* tp_richcompare */
3764 0, /* tp_weaklistoffset */
3765 PyObject_SelfIter, /* tp_iter */
3766 (iternextfunc)repeat_next, /* tp_iternext */
3767 repeat_methods, /* tp_methods */
3768 0, /* tp_members */
3769 0, /* tp_getset */
3770 0, /* tp_base */
3771 0, /* tp_dict */
3772 0, /* tp_descr_get */
3773 0, /* tp_descr_set */
3774 0, /* tp_dictoffset */
3775 0, /* tp_init */
3776 0, /* tp_alloc */
3777 repeat_new, /* tp_new */
3778 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003779};
3780
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003781/* iziplongest object ************************************************************/
3782
3783#include "Python.h"
3784
3785typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003786 PyObject_HEAD
3787 Py_ssize_t tuplesize;
3788 Py_ssize_t numactive;
3789 PyObject *ittuple; /* tuple of iterators */
3790 PyObject *result;
3791 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003792} iziplongestobject;
3793
3794static PyTypeObject iziplongest_type;
3795
3796static PyObject *
3797izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3798{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003799 iziplongestobject *lz;
3800 Py_ssize_t i;
3801 PyObject *ittuple; /* tuple of iterators */
3802 PyObject *result;
3803 PyObject *fillvalue = Py_None;
3804 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003805
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003806 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3807 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3808 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3809 PyErr_SetString(PyExc_TypeError,
3810 "izip_longest() got an unexpected keyword argument");
3811 return NULL;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003812 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003813 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003814
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003815 /* args must be a tuple */
3816 assert(PyTuple_Check(args));
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003817
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003818 /* obtain iterators */
3819 ittuple = PyTuple_New(tuplesize);
3820 if (ittuple == NULL)
3821 return NULL;
3822 for (i=0; i < tuplesize; ++i) {
3823 PyObject *item = PyTuple_GET_ITEM(args, i);
3824 PyObject *it = PyObject_GetIter(item);
3825 if (it == NULL) {
3826 if (PyErr_ExceptionMatches(PyExc_TypeError))
3827 PyErr_Format(PyExc_TypeError,
3828 "izip_longest argument #%zd must support iteration",
3829 i+1);
3830 Py_DECREF(ittuple);
3831 return NULL;
3832 }
3833 PyTuple_SET_ITEM(ittuple, i, it);
3834 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003835
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003836 /* create a result holder */
3837 result = PyTuple_New(tuplesize);
3838 if (result == NULL) {
3839 Py_DECREF(ittuple);
3840 return NULL;
3841 }
3842 for (i=0 ; i < tuplesize ; i++) {
3843 Py_INCREF(Py_None);
3844 PyTuple_SET_ITEM(result, i, Py_None);
3845 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003846
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003847 /* create iziplongestobject structure */
3848 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3849 if (lz == NULL) {
3850 Py_DECREF(ittuple);
3851 Py_DECREF(result);
3852 return NULL;
3853 }
3854 lz->ittuple = ittuple;
3855 lz->tuplesize = tuplesize;
3856 lz->numactive = tuplesize;
3857 lz->result = result;
3858 Py_INCREF(fillvalue);
3859 lz->fillvalue = fillvalue;
3860 return (PyObject *)lz;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003861}
3862
3863static void
3864izip_longest_dealloc(iziplongestobject *lz)
3865{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003866 PyObject_GC_UnTrack(lz);
3867 Py_XDECREF(lz->ittuple);
3868 Py_XDECREF(lz->result);
3869 Py_XDECREF(lz->fillvalue);
3870 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003871}
3872
3873static int
3874izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3875{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003876 Py_VISIT(lz->ittuple);
3877 Py_VISIT(lz->result);
3878 Py_VISIT(lz->fillvalue);
3879 return 0;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003880}
3881
3882static PyObject *
3883izip_longest_next(iziplongestobject *lz)
3884{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003885 Py_ssize_t i;
3886 Py_ssize_t tuplesize = lz->tuplesize;
3887 PyObject *result = lz->result;
3888 PyObject *it;
3889 PyObject *item;
3890 PyObject *olditem;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003891
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003892 if (tuplesize == 0)
3893 return NULL;
3894 if (lz->numactive == 0)
3895 return NULL;
3896 if (Py_REFCNT(result) == 1) {
3897 Py_INCREF(result);
3898 for (i=0 ; i < tuplesize ; i++) {
3899 it = PyTuple_GET_ITEM(lz->ittuple, i);
3900 if (it == NULL) {
3901 Py_INCREF(lz->fillvalue);
3902 item = lz->fillvalue;
3903 } else {
3904 item = PyIter_Next(it);
3905 if (item == NULL) {
3906 lz->numactive -= 1;
3907 if (lz->numactive == 0 || PyErr_Occurred()) {
3908 lz->numactive = 0;
3909 Py_DECREF(result);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003910 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003911 } else {
3912 Py_INCREF(lz->fillvalue);
3913 item = lz->fillvalue;
3914 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3915 Py_DECREF(it);
3916 }
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003917 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003918 }
3919 olditem = PyTuple_GET_ITEM(result, i);
3920 PyTuple_SET_ITEM(result, i, item);
3921 Py_DECREF(olditem);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003922 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003923 } else {
3924 result = PyTuple_New(tuplesize);
3925 if (result == NULL)
3926 return NULL;
3927 for (i=0 ; i < tuplesize ; i++) {
3928 it = PyTuple_GET_ITEM(lz->ittuple, i);
3929 if (it == NULL) {
3930 Py_INCREF(lz->fillvalue);
3931 item = lz->fillvalue;
3932 } else {
3933 item = PyIter_Next(it);
3934 if (item == NULL) {
3935 lz->numactive -= 1;
3936 if (lz->numactive == 0 || PyErr_Occurred()) {
3937 lz->numactive = 0;
3938 Py_DECREF(result);
3939 return NULL;
3940 } else {
3941 Py_INCREF(lz->fillvalue);
3942 item = lz->fillvalue;
3943 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3944 Py_DECREF(it);
3945 }
3946 }
3947 }
3948 PyTuple_SET_ITEM(result, i, item);
3949 }
3950 }
3951 return result;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003952}
3953
3954PyDoc_STRVAR(izip_longest_doc,
3955"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3956\n\
3957Return an izip_longest object whose .next() method returns a tuple where\n\
3958the i-th element comes from the i-th iterable argument. The .next()\n\
3959method continues until the longest iterable in the argument sequence\n\
3960is exhausted and then it raises StopIteration. When the shorter iterables\n\
3961are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3962defaults to None or can be specified by a keyword argument.\n\
3963");
3964
3965static PyTypeObject iziplongest_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003966 PyVarObject_HEAD_INIT(NULL, 0)
3967 "itertools.izip_longest", /* tp_name */
3968 sizeof(iziplongestobject), /* tp_basicsize */
3969 0, /* tp_itemsize */
3970 /* methods */
3971 (destructor)izip_longest_dealloc, /* tp_dealloc */
3972 0, /* tp_print */
3973 0, /* tp_getattr */
3974 0, /* tp_setattr */
3975 0, /* tp_compare */
3976 0, /* tp_repr */
3977 0, /* tp_as_number */
3978 0, /* tp_as_sequence */
3979 0, /* tp_as_mapping */
3980 0, /* tp_hash */
3981 0, /* tp_call */
3982 0, /* tp_str */
3983 PyObject_GenericGetAttr, /* tp_getattro */
3984 0, /* tp_setattro */
3985 0, /* tp_as_buffer */
3986 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3987 Py_TPFLAGS_BASETYPE, /* tp_flags */
3988 izip_longest_doc, /* tp_doc */
3989 (traverseproc)izip_longest_traverse, /* tp_traverse */
3990 0, /* tp_clear */
3991 0, /* tp_richcompare */
3992 0, /* tp_weaklistoffset */
3993 PyObject_SelfIter, /* tp_iter */
3994 (iternextfunc)izip_longest_next, /* tp_iternext */
3995 0, /* tp_methods */
3996 0, /* tp_members */
3997 0, /* tp_getset */
3998 0, /* tp_base */
3999 0, /* tp_dict */
4000 0, /* tp_descr_get */
4001 0, /* tp_descr_set */
4002 0, /* tp_dictoffset */
4003 0, /* tp_init */
4004 0, /* tp_alloc */
4005 izip_longest_new, /* tp_new */
4006 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd36862c2007-02-21 05:20:38 +00004007};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004008
4009/* module level code ********************************************************/
4010
4011PyDoc_STRVAR(module_doc,
4012"Functional tools for creating and using iterators.\n\
4013\n\
4014Infinite iterators:\n\
4015count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004016cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004017repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004018\n\
4019Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004020chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettingere76816b2009-01-29 03:43:44 +00004021compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004022dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4023groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00004024ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
4025ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004026islice(seq, [start,] stop [, step]) --> elements from\n\
4027 seq[start:stop:step]\n\
4028imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
4029starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004030tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004031takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004032izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4033izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4034\n\
4035Combinatoric generators:\n\
4036product(p, q, ... [repeat=1]) --> cartesian product\n\
4037permutations(p[, r])\n\
Raymond Hettinger9eac1192009-11-19 01:22:04 +00004038combinations(p, r)\n\
4039combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004040");
4041
4042
Raymond Hettingerad983e72003-11-12 14:32:26 +00004043static PyMethodDef module_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004044 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4045 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004046};
4047
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004048PyMODINIT_FUNC
4049inititertools(void)
4050{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004051 int i;
4052 PyObject *m;
4053 char *name;
4054 PyTypeObject *typelist[] = {
4055 &combinations_type,
4056 &cwr_type,
4057 &cycle_type,
4058 &dropwhile_type,
4059 &takewhile_type,
4060 &islice_type,
4061 &starmap_type,
4062 &imap_type,
4063 &chain_type,
4064 &compress_type,
4065 &ifilter_type,
4066 &ifilterfalse_type,
4067 &count_type,
4068 &izip_type,
4069 &iziplongest_type,
4070 &permutations_type,
4071 &product_type,
4072 &repeat_type,
4073 &groupby_type,
4074 NULL
4075 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004076
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004077 Py_TYPE(&teedataobject_type) = &PyType_Type;
4078 m = Py_InitModule3("itertools", module_methods, module_doc);
4079 if (m == NULL)
4080 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004081
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004082 for (i=0 ; typelist[i] != NULL ; i++) {
4083 if (PyType_Ready(typelist[i]) < 0)
4084 return;
4085 name = strchr(typelist[i]->tp_name, '.');
4086 assert (name != NULL);
4087 Py_INCREF(typelist[i]);
4088 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4089 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004090
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004091 if (PyType_Ready(&teedataobject_type) < 0)
4092 return;
4093 if (PyType_Ready(&tee_type) < 0)
4094 return;
4095 if (PyType_Ready(&_grouper_type) < 0)
4096 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004097}