blob: ea822913662bf241f41ad2862857dfea61a54382 [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
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200404static void
405teedataobject_safe_decref(PyObject *obj)
406{
407 while (obj && Py_TYPE(obj) == &teedataobject_type &&
408 Py_REFCNT(obj) == 1) {
409 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
410 ((teedataobject *)obj)->nextlink = NULL;
411 Py_DECREF(obj);
412 obj = nextlink;
413 }
414 Py_XDECREF(obj);
415}
416
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000417static int
418teedataobject_clear(teedataobject *tdo)
419{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000420 int i;
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200421 PyObject *tmp;
422
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000423 Py_CLEAR(tdo->it);
424 for (i=0 ; i<tdo->numread ; i++)
425 Py_CLEAR(tdo->values[i]);
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200426 tmp = tdo->nextlink;
427 tdo->nextlink = NULL;
428 teedataobject_safe_decref(tmp);
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000429 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000430}
431
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000432static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000433teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000434{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000435 PyObject_GC_UnTrack(tdo);
436 teedataobject_clear(tdo);
437 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000438}
439
Raymond Hettingerad983e72003-11-12 14:32:26 +0000440PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000441
Raymond Hettingerad983e72003-11-12 14:32:26 +0000442static PyTypeObject teedataobject_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000443 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
444 "itertools.tee_dataobject", /* tp_name */
445 sizeof(teedataobject), /* tp_basicsize */
446 0, /* tp_itemsize */
447 /* methods */
448 (destructor)teedataobject_dealloc, /* tp_dealloc */
449 0, /* tp_print */
450 0, /* tp_getattr */
451 0, /* tp_setattr */
452 0, /* tp_compare */
453 0, /* tp_repr */
454 0, /* tp_as_number */
455 0, /* tp_as_sequence */
456 0, /* tp_as_mapping */
457 0, /* tp_hash */
458 0, /* tp_call */
459 0, /* tp_str */
460 PyObject_GenericGetAttr, /* tp_getattro */
461 0, /* tp_setattro */
462 0, /* tp_as_buffer */
463 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
464 teedataobject_doc, /* tp_doc */
465 (traverseproc)teedataobject_traverse, /* tp_traverse */
466 (inquiry)teedataobject_clear, /* tp_clear */
467 0, /* tp_richcompare */
468 0, /* tp_weaklistoffset */
469 0, /* tp_iter */
470 0, /* tp_iternext */
471 0, /* tp_methods */
472 0, /* tp_members */
473 0, /* tp_getset */
474 0, /* tp_base */
475 0, /* tp_dict */
476 0, /* tp_descr_get */
477 0, /* tp_descr_set */
478 0, /* tp_dictoffset */
479 0, /* tp_init */
480 0, /* tp_alloc */
481 0, /* tp_new */
482 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000483};
484
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000485
486static PyTypeObject tee_type;
487
488static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000489tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000490{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000491 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000492
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000493 if (to->index >= LINKCELLS) {
494 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200495 if (link == NULL)
496 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000497 Py_DECREF(to->dataobj);
498 to->dataobj = (teedataobject *)link;
499 to->index = 0;
500 }
501 value = teedataobject_getitem(to->dataobj, to->index);
502 if (value == NULL)
503 return NULL;
504 to->index++;
505 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000506}
507
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000508static int
509tee_traverse(teeobject *to, visitproc visit, void *arg)
510{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000511 Py_VISIT((PyObject *)to->dataobj);
512 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000513}
514
Raymond Hettingerad983e72003-11-12 14:32:26 +0000515static PyObject *
516tee_copy(teeobject *to)
517{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000518 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000519
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000520 newto = PyObject_GC_New(teeobject, &tee_type);
521 if (newto == NULL)
522 return NULL;
523 Py_INCREF(to->dataobj);
524 newto->dataobj = to->dataobj;
525 newto->index = to->index;
526 newto->weakreflist = NULL;
527 PyObject_GC_Track(newto);
528 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000529}
530
531PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
532
533static PyObject *
534tee_fromiterable(PyObject *iterable)
535{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000536 teeobject *to;
537 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000538
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000539 it = PyObject_GetIter(iterable);
540 if (it == NULL)
541 return NULL;
542 if (PyObject_TypeCheck(it, &tee_type)) {
543 to = (teeobject *)tee_copy((teeobject *)it);
544 goto done;
545 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000546
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000547 to = PyObject_GC_New(teeobject, &tee_type);
548 if (to == NULL)
549 goto done;
550 to->dataobj = (teedataobject *)teedataobject_new(it);
551 if (!to->dataobj) {
552 PyObject_GC_Del(to);
553 to = NULL;
554 goto done;
555 }
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000556
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000557 to->index = 0;
558 to->weakreflist = NULL;
559 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000560done:
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000561 Py_XDECREF(it);
562 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000563}
564
565static PyObject *
566tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
567{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000568 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000569
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000570 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
571 return NULL;
572 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000573}
574
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000575static int
576tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000577{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000578 if (to->weakreflist != NULL)
579 PyObject_ClearWeakRefs((PyObject *) to);
580 Py_CLEAR(to->dataobj);
581 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000582}
583
584static void
585tee_dealloc(teeobject *to)
586{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000587 PyObject_GC_UnTrack(to);
588 tee_clear(to);
589 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000590}
591
Raymond Hettingerad983e72003-11-12 14:32:26 +0000592PyDoc_STRVAR(teeobject_doc,
593"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000594
Raymond Hettingerad983e72003-11-12 14:32:26 +0000595static PyMethodDef tee_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000596 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
597 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000598};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000599
600static PyTypeObject tee_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000601 PyVarObject_HEAD_INIT(NULL, 0)
602 "itertools.tee", /* tp_name */
603 sizeof(teeobject), /* tp_basicsize */
604 0, /* tp_itemsize */
605 /* methods */
606 (destructor)tee_dealloc, /* tp_dealloc */
607 0, /* tp_print */
608 0, /* tp_getattr */
609 0, /* tp_setattr */
610 0, /* tp_compare */
611 0, /* tp_repr */
612 0, /* tp_as_number */
613 0, /* tp_as_sequence */
614 0, /* tp_as_mapping */
615 0, /* tp_hash */
616 0, /* tp_call */
617 0, /* tp_str */
618 0, /* tp_getattro */
619 0, /* tp_setattro */
620 0, /* tp_as_buffer */
621 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
622 teeobject_doc, /* tp_doc */
623 (traverseproc)tee_traverse, /* tp_traverse */
624 (inquiry)tee_clear, /* tp_clear */
625 0, /* tp_richcompare */
626 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
627 PyObject_SelfIter, /* tp_iter */
628 (iternextfunc)tee_next, /* tp_iternext */
629 tee_methods, /* tp_methods */
630 0, /* tp_members */
631 0, /* tp_getset */
632 0, /* tp_base */
633 0, /* tp_dict */
634 0, /* tp_descr_get */
635 0, /* tp_descr_set */
636 0, /* tp_dictoffset */
637 0, /* tp_init */
638 0, /* tp_alloc */
639 tee_new, /* tp_new */
640 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000641};
642
Raymond Hettingerad983e72003-11-12 14:32:26 +0000643static PyObject *
644tee(PyObject *self, PyObject *args)
645{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000646 Py_ssize_t i, n=2;
647 PyObject *it, *iterable, *copyable, *result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000648
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000649 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
650 return NULL;
651 if (n < 0) {
652 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
653 return NULL;
654 }
655 result = PyTuple_New(n);
656 if (result == NULL)
657 return NULL;
658 if (n == 0)
659 return result;
660 it = PyObject_GetIter(iterable);
661 if (it == NULL) {
662 Py_DECREF(result);
663 return NULL;
664 }
665 if (!PyObject_HasAttrString(it, "__copy__")) {
666 copyable = tee_fromiterable(it);
667 Py_DECREF(it);
668 if (copyable == NULL) {
669 Py_DECREF(result);
670 return NULL;
671 }
672 } else
673 copyable = it;
674 PyTuple_SET_ITEM(result, 0, copyable);
675 for (i=1 ; i<n ; i++) {
676 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
677 if (copyable == NULL) {
678 Py_DECREF(result);
679 return NULL;
680 }
681 PyTuple_SET_ITEM(result, i, copyable);
682 }
683 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000684}
685
686PyDoc_STRVAR(tee_doc,
687"tee(iterable, n=2) --> tuple of n independent iterators.");
688
689
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000690/* cycle object **********************************************************/
691
692typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000693 PyObject_HEAD
694 PyObject *it;
695 PyObject *saved;
696 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000697} cycleobject;
698
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000699static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000700
701static PyObject *
702cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
703{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000704 PyObject *it;
705 PyObject *iterable;
706 PyObject *saved;
707 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000708
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000709 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
710 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000711
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000712 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
713 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000714
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000715 /* Get iterator. */
716 it = PyObject_GetIter(iterable);
717 if (it == NULL)
718 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000719
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000720 saved = PyList_New(0);
721 if (saved == NULL) {
722 Py_DECREF(it);
723 return NULL;
724 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000725
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000726 /* create cycleobject structure */
727 lz = (cycleobject *)type->tp_alloc(type, 0);
728 if (lz == NULL) {
729 Py_DECREF(it);
730 Py_DECREF(saved);
731 return NULL;
732 }
733 lz->it = it;
734 lz->saved = saved;
735 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000736
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000737 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000738}
739
740static void
741cycle_dealloc(cycleobject *lz)
742{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000743 PyObject_GC_UnTrack(lz);
744 Py_XDECREF(lz->saved);
745 Py_XDECREF(lz->it);
746 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000747}
748
749static int
750cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
751{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000752 Py_VISIT(lz->it);
753 Py_VISIT(lz->saved);
754 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000755}
756
757static PyObject *
758cycle_next(cycleobject *lz)
759{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000760 PyObject *item;
761 PyObject *it;
762 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000763
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000764 while (1) {
765 item = PyIter_Next(lz->it);
766 if (item != NULL) {
767 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
768 Py_DECREF(item);
769 return NULL;
770 }
771 return item;
772 }
773 if (PyErr_Occurred()) {
774 if (PyErr_ExceptionMatches(PyExc_StopIteration))
775 PyErr_Clear();
776 else
777 return NULL;
778 }
779 if (PyList_Size(lz->saved) == 0)
780 return NULL;
781 it = PyObject_GetIter(lz->saved);
782 if (it == NULL)
783 return NULL;
784 tmp = lz->it;
785 lz->it = it;
786 lz->firstpass = 1;
787 Py_DECREF(tmp);
788 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000789}
790
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000791PyDoc_STRVAR(cycle_doc,
792"cycle(iterable) --> cycle object\n\
793\n\
794Return elements from the iterable until it is exhausted.\n\
795Then repeat the sequence indefinitely.");
796
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000797static PyTypeObject cycle_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000798 PyVarObject_HEAD_INIT(NULL, 0)
799 "itertools.cycle", /* tp_name */
800 sizeof(cycleobject), /* tp_basicsize */
801 0, /* tp_itemsize */
802 /* methods */
803 (destructor)cycle_dealloc, /* tp_dealloc */
804 0, /* tp_print */
805 0, /* tp_getattr */
806 0, /* tp_setattr */
807 0, /* tp_compare */
808 0, /* tp_repr */
809 0, /* tp_as_number */
810 0, /* tp_as_sequence */
811 0, /* tp_as_mapping */
812 0, /* tp_hash */
813 0, /* tp_call */
814 0, /* tp_str */
815 PyObject_GenericGetAttr, /* tp_getattro */
816 0, /* tp_setattro */
817 0, /* tp_as_buffer */
818 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
819 Py_TPFLAGS_BASETYPE, /* tp_flags */
820 cycle_doc, /* tp_doc */
821 (traverseproc)cycle_traverse, /* tp_traverse */
822 0, /* tp_clear */
823 0, /* tp_richcompare */
824 0, /* tp_weaklistoffset */
825 PyObject_SelfIter, /* tp_iter */
826 (iternextfunc)cycle_next, /* tp_iternext */
827 0, /* tp_methods */
828 0, /* tp_members */
829 0, /* tp_getset */
830 0, /* tp_base */
831 0, /* tp_dict */
832 0, /* tp_descr_get */
833 0, /* tp_descr_set */
834 0, /* tp_dictoffset */
835 0, /* tp_init */
836 0, /* tp_alloc */
837 cycle_new, /* tp_new */
838 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000839};
840
841
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000842/* dropwhile object **********************************************************/
843
844typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000845 PyObject_HEAD
846 PyObject *func;
847 PyObject *it;
848 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000849} dropwhileobject;
850
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000851static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000852
853static PyObject *
854dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
855{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000856 PyObject *func, *seq;
857 PyObject *it;
858 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000859
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000860 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
861 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000862
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000863 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
864 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000865
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000866 /* Get iterator. */
867 it = PyObject_GetIter(seq);
868 if (it == NULL)
869 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000870
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000871 /* create dropwhileobject structure */
872 lz = (dropwhileobject *)type->tp_alloc(type, 0);
873 if (lz == NULL) {
874 Py_DECREF(it);
875 return NULL;
876 }
877 Py_INCREF(func);
878 lz->func = func;
879 lz->it = it;
880 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000881
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000882 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000883}
884
885static void
886dropwhile_dealloc(dropwhileobject *lz)
887{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000888 PyObject_GC_UnTrack(lz);
889 Py_XDECREF(lz->func);
890 Py_XDECREF(lz->it);
891 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000892}
893
894static int
895dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
896{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000897 Py_VISIT(lz->it);
898 Py_VISIT(lz->func);
899 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000900}
901
902static PyObject *
903dropwhile_next(dropwhileobject *lz)
904{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000905 PyObject *item, *good;
906 PyObject *it = lz->it;
907 long ok;
908 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000909
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000910 iternext = *Py_TYPE(it)->tp_iternext;
911 for (;;) {
912 item = iternext(it);
913 if (item == NULL)
914 return NULL;
915 if (lz->start == 1)
916 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000917
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000918 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
919 if (good == NULL) {
920 Py_DECREF(item);
921 return NULL;
922 }
923 ok = PyObject_IsTrue(good);
924 Py_DECREF(good);
Antoine Pitrouc5bef752012-08-15 23:16:51 +0200925 if (ok == 0) {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000926 lz->start = 1;
927 return item;
928 }
929 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +0200930 if (ok < 0)
931 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000932 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000933}
934
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000935PyDoc_STRVAR(dropwhile_doc,
936"dropwhile(predicate, iterable) --> dropwhile object\n\
937\n\
938Drop items from the iterable while predicate(item) is true.\n\
939Afterwards, return every element until the iterable is exhausted.");
940
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000941static PyTypeObject dropwhile_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000942 PyVarObject_HEAD_INIT(NULL, 0)
943 "itertools.dropwhile", /* tp_name */
944 sizeof(dropwhileobject), /* tp_basicsize */
945 0, /* tp_itemsize */
946 /* methods */
947 (destructor)dropwhile_dealloc, /* tp_dealloc */
948 0, /* tp_print */
949 0, /* tp_getattr */
950 0, /* tp_setattr */
951 0, /* tp_compare */
952 0, /* tp_repr */
953 0, /* tp_as_number */
954 0, /* tp_as_sequence */
955 0, /* tp_as_mapping */
956 0, /* tp_hash */
957 0, /* tp_call */
958 0, /* tp_str */
959 PyObject_GenericGetAttr, /* tp_getattro */
960 0, /* tp_setattro */
961 0, /* tp_as_buffer */
962 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
963 Py_TPFLAGS_BASETYPE, /* tp_flags */
964 dropwhile_doc, /* tp_doc */
965 (traverseproc)dropwhile_traverse, /* tp_traverse */
966 0, /* tp_clear */
967 0, /* tp_richcompare */
968 0, /* tp_weaklistoffset */
969 PyObject_SelfIter, /* tp_iter */
970 (iternextfunc)dropwhile_next, /* tp_iternext */
971 0, /* tp_methods */
972 0, /* tp_members */
973 0, /* tp_getset */
974 0, /* tp_base */
975 0, /* tp_dict */
976 0, /* tp_descr_get */
977 0, /* tp_descr_set */
978 0, /* tp_dictoffset */
979 0, /* tp_init */
980 0, /* tp_alloc */
981 dropwhile_new, /* tp_new */
982 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000983};
984
985
986/* takewhile object **********************************************************/
987
988typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000989 PyObject_HEAD
990 PyObject *func;
991 PyObject *it;
992 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000993} takewhileobject;
994
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000995static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000996
997static PyObject *
998takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
999{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001000 PyObject *func, *seq;
1001 PyObject *it;
1002 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001003
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001004 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1005 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001006
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001007 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1008 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001009
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001010 /* Get iterator. */
1011 it = PyObject_GetIter(seq);
1012 if (it == NULL)
1013 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001014
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001015 /* create takewhileobject structure */
1016 lz = (takewhileobject *)type->tp_alloc(type, 0);
1017 if (lz == NULL) {
1018 Py_DECREF(it);
1019 return NULL;
1020 }
1021 Py_INCREF(func);
1022 lz->func = func;
1023 lz->it = it;
1024 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001025
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001026 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001027}
1028
1029static void
1030takewhile_dealloc(takewhileobject *lz)
1031{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001032 PyObject_GC_UnTrack(lz);
1033 Py_XDECREF(lz->func);
1034 Py_XDECREF(lz->it);
1035 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001036}
1037
1038static int
1039takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1040{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001041 Py_VISIT(lz->it);
1042 Py_VISIT(lz->func);
1043 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001044}
1045
1046static PyObject *
1047takewhile_next(takewhileobject *lz)
1048{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001049 PyObject *item, *good;
1050 PyObject *it = lz->it;
1051 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001052
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001053 if (lz->stop == 1)
1054 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001055
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001056 item = (*Py_TYPE(it)->tp_iternext)(it);
1057 if (item == NULL)
1058 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001059
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001060 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1061 if (good == NULL) {
1062 Py_DECREF(item);
1063 return NULL;
1064 }
1065 ok = PyObject_IsTrue(good);
1066 Py_DECREF(good);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02001067 if (ok > 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001068 return item;
1069 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02001070 if (ok == 0)
1071 lz->stop = 1;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001072 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001073}
1074
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001075PyDoc_STRVAR(takewhile_doc,
1076"takewhile(predicate, iterable) --> takewhile object\n\
1077\n\
1078Return successive entries from an iterable as long as the \n\
1079predicate evaluates to true for each entry.");
1080
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001081static PyTypeObject takewhile_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001082 PyVarObject_HEAD_INIT(NULL, 0)
1083 "itertools.takewhile", /* tp_name */
1084 sizeof(takewhileobject), /* tp_basicsize */
1085 0, /* tp_itemsize */
1086 /* methods */
1087 (destructor)takewhile_dealloc, /* tp_dealloc */
1088 0, /* tp_print */
1089 0, /* tp_getattr */
1090 0, /* tp_setattr */
1091 0, /* tp_compare */
1092 0, /* tp_repr */
1093 0, /* tp_as_number */
1094 0, /* tp_as_sequence */
1095 0, /* tp_as_mapping */
1096 0, /* tp_hash */
1097 0, /* tp_call */
1098 0, /* tp_str */
1099 PyObject_GenericGetAttr, /* tp_getattro */
1100 0, /* tp_setattro */
1101 0, /* tp_as_buffer */
1102 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1103 Py_TPFLAGS_BASETYPE, /* tp_flags */
1104 takewhile_doc, /* tp_doc */
1105 (traverseproc)takewhile_traverse, /* tp_traverse */
1106 0, /* tp_clear */
1107 0, /* tp_richcompare */
1108 0, /* tp_weaklistoffset */
1109 PyObject_SelfIter, /* tp_iter */
1110 (iternextfunc)takewhile_next, /* tp_iternext */
1111 0, /* tp_methods */
1112 0, /* tp_members */
1113 0, /* tp_getset */
1114 0, /* tp_base */
1115 0, /* tp_dict */
1116 0, /* tp_descr_get */
1117 0, /* tp_descr_set */
1118 0, /* tp_dictoffset */
1119 0, /* tp_init */
1120 0, /* tp_alloc */
1121 takewhile_new, /* tp_new */
1122 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001123};
1124
1125
1126/* islice object ************************************************************/
1127
1128typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001129 PyObject_HEAD
1130 PyObject *it;
1131 Py_ssize_t next;
1132 Py_ssize_t stop;
1133 Py_ssize_t step;
1134 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001135} isliceobject;
1136
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001137static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001138
1139static PyObject *
1140islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1141{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001142 PyObject *seq;
1143 Py_ssize_t start=0, stop=-1, step=1;
1144 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1145 Py_ssize_t numargs;
1146 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001147
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001148 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1149 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001150
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001151 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1152 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001153
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001154 numargs = PyTuple_Size(args);
1155 if (numargs == 2) {
1156 if (a1 != Py_None) {
1157 stop = PyInt_AsSsize_t(a1);
1158 if (stop == -1) {
1159 if (PyErr_Occurred())
1160 PyErr_Clear();
1161 PyErr_SetString(PyExc_ValueError,
1162 "Stop argument for islice() must be None or an integer: 0 <= x <= maxint.");
1163 return NULL;
1164 }
1165 }
1166 } else {
1167 if (a1 != Py_None)
1168 start = PyInt_AsSsize_t(a1);
1169 if (start == -1 && PyErr_Occurred())
1170 PyErr_Clear();
1171 if (a2 != Py_None) {
1172 stop = PyInt_AsSsize_t(a2);
1173 if (stop == -1) {
1174 if (PyErr_Occurred())
1175 PyErr_Clear();
1176 PyErr_SetString(PyExc_ValueError,
1177 "Stop argument for islice() must be None or an integer: 0 <= x <= maxint.");
1178 return NULL;
1179 }
1180 }
1181 }
1182 if (start<0 || stop<-1) {
1183 PyErr_SetString(PyExc_ValueError,
1184 "Indices for islice() must be None or an integer: 0 <= x <= maxint.");
1185 return NULL;
1186 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001187
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001188 if (a3 != NULL) {
1189 if (a3 != Py_None)
1190 step = PyInt_AsSsize_t(a3);
1191 if (step == -1 && PyErr_Occurred())
1192 PyErr_Clear();
1193 }
1194 if (step<1) {
1195 PyErr_SetString(PyExc_ValueError,
1196 "Step for islice() must be a positive integer or None.");
1197 return NULL;
1198 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001199
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001200 /* Get iterator. */
1201 it = PyObject_GetIter(seq);
1202 if (it == NULL)
1203 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001205 /* create isliceobject structure */
1206 lz = (isliceobject *)type->tp_alloc(type, 0);
1207 if (lz == NULL) {
1208 Py_DECREF(it);
1209 return NULL;
1210 }
1211 lz->it = it;
1212 lz->next = start;
1213 lz->stop = stop;
1214 lz->step = step;
1215 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001216
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001217 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001218}
1219
1220static void
1221islice_dealloc(isliceobject *lz)
1222{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001223 PyObject_GC_UnTrack(lz);
1224 Py_XDECREF(lz->it);
1225 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001226}
1227
1228static int
1229islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1230{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001231 Py_VISIT(lz->it);
1232 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001233}
1234
1235static PyObject *
1236islice_next(isliceobject *lz)
1237{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001238 PyObject *item;
1239 PyObject *it = lz->it;
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001240 Py_ssize_t stop = lz->stop;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001241 Py_ssize_t oldnext;
1242 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001243
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001244 iternext = *Py_TYPE(it)->tp_iternext;
1245 while (lz->cnt < lz->next) {
1246 item = iternext(it);
1247 if (item == NULL)
1248 return NULL;
1249 Py_DECREF(item);
1250 lz->cnt++;
1251 }
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001252 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001253 return NULL;
1254 item = iternext(it);
1255 if (item == NULL)
1256 return NULL;
1257 lz->cnt++;
1258 oldnext = lz->next;
Mark Dickinsona96b0d12011-09-24 09:01:16 +01001259 /* The (size_t) cast below avoids the danger of undefined
1260 behaviour from signed integer overflow. */
1261 lz->next += (size_t)lz->step;
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001262 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1263 lz->next = stop;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001264 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001265}
1266
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001267PyDoc_STRVAR(islice_doc,
1268"islice(iterable, [start,] stop [, step]) --> islice object\n\
1269\n\
1270Return an iterator whose next() method returns selected values from an\n\
1271iterable. If start is specified, will skip all preceding elements;\n\
1272otherwise, start defaults to zero. Step defaults to one. If\n\
1273specified as another value, step determines how many values are \n\
1274skipped between successive calls. Works like a slice() on a list\n\
1275but returns an iterator.");
1276
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001277static PyTypeObject islice_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001278 PyVarObject_HEAD_INIT(NULL, 0)
1279 "itertools.islice", /* tp_name */
1280 sizeof(isliceobject), /* tp_basicsize */
1281 0, /* tp_itemsize */
1282 /* methods */
1283 (destructor)islice_dealloc, /* tp_dealloc */
1284 0, /* tp_print */
1285 0, /* tp_getattr */
1286 0, /* tp_setattr */
1287 0, /* tp_compare */
1288 0, /* tp_repr */
1289 0, /* tp_as_number */
1290 0, /* tp_as_sequence */
1291 0, /* tp_as_mapping */
1292 0, /* tp_hash */
1293 0, /* tp_call */
1294 0, /* tp_str */
1295 PyObject_GenericGetAttr, /* tp_getattro */
1296 0, /* tp_setattro */
1297 0, /* tp_as_buffer */
1298 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1299 Py_TPFLAGS_BASETYPE, /* tp_flags */
1300 islice_doc, /* tp_doc */
1301 (traverseproc)islice_traverse, /* tp_traverse */
1302 0, /* tp_clear */
1303 0, /* tp_richcompare */
1304 0, /* tp_weaklistoffset */
1305 PyObject_SelfIter, /* tp_iter */
1306 (iternextfunc)islice_next, /* tp_iternext */
1307 0, /* tp_methods */
1308 0, /* tp_members */
1309 0, /* tp_getset */
1310 0, /* tp_base */
1311 0, /* tp_dict */
1312 0, /* tp_descr_get */
1313 0, /* tp_descr_set */
1314 0, /* tp_dictoffset */
1315 0, /* tp_init */
1316 0, /* tp_alloc */
1317 islice_new, /* tp_new */
1318 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001319};
1320
1321
1322/* starmap object ************************************************************/
1323
1324typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001325 PyObject_HEAD
1326 PyObject *func;
1327 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001328} starmapobject;
1329
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001330static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001331
1332static PyObject *
1333starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1334{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001335 PyObject *func, *seq;
1336 PyObject *it;
1337 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001338
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001339 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1340 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001341
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001342 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1343 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001344
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001345 /* Get iterator. */
1346 it = PyObject_GetIter(seq);
1347 if (it == NULL)
1348 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001349
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001350 /* create starmapobject structure */
1351 lz = (starmapobject *)type->tp_alloc(type, 0);
1352 if (lz == NULL) {
1353 Py_DECREF(it);
1354 return NULL;
1355 }
1356 Py_INCREF(func);
1357 lz->func = func;
1358 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001359
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001360 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001361}
1362
1363static void
1364starmap_dealloc(starmapobject *lz)
1365{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001366 PyObject_GC_UnTrack(lz);
1367 Py_XDECREF(lz->func);
1368 Py_XDECREF(lz->it);
1369 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001370}
1371
1372static int
1373starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1374{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001375 Py_VISIT(lz->it);
1376 Py_VISIT(lz->func);
1377 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001378}
1379
1380static PyObject *
1381starmap_next(starmapobject *lz)
1382{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001383 PyObject *args;
1384 PyObject *result;
1385 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001386
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001387 args = (*Py_TYPE(it)->tp_iternext)(it);
1388 if (args == NULL)
1389 return NULL;
1390 if (!PyTuple_CheckExact(args)) {
1391 PyObject *newargs = PySequence_Tuple(args);
1392 Py_DECREF(args);
1393 if (newargs == NULL)
1394 return NULL;
1395 args = newargs;
1396 }
1397 result = PyObject_Call(lz->func, args, NULL);
1398 Py_DECREF(args);
1399 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001400}
1401
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001402PyDoc_STRVAR(starmap_doc,
1403"starmap(function, sequence) --> starmap object\n\
1404\n\
1405Return an iterator whose values are returned from the function evaluated\n\
1406with a argument tuple taken from the given sequence.");
1407
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001408static PyTypeObject starmap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001409 PyVarObject_HEAD_INIT(NULL, 0)
1410 "itertools.starmap", /* tp_name */
1411 sizeof(starmapobject), /* tp_basicsize */
1412 0, /* tp_itemsize */
1413 /* methods */
1414 (destructor)starmap_dealloc, /* tp_dealloc */
1415 0, /* tp_print */
1416 0, /* tp_getattr */
1417 0, /* tp_setattr */
1418 0, /* tp_compare */
1419 0, /* tp_repr */
1420 0, /* tp_as_number */
1421 0, /* tp_as_sequence */
1422 0, /* tp_as_mapping */
1423 0, /* tp_hash */
1424 0, /* tp_call */
1425 0, /* tp_str */
1426 PyObject_GenericGetAttr, /* tp_getattro */
1427 0, /* tp_setattro */
1428 0, /* tp_as_buffer */
1429 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1430 Py_TPFLAGS_BASETYPE, /* tp_flags */
1431 starmap_doc, /* tp_doc */
1432 (traverseproc)starmap_traverse, /* tp_traverse */
1433 0, /* tp_clear */
1434 0, /* tp_richcompare */
1435 0, /* tp_weaklistoffset */
1436 PyObject_SelfIter, /* tp_iter */
1437 (iternextfunc)starmap_next, /* tp_iternext */
1438 0, /* tp_methods */
1439 0, /* tp_members */
1440 0, /* tp_getset */
1441 0, /* tp_base */
1442 0, /* tp_dict */
1443 0, /* tp_descr_get */
1444 0, /* tp_descr_set */
1445 0, /* tp_dictoffset */
1446 0, /* tp_init */
1447 0, /* tp_alloc */
1448 starmap_new, /* tp_new */
1449 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001450};
1451
1452
1453/* imap object ************************************************************/
1454
1455typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001456 PyObject_HEAD
1457 PyObject *iters;
1458 PyObject *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459} imapobject;
1460
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001461static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001462
1463static PyObject *
1464imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1465{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001466 PyObject *it, *iters, *func;
1467 imapobject *lz;
1468 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001470 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
1471 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001472
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001473 numargs = PyTuple_Size(args);
1474 if (numargs < 2) {
1475 PyErr_SetString(PyExc_TypeError,
1476 "imap() must have at least two arguments.");
1477 return NULL;
1478 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001479
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001480 iters = PyTuple_New(numargs-1);
1481 if (iters == NULL)
1482 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001483
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001484 for (i=1 ; i<numargs ; i++) {
1485 /* Get iterator. */
1486 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1487 if (it == NULL) {
1488 Py_DECREF(iters);
1489 return NULL;
1490 }
1491 PyTuple_SET_ITEM(iters, i-1, it);
1492 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001493
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001494 /* create imapobject structure */
1495 lz = (imapobject *)type->tp_alloc(type, 0);
1496 if (lz == NULL) {
1497 Py_DECREF(iters);
1498 return NULL;
1499 }
1500 lz->iters = iters;
1501 func = PyTuple_GET_ITEM(args, 0);
1502 Py_INCREF(func);
1503 lz->func = func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001504
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001505 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001506}
1507
1508static void
1509imap_dealloc(imapobject *lz)
1510{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001511 PyObject_GC_UnTrack(lz);
1512 Py_XDECREF(lz->iters);
1513 Py_XDECREF(lz->func);
1514 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001515}
1516
1517static int
1518imap_traverse(imapobject *lz, visitproc visit, void *arg)
1519{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001520 Py_VISIT(lz->iters);
1521 Py_VISIT(lz->func);
1522 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001523}
1524
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001525/*
Raymond Hettinger2012f172003-02-07 05:32:58 +00001526imap() is an iterator version of __builtins__.map() except that it does
1527not have the None fill-in feature. That was intentionally left out for
1528the following reasons:
1529
1530 1) Itertools are designed to be easily combined and chained together.
1531 Having all tools stop with the shortest input is a unifying principle
1532 that makes it easier to combine finite iterators (supplying data) with
1533 infinite iterators like count() and repeat() (for supplying sequential
1534 or constant arguments to a function).
1535
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001536 2) In typical use cases for combining itertools, having one finite data
Raymond Hettinger2012f172003-02-07 05:32:58 +00001537 supplier run out before another is likely to be an error condition which
1538 should not pass silently by automatically supplying None.
1539
1540 3) The use cases for automatic None fill-in are rare -- not many functions
1541 do something useful when a parameter suddenly switches type and becomes
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001542 None.
Raymond Hettinger2012f172003-02-07 05:32:58 +00001543
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001544 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001545 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001546
1547 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1548*/
1549
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001550static PyObject *
1551imap_next(imapobject *lz)
1552{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001553 PyObject *val;
1554 PyObject *argtuple;
1555 PyObject *result;
1556 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001557
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001558 numargs = PyTuple_Size(lz->iters);
1559 argtuple = PyTuple_New(numargs);
1560 if (argtuple == NULL)
1561 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001562
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001563 for (i=0 ; i<numargs ; i++) {
1564 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1565 if (val == NULL) {
1566 Py_DECREF(argtuple);
1567 return NULL;
1568 }
1569 PyTuple_SET_ITEM(argtuple, i, val);
1570 }
1571 if (lz->func == Py_None)
1572 return argtuple;
1573 result = PyObject_Call(lz->func, argtuple, NULL);
1574 Py_DECREF(argtuple);
1575 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001576}
1577
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001578PyDoc_STRVAR(imap_doc,
1579"imap(func, *iterables) --> imap object\n\
1580\n\
1581Make an iterator that computes the function using arguments from\n\
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001582each of the iterables. Like map() except that it returns\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001583an iterator instead of a list and that it stops when the shortest\n\
1584iterable is exhausted instead of filling in None for shorter\n\
1585iterables.");
1586
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001587static PyTypeObject imap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001588 PyVarObject_HEAD_INIT(NULL, 0)
1589 "itertools.imap", /* tp_name */
1590 sizeof(imapobject), /* tp_basicsize */
1591 0, /* tp_itemsize */
1592 /* methods */
1593 (destructor)imap_dealloc, /* tp_dealloc */
1594 0, /* tp_print */
1595 0, /* tp_getattr */
1596 0, /* tp_setattr */
1597 0, /* tp_compare */
1598 0, /* tp_repr */
1599 0, /* tp_as_number */
1600 0, /* tp_as_sequence */
1601 0, /* tp_as_mapping */
1602 0, /* tp_hash */
1603 0, /* tp_call */
1604 0, /* tp_str */
1605 PyObject_GenericGetAttr, /* tp_getattro */
1606 0, /* tp_setattro */
1607 0, /* tp_as_buffer */
1608 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1609 Py_TPFLAGS_BASETYPE, /* tp_flags */
1610 imap_doc, /* tp_doc */
1611 (traverseproc)imap_traverse, /* tp_traverse */
1612 0, /* tp_clear */
1613 0, /* tp_richcompare */
1614 0, /* tp_weaklistoffset */
1615 PyObject_SelfIter, /* tp_iter */
1616 (iternextfunc)imap_next, /* tp_iternext */
1617 0, /* tp_methods */
1618 0, /* tp_members */
1619 0, /* tp_getset */
1620 0, /* tp_base */
1621 0, /* tp_dict */
1622 0, /* tp_descr_get */
1623 0, /* tp_descr_set */
1624 0, /* tp_dictoffset */
1625 0, /* tp_init */
1626 0, /* tp_alloc */
1627 imap_new, /* tp_new */
1628 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001629};
1630
1631
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001632/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001633
1634typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001635 PyObject_HEAD
1636 PyObject *source; /* Iterator over input iterables */
1637 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001638} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001639
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001640static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001641
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001642static PyObject *
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001643chain_new_internal(PyTypeObject *type, PyObject *source)
1644{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001645 chainobject *lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001646
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001647 lz = (chainobject *)type->tp_alloc(type, 0);
1648 if (lz == NULL) {
1649 Py_DECREF(source);
1650 return NULL;
1651 }
1652
1653 lz->source = source;
1654 lz->active = NULL;
1655 return (PyObject *)lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001656}
1657
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001658static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001659chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001660{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001661 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001662
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001663 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1664 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001665
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001666 source = PyObject_GetIter(args);
1667 if (source == NULL)
1668 return NULL;
1669
1670 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001671}
1672
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001673static PyObject *
1674chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1675{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001676 PyObject *source;
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001677
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001678 source = PyObject_GetIter(arg);
1679 if (source == NULL)
1680 return NULL;
1681
1682 return chain_new_internal(type, source);
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001683}
1684
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001685static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001686chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001687{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001688 PyObject_GC_UnTrack(lz);
1689 Py_XDECREF(lz->active);
1690 Py_XDECREF(lz->source);
1691 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001692}
1693
Raymond Hettinger2012f172003-02-07 05:32:58 +00001694static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001695chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001696{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001697 Py_VISIT(lz->source);
1698 Py_VISIT(lz->active);
1699 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001700}
1701
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001702static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001703chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001704{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001705 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001706
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001707 if (lz->source == NULL)
1708 return NULL; /* already stopped */
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001709
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001710 if (lz->active == NULL) {
1711 PyObject *iterable = PyIter_Next(lz->source);
1712 if (iterable == NULL) {
1713 Py_CLEAR(lz->source);
1714 return NULL; /* no more input sources */
1715 }
1716 lz->active = PyObject_GetIter(iterable);
1717 Py_DECREF(iterable);
1718 if (lz->active == NULL) {
1719 Py_CLEAR(lz->source);
1720 return NULL; /* input not iterable */
1721 }
1722 }
1723 item = PyIter_Next(lz->active);
1724 if (item != NULL)
1725 return item;
1726 if (PyErr_Occurred()) {
1727 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1728 PyErr_Clear();
1729 else
1730 return NULL; /* input raised an exception */
1731 }
1732 Py_CLEAR(lz->active);
1733 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001734}
1735
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001736PyDoc_STRVAR(chain_doc,
1737"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001738\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001739Return a chain object whose .next() method returns elements from the\n\
1740first iterable until it is exhausted, then elements from the next\n\
1741iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001742
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001743PyDoc_STRVAR(chain_from_iterable_doc,
1744"chain.from_iterable(iterable) --> chain object\n\
1745\n\
1746Alternate chain() contructor taking a single iterable argument\n\
1747that evaluates lazily.");
1748
1749static PyMethodDef chain_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001750 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1751 chain_from_iterable_doc},
1752 {NULL, NULL} /* sentinel */
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001753};
1754
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001755static PyTypeObject chain_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001756 PyVarObject_HEAD_INIT(NULL, 0)
1757 "itertools.chain", /* tp_name */
1758 sizeof(chainobject), /* tp_basicsize */
1759 0, /* tp_itemsize */
1760 /* methods */
1761 (destructor)chain_dealloc, /* tp_dealloc */
1762 0, /* tp_print */
1763 0, /* tp_getattr */
1764 0, /* tp_setattr */
1765 0, /* tp_compare */
1766 0, /* tp_repr */
1767 0, /* tp_as_number */
1768 0, /* tp_as_sequence */
1769 0, /* tp_as_mapping */
1770 0, /* tp_hash */
1771 0, /* tp_call */
1772 0, /* tp_str */
1773 PyObject_GenericGetAttr, /* tp_getattro */
1774 0, /* tp_setattro */
1775 0, /* tp_as_buffer */
1776 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1777 Py_TPFLAGS_BASETYPE, /* tp_flags */
1778 chain_doc, /* tp_doc */
1779 (traverseproc)chain_traverse, /* tp_traverse */
1780 0, /* tp_clear */
1781 0, /* tp_richcompare */
1782 0, /* tp_weaklistoffset */
1783 PyObject_SelfIter, /* tp_iter */
1784 (iternextfunc)chain_next, /* tp_iternext */
1785 chain_methods, /* tp_methods */
1786 0, /* tp_members */
1787 0, /* tp_getset */
1788 0, /* tp_base */
1789 0, /* tp_dict */
1790 0, /* tp_descr_get */
1791 0, /* tp_descr_set */
1792 0, /* tp_dictoffset */
1793 0, /* tp_init */
1794 0, /* tp_alloc */
1795 chain_new, /* tp_new */
1796 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001797};
1798
1799
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001800/* product object ************************************************************/
1801
1802typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001803 PyObject_HEAD
1804 PyObject *pools; /* tuple of pool tuples */
1805 Py_ssize_t *indices; /* one index per pool */
1806 PyObject *result; /* most recently returned result tuple */
1807 int stopped; /* set to 1 when the product iterator is exhausted */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001808} productobject;
1809
1810static PyTypeObject product_type;
1811
1812static PyObject *
1813product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1814{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001815 productobject *lz;
1816 Py_ssize_t nargs, npools, repeat=1;
1817 PyObject *pools = NULL;
1818 Py_ssize_t *indices = NULL;
1819 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001820
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001821 if (kwds != NULL) {
1822 char *kwlist[] = {"repeat", 0};
1823 PyObject *tmpargs = PyTuple_New(0);
1824 if (tmpargs == NULL)
1825 return NULL;
1826 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1827 Py_DECREF(tmpargs);
1828 return NULL;
1829 }
1830 Py_DECREF(tmpargs);
1831 if (repeat < 0) {
1832 PyErr_SetString(PyExc_ValueError,
1833 "repeat argument cannot be negative");
1834 return NULL;
1835 }
1836 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001837
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001838 assert(PyTuple_Check(args));
1839 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1840 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001841
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001842 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1843 if (indices == NULL) {
1844 PyErr_NoMemory();
1845 goto error;
1846 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001847
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001848 pools = PyTuple_New(npools);
1849 if (pools == NULL)
1850 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001851
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001852 for (i=0; i < nargs ; ++i) {
1853 PyObject *item = PyTuple_GET_ITEM(args, i);
1854 PyObject *pool = PySequence_Tuple(item);
1855 if (pool == NULL)
1856 goto error;
1857 PyTuple_SET_ITEM(pools, i, pool);
1858 indices[i] = 0;
1859 }
1860 for ( ; i < npools; ++i) {
1861 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1862 Py_INCREF(pool);
1863 PyTuple_SET_ITEM(pools, i, pool);
1864 indices[i] = 0;
1865 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001866
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001867 /* create productobject structure */
1868 lz = (productobject *)type->tp_alloc(type, 0);
1869 if (lz == NULL)
1870 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001871
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001872 lz->pools = pools;
1873 lz->indices = indices;
1874 lz->result = NULL;
1875 lz->stopped = 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001876
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001877 return (PyObject *)lz;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001878
1879error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001880 if (indices != NULL)
1881 PyMem_Free(indices);
1882 Py_XDECREF(pools);
1883 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001884}
1885
1886static void
1887product_dealloc(productobject *lz)
1888{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001889 PyObject_GC_UnTrack(lz);
1890 Py_XDECREF(lz->pools);
1891 Py_XDECREF(lz->result);
1892 if (lz->indices != NULL)
1893 PyMem_Free(lz->indices);
1894 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001895}
1896
1897static int
1898product_traverse(productobject *lz, visitproc visit, void *arg)
1899{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001900 Py_VISIT(lz->pools);
1901 Py_VISIT(lz->result);
1902 return 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001903}
1904
1905static PyObject *
1906product_next(productobject *lz)
1907{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001908 PyObject *pool;
1909 PyObject *elem;
1910 PyObject *oldelem;
1911 PyObject *pools = lz->pools;
1912 PyObject *result = lz->result;
1913 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1914 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001915
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001916 if (lz->stopped)
1917 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001918
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001919 if (result == NULL) {
1920 /* On the first pass, return an initial tuple filled with the
1921 first element from each pool. */
1922 result = PyTuple_New(npools);
1923 if (result == NULL)
1924 goto empty;
1925 lz->result = result;
1926 for (i=0; i < npools; i++) {
1927 pool = PyTuple_GET_ITEM(pools, i);
1928 if (PyTuple_GET_SIZE(pool) == 0)
1929 goto empty;
1930 elem = PyTuple_GET_ITEM(pool, 0);
1931 Py_INCREF(elem);
1932 PyTuple_SET_ITEM(result, i, elem);
1933 }
1934 } else {
1935 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001936
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001937 /* Copy the previous result tuple or re-use it if available */
1938 if (Py_REFCNT(result) > 1) {
1939 PyObject *old_result = result;
1940 result = PyTuple_New(npools);
1941 if (result == NULL)
1942 goto empty;
1943 lz->result = result;
1944 for (i=0; i < npools; i++) {
1945 elem = PyTuple_GET_ITEM(old_result, i);
1946 Py_INCREF(elem);
1947 PyTuple_SET_ITEM(result, i, elem);
1948 }
1949 Py_DECREF(old_result);
1950 }
1951 /* Now, we've got the only copy so we can update it in-place */
1952 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001953
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001954 /* Update the pool indices right-to-left. Only advance to the
1955 next pool when the previous one rolls-over */
1956 for (i=npools-1 ; i >= 0 ; i--) {
1957 pool = PyTuple_GET_ITEM(pools, i);
1958 indices[i]++;
1959 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1960 /* Roll-over and advance to next pool */
1961 indices[i] = 0;
1962 elem = PyTuple_GET_ITEM(pool, 0);
1963 Py_INCREF(elem);
1964 oldelem = PyTuple_GET_ITEM(result, i);
1965 PyTuple_SET_ITEM(result, i, elem);
1966 Py_DECREF(oldelem);
1967 } else {
1968 /* No rollover. Just increment and stop here. */
1969 elem = PyTuple_GET_ITEM(pool, indices[i]);
1970 Py_INCREF(elem);
1971 oldelem = PyTuple_GET_ITEM(result, i);
1972 PyTuple_SET_ITEM(result, i, elem);
1973 Py_DECREF(oldelem);
1974 break;
1975 }
1976 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00001977
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001978 /* If i is negative, then the indices have all rolled-over
1979 and we're done. */
1980 if (i < 0)
1981 goto empty;
1982 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001983
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001984 Py_INCREF(result);
1985 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001986
1987empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001988 lz->stopped = 1;
1989 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001990}
1991
1992PyDoc_STRVAR(product_doc,
1993"product(*iterables) --> product object\n\
1994\n\
Raymond Hettinger73d79632008-02-23 02:20:41 +00001995Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001996For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1997The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1998cycle in a manner similar to an odometer (with the rightmost element changing\n\
1999on every iteration).\n\n\
Raymond Hettinger94034ea2009-02-09 18:39:41 +00002000To compute the product of an iterable with itself, specify the number\n\
2001of repetitions with the optional repeat keyword argument. For example,\n\
2002product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002003product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2004product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2005
2006static PyTypeObject product_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002007 PyVarObject_HEAD_INIT(NULL, 0)
2008 "itertools.product", /* tp_name */
2009 sizeof(productobject), /* tp_basicsize */
2010 0, /* tp_itemsize */
2011 /* methods */
2012 (destructor)product_dealloc, /* tp_dealloc */
2013 0, /* tp_print */
2014 0, /* tp_getattr */
2015 0, /* tp_setattr */
2016 0, /* tp_compare */
2017 0, /* tp_repr */
2018 0, /* tp_as_number */
2019 0, /* tp_as_sequence */
2020 0, /* tp_as_mapping */
2021 0, /* tp_hash */
2022 0, /* tp_call */
2023 0, /* tp_str */
2024 PyObject_GenericGetAttr, /* tp_getattro */
2025 0, /* tp_setattro */
2026 0, /* tp_as_buffer */
2027 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2028 Py_TPFLAGS_BASETYPE, /* tp_flags */
2029 product_doc, /* tp_doc */
2030 (traverseproc)product_traverse, /* tp_traverse */
2031 0, /* tp_clear */
2032 0, /* tp_richcompare */
2033 0, /* tp_weaklistoffset */
2034 PyObject_SelfIter, /* tp_iter */
2035 (iternextfunc)product_next, /* tp_iternext */
2036 0, /* tp_methods */
2037 0, /* tp_members */
2038 0, /* tp_getset */
2039 0, /* tp_base */
2040 0, /* tp_dict */
2041 0, /* tp_descr_get */
2042 0, /* tp_descr_set */
2043 0, /* tp_dictoffset */
2044 0, /* tp_init */
2045 0, /* tp_alloc */
2046 product_new, /* tp_new */
2047 PyObject_GC_Del, /* tp_free */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002048};
2049
2050
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002051/* combinations object ************************************************************/
2052
2053typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002054 PyObject_HEAD
2055 PyObject *pool; /* input converted to a tuple */
2056 Py_ssize_t *indices; /* one index per result element */
2057 PyObject *result; /* most recently returned result tuple */
2058 Py_ssize_t r; /* size of result tuple */
2059 int stopped; /* set to 1 when the combinations iterator is exhausted */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002060} combinationsobject;
2061
2062static PyTypeObject combinations_type;
2063
2064static PyObject *
2065combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2066{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002067 combinationsobject *co;
2068 Py_ssize_t n;
2069 Py_ssize_t r;
2070 PyObject *pool = NULL;
2071 PyObject *iterable = NULL;
2072 Py_ssize_t *indices = NULL;
2073 Py_ssize_t i;
2074 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002075
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002076 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2077 &iterable, &r))
2078 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002079
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002080 pool = PySequence_Tuple(iterable);
2081 if (pool == NULL)
2082 goto error;
2083 n = PyTuple_GET_SIZE(pool);
2084 if (r < 0) {
2085 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2086 goto error;
2087 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002088
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002089 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2090 if (indices == NULL) {
2091 PyErr_NoMemory();
2092 goto error;
2093 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002094
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002095 for (i=0 ; i<r ; i++)
2096 indices[i] = i;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002097
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002098 /* create combinationsobject structure */
2099 co = (combinationsobject *)type->tp_alloc(type, 0);
2100 if (co == NULL)
2101 goto error;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002102
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002103 co->pool = pool;
2104 co->indices = indices;
2105 co->result = NULL;
2106 co->r = r;
2107 co->stopped = r > n ? 1 : 0;
2108
2109 return (PyObject *)co;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002110
2111error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002112 if (indices != NULL)
2113 PyMem_Free(indices);
2114 Py_XDECREF(pool);
2115 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002116}
2117
2118static void
2119combinations_dealloc(combinationsobject *co)
2120{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002121 PyObject_GC_UnTrack(co);
2122 Py_XDECREF(co->pool);
2123 Py_XDECREF(co->result);
2124 if (co->indices != NULL)
2125 PyMem_Free(co->indices);
2126 Py_TYPE(co)->tp_free(co);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002127}
2128
2129static int
2130combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2131{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002132 Py_VISIT(co->pool);
2133 Py_VISIT(co->result);
2134 return 0;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002135}
2136
2137static PyObject *
2138combinations_next(combinationsobject *co)
2139{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002140 PyObject *elem;
2141 PyObject *oldelem;
2142 PyObject *pool = co->pool;
2143 Py_ssize_t *indices = co->indices;
2144 PyObject *result = co->result;
2145 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2146 Py_ssize_t r = co->r;
2147 Py_ssize_t i, j, index;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002148
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002149 if (co->stopped)
2150 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002151
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002152 if (result == NULL) {
2153 /* On the first pass, initialize result tuple using the indices */
2154 result = PyTuple_New(r);
2155 if (result == NULL)
2156 goto empty;
2157 co->result = result;
2158 for (i=0; i<r ; i++) {
2159 index = indices[i];
2160 elem = PyTuple_GET_ITEM(pool, index);
2161 Py_INCREF(elem);
2162 PyTuple_SET_ITEM(result, i, elem);
2163 }
2164 } else {
2165 /* Copy the previous result tuple or re-use it if available */
2166 if (Py_REFCNT(result) > 1) {
2167 PyObject *old_result = result;
2168 result = PyTuple_New(r);
2169 if (result == NULL)
2170 goto empty;
2171 co->result = result;
2172 for (i=0; i<r ; i++) {
2173 elem = PyTuple_GET_ITEM(old_result, i);
2174 Py_INCREF(elem);
2175 PyTuple_SET_ITEM(result, i, elem);
2176 }
2177 Py_DECREF(old_result);
2178 }
2179 /* Now, we've got the only copy so we can update it in-place
2180 * CPython's empty tuple is a singleton and cached in
2181 * PyTuple's freelist.
2182 */
2183 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002184
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002185 /* Scan indices right-to-left until finding one that is not
2186 at its maximum (i + n - r). */
2187 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2188 ;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002189
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002190 /* If i is negative, then the indices are all at
2191 their maximum value and we're done. */
2192 if (i < 0)
2193 goto empty;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002194
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002195 /* Increment the current index which we know is not at its
2196 maximum. Then move back to the right setting each index
2197 to its lowest possible value (one higher than the index
2198 to its left -- this maintains the sort order invariant). */
2199 indices[i]++;
2200 for (j=i+1 ; j<r ; j++)
2201 indices[j] = indices[j-1] + 1;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002202
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002203 /* Update the result tuple for the new indices
2204 starting with i, the leftmost index that changed */
2205 for ( ; i<r ; i++) {
2206 index = indices[i];
2207 elem = PyTuple_GET_ITEM(pool, index);
2208 Py_INCREF(elem);
2209 oldelem = PyTuple_GET_ITEM(result, i);
2210 PyTuple_SET_ITEM(result, i, elem);
2211 Py_DECREF(oldelem);
2212 }
2213 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002214
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002215 Py_INCREF(result);
2216 return result;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002217
2218empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002219 co->stopped = 1;
2220 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002221}
2222
2223PyDoc_STRVAR(combinations_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002224"combinations(iterable, r) --> combinations object\n\
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002225\n\
2226Return successive r-length combinations of elements in the iterable.\n\n\
2227combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2228
2229static PyTypeObject combinations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002230 PyVarObject_HEAD_INIT(NULL, 0)
2231 "itertools.combinations", /* tp_name */
2232 sizeof(combinationsobject), /* tp_basicsize */
2233 0, /* tp_itemsize */
2234 /* methods */
2235 (destructor)combinations_dealloc, /* tp_dealloc */
2236 0, /* tp_print */
2237 0, /* tp_getattr */
2238 0, /* tp_setattr */
2239 0, /* tp_compare */
2240 0, /* tp_repr */
2241 0, /* tp_as_number */
2242 0, /* tp_as_sequence */
2243 0, /* tp_as_mapping */
2244 0, /* tp_hash */
2245 0, /* tp_call */
2246 0, /* tp_str */
2247 PyObject_GenericGetAttr, /* tp_getattro */
2248 0, /* tp_setattro */
2249 0, /* tp_as_buffer */
2250 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2251 Py_TPFLAGS_BASETYPE, /* tp_flags */
2252 combinations_doc, /* tp_doc */
2253 (traverseproc)combinations_traverse, /* tp_traverse */
2254 0, /* tp_clear */
2255 0, /* tp_richcompare */
2256 0, /* tp_weaklistoffset */
2257 PyObject_SelfIter, /* tp_iter */
2258 (iternextfunc)combinations_next, /* tp_iternext */
2259 0, /* tp_methods */
2260 0, /* tp_members */
2261 0, /* tp_getset */
2262 0, /* tp_base */
2263 0, /* tp_dict */
2264 0, /* tp_descr_get */
2265 0, /* tp_descr_set */
2266 0, /* tp_dictoffset */
2267 0, /* tp_init */
2268 0, /* tp_alloc */
2269 combinations_new, /* tp_new */
2270 PyObject_GC_Del, /* tp_free */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002271};
2272
2273
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002274/* combinations with replacement object *******************************************/
2275
2276/* Equivalent to:
2277
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002278 def combinations_with_replacement(iterable, r):
2279 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2280 # number items returned: (n+r-1)! / r! / (n-1)!
2281 pool = tuple(iterable)
2282 n = len(pool)
2283 indices = [0] * r
2284 yield tuple(pool[i] for i in indices)
2285 while 1:
2286 for i in reversed(range(r)):
2287 if indices[i] != n - 1:
2288 break
2289 else:
2290 return
2291 indices[i:] = [indices[i] + 1] * (r - i)
2292 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002293
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002294 def combinations_with_replacement2(iterable, r):
2295 'Alternate version that filters from product()'
2296 pool = tuple(iterable)
2297 n = len(pool)
2298 for indices in product(range(n), repeat=r):
2299 if sorted(indices) == list(indices):
2300 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002301*/
2302typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002303 PyObject_HEAD
2304 PyObject *pool; /* input converted to a tuple */
2305 Py_ssize_t *indices; /* one index per result element */
2306 PyObject *result; /* most recently returned result tuple */
2307 Py_ssize_t r; /* size of result tuple */
2308 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002309} cwrobject;
2310
2311static PyTypeObject cwr_type;
2312
2313static PyObject *
2314cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2315{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002316 cwrobject *co;
2317 Py_ssize_t n;
2318 Py_ssize_t r;
2319 PyObject *pool = NULL;
2320 PyObject *iterable = NULL;
2321 Py_ssize_t *indices = NULL;
2322 Py_ssize_t i;
2323 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002324
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002325 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2326 &iterable, &r))
2327 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002328
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002329 pool = PySequence_Tuple(iterable);
2330 if (pool == NULL)
2331 goto error;
2332 n = PyTuple_GET_SIZE(pool);
2333 if (r < 0) {
2334 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2335 goto error;
2336 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002337
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002338 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2339 if (indices == NULL) {
2340 PyErr_NoMemory();
2341 goto error;
2342 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002343
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002344 for (i=0 ; i<r ; i++)
2345 indices[i] = 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002346
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002347 /* create cwrobject structure */
2348 co = (cwrobject *)type->tp_alloc(type, 0);
2349 if (co == NULL)
2350 goto error;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002351
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002352 co->pool = pool;
2353 co->indices = indices;
2354 co->result = NULL;
2355 co->r = r;
2356 co->stopped = !n && r;
2357
2358 return (PyObject *)co;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002359
2360error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002361 if (indices != NULL)
2362 PyMem_Free(indices);
2363 Py_XDECREF(pool);
2364 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002365}
2366
2367static void
2368cwr_dealloc(cwrobject *co)
2369{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002370 PyObject_GC_UnTrack(co);
2371 Py_XDECREF(co->pool);
2372 Py_XDECREF(co->result);
2373 if (co->indices != NULL)
2374 PyMem_Free(co->indices);
2375 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002376}
2377
2378static int
2379cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2380{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002381 Py_VISIT(co->pool);
2382 Py_VISIT(co->result);
2383 return 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002384}
2385
2386static PyObject *
2387cwr_next(cwrobject *co)
2388{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002389 PyObject *elem;
2390 PyObject *oldelem;
2391 PyObject *pool = co->pool;
2392 Py_ssize_t *indices = co->indices;
2393 PyObject *result = co->result;
2394 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2395 Py_ssize_t r = co->r;
2396 Py_ssize_t i, j, index;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002397
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002398 if (co->stopped)
2399 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002400
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002401 if (result == NULL) {
2402 /* On the first pass, initialize result tuple using the indices */
2403 result = PyTuple_New(r);
2404 if (result == NULL)
2405 goto empty;
2406 co->result = result;
2407 for (i=0; i<r ; i++) {
2408 index = indices[i];
2409 elem = PyTuple_GET_ITEM(pool, index);
2410 Py_INCREF(elem);
2411 PyTuple_SET_ITEM(result, i, elem);
2412 }
2413 } else {
2414 /* Copy the previous result tuple or re-use it if available */
2415 if (Py_REFCNT(result) > 1) {
2416 PyObject *old_result = result;
2417 result = PyTuple_New(r);
2418 if (result == NULL)
2419 goto empty;
2420 co->result = result;
2421 for (i=0; i<r ; i++) {
2422 elem = PyTuple_GET_ITEM(old_result, i);
2423 Py_INCREF(elem);
2424 PyTuple_SET_ITEM(result, i, elem);
2425 }
2426 Py_DECREF(old_result);
2427 }
2428 /* Now, we've got the only copy so we can update it in-place CPython's
2429 empty tuple is a singleton and cached in PyTuple's freelist. */
2430 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002431
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002432 /* Scan indices right-to-left until finding one that is not
2433 * at its maximum (n-1). */
2434 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2435 ;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002436
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002437 /* If i is negative, then the indices are all at
2438 their maximum value and we're done. */
2439 if (i < 0)
2440 goto empty;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002441
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002442 /* Increment the current index which we know is not at its
2443 maximum. Then set all to the right to the same value. */
2444 indices[i]++;
2445 for (j=i+1 ; j<r ; j++)
2446 indices[j] = indices[j-1];
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002447
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002448 /* Update the result tuple for the new indices
2449 starting with i, the leftmost index that changed */
2450 for ( ; i<r ; i++) {
2451 index = indices[i];
2452 elem = PyTuple_GET_ITEM(pool, index);
2453 Py_INCREF(elem);
2454 oldelem = PyTuple_GET_ITEM(result, i);
2455 PyTuple_SET_ITEM(result, i, elem);
2456 Py_DECREF(oldelem);
2457 }
2458 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002459
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002460 Py_INCREF(result);
2461 return result;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002462
2463empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002464 co->stopped = 1;
2465 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002466}
2467
2468PyDoc_STRVAR(cwr_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002469"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002470\n\
2471Return successive r-length combinations of elements in the iterable\n\
2472allowing individual elements to have successive repeats.\n\
2473combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2474
2475static PyTypeObject cwr_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002476 PyVarObject_HEAD_INIT(NULL, 0)
2477 "itertools.combinations_with_replacement", /* tp_name */
2478 sizeof(cwrobject), /* tp_basicsize */
2479 0, /* tp_itemsize */
2480 /* methods */
2481 (destructor)cwr_dealloc, /* tp_dealloc */
2482 0, /* tp_print */
2483 0, /* tp_getattr */
2484 0, /* tp_setattr */
2485 0, /* tp_compare */
2486 0, /* tp_repr */
2487 0, /* tp_as_number */
2488 0, /* tp_as_sequence */
2489 0, /* tp_as_mapping */
2490 0, /* tp_hash */
2491 0, /* tp_call */
2492 0, /* tp_str */
2493 PyObject_GenericGetAttr, /* tp_getattro */
2494 0, /* tp_setattro */
2495 0, /* tp_as_buffer */
2496 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2497 Py_TPFLAGS_BASETYPE, /* tp_flags */
2498 cwr_doc, /* tp_doc */
2499 (traverseproc)cwr_traverse, /* tp_traverse */
2500 0, /* tp_clear */
2501 0, /* tp_richcompare */
2502 0, /* tp_weaklistoffset */
2503 PyObject_SelfIter, /* tp_iter */
2504 (iternextfunc)cwr_next, /* tp_iternext */
2505 0, /* tp_methods */
2506 0, /* tp_members */
2507 0, /* tp_getset */
2508 0, /* tp_base */
2509 0, /* tp_dict */
2510 0, /* tp_descr_get */
2511 0, /* tp_descr_set */
2512 0, /* tp_dictoffset */
2513 0, /* tp_init */
2514 0, /* tp_alloc */
2515 cwr_new, /* tp_new */
2516 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002517};
2518
2519
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002520/* permutations object ************************************************************
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002521
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002522def permutations(iterable, r=None):
2523 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2524 pool = tuple(iterable)
2525 n = len(pool)
2526 r = n if r is None else r
2527 indices = range(n)
2528 cycles = range(n-r+1, n+1)[::-1]
2529 yield tuple(pool[i] for i in indices[:r])
2530 while n:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002531 for i in reversed(range(r)):
2532 cycles[i] -= 1
2533 if cycles[i] == 0:
2534 indices[i:] = indices[i+1:] + indices[i:i+1]
2535 cycles[i] = n - i
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002536 else:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002537 j = cycles[i]
2538 indices[i], indices[-j] = indices[-j], indices[i]
2539 yield tuple(pool[i] for i in indices[:r])
2540 break
2541 else:
2542 return
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002543*/
2544
2545typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002546 PyObject_HEAD
2547 PyObject *pool; /* input converted to a tuple */
2548 Py_ssize_t *indices; /* one index per element in the pool */
2549 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2550 PyObject *result; /* most recently returned result tuple */
2551 Py_ssize_t r; /* size of result tuple */
2552 int stopped; /* set to 1 when the permutations iterator is exhausted */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002553} permutationsobject;
2554
2555static PyTypeObject permutations_type;
2556
2557static PyObject *
2558permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2559{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002560 permutationsobject *po;
2561 Py_ssize_t n;
2562 Py_ssize_t r;
2563 PyObject *robj = Py_None;
2564 PyObject *pool = NULL;
2565 PyObject *iterable = NULL;
2566 Py_ssize_t *indices = NULL;
2567 Py_ssize_t *cycles = NULL;
2568 Py_ssize_t i;
2569 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002570
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002571 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2572 &iterable, &robj))
2573 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002574
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002575 pool = PySequence_Tuple(iterable);
2576 if (pool == NULL)
2577 goto error;
2578 n = PyTuple_GET_SIZE(pool);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002579
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002580 r = n;
2581 if (robj != Py_None) {
2582 r = PyInt_AsSsize_t(robj);
2583 if (r == -1 && PyErr_Occurred())
2584 goto error;
2585 }
2586 if (r < 0) {
2587 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2588 goto error;
2589 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002590
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002591 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2592 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2593 if (indices == NULL || cycles == NULL) {
2594 PyErr_NoMemory();
2595 goto error;
2596 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002597
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002598 for (i=0 ; i<n ; i++)
2599 indices[i] = i;
2600 for (i=0 ; i<r ; i++)
2601 cycles[i] = n - i;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002602
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002603 /* create permutationsobject structure */
2604 po = (permutationsobject *)type->tp_alloc(type, 0);
2605 if (po == NULL)
2606 goto error;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002607
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002608 po->pool = pool;
2609 po->indices = indices;
2610 po->cycles = cycles;
2611 po->result = NULL;
2612 po->r = r;
2613 po->stopped = r > n ? 1 : 0;
2614
2615 return (PyObject *)po;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002616
2617error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002618 if (indices != NULL)
2619 PyMem_Free(indices);
2620 if (cycles != NULL)
2621 PyMem_Free(cycles);
2622 Py_XDECREF(pool);
2623 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002624}
2625
2626static void
2627permutations_dealloc(permutationsobject *po)
2628{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002629 PyObject_GC_UnTrack(po);
2630 Py_XDECREF(po->pool);
2631 Py_XDECREF(po->result);
2632 PyMem_Free(po->indices);
2633 PyMem_Free(po->cycles);
2634 Py_TYPE(po)->tp_free(po);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002635}
2636
2637static int
2638permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2639{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002640 Py_VISIT(po->pool);
2641 Py_VISIT(po->result);
2642 return 0;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002643}
2644
2645static PyObject *
2646permutations_next(permutationsobject *po)
2647{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002648 PyObject *elem;
2649 PyObject *oldelem;
2650 PyObject *pool = po->pool;
2651 Py_ssize_t *indices = po->indices;
2652 Py_ssize_t *cycles = po->cycles;
2653 PyObject *result = po->result;
2654 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2655 Py_ssize_t r = po->r;
2656 Py_ssize_t i, j, k, index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002657
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002658 if (po->stopped)
2659 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002660
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002661 if (result == NULL) {
2662 /* On the first pass, initialize result tuple using the indices */
2663 result = PyTuple_New(r);
2664 if (result == NULL)
2665 goto empty;
2666 po->result = result;
2667 for (i=0; i<r ; i++) {
2668 index = indices[i];
2669 elem = PyTuple_GET_ITEM(pool, index);
2670 Py_INCREF(elem);
2671 PyTuple_SET_ITEM(result, i, elem);
2672 }
2673 } else {
2674 if (n == 0)
2675 goto empty;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002676
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002677 /* Copy the previous result tuple or re-use it if available */
2678 if (Py_REFCNT(result) > 1) {
2679 PyObject *old_result = result;
2680 result = PyTuple_New(r);
2681 if (result == NULL)
2682 goto empty;
2683 po->result = result;
2684 for (i=0; i<r ; i++) {
2685 elem = PyTuple_GET_ITEM(old_result, i);
2686 Py_INCREF(elem);
2687 PyTuple_SET_ITEM(result, i, elem);
2688 }
2689 Py_DECREF(old_result);
2690 }
2691 /* Now, we've got the only copy so we can update it in-place */
2692 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002693
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002694 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2695 for (i=r-1 ; i>=0 ; i--) {
2696 cycles[i] -= 1;
2697 if (cycles[i] == 0) {
2698 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2699 index = indices[i];
2700 for (j=i ; j<n-1 ; j++)
2701 indices[j] = indices[j+1];
2702 indices[n-1] = index;
2703 cycles[i] = n - i;
2704 } else {
2705 j = cycles[i];
2706 index = indices[i];
2707 indices[i] = indices[n-j];
2708 indices[n-j] = index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002709
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002710 for (k=i; k<r ; k++) {
2711 /* start with i, the leftmost element that changed */
2712 /* yield tuple(pool[k] for k in indices[:r]) */
2713 index = indices[k];
2714 elem = PyTuple_GET_ITEM(pool, index);
2715 Py_INCREF(elem);
2716 oldelem = PyTuple_GET_ITEM(result, k);
2717 PyTuple_SET_ITEM(result, k, elem);
2718 Py_DECREF(oldelem);
2719 }
2720 break;
2721 }
2722 }
2723 /* If i is negative, then the cycles have all
2724 rolled-over and we're done. */
2725 if (i < 0)
2726 goto empty;
2727 }
2728 Py_INCREF(result);
2729 return result;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002730
2731empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002732 po->stopped = 1;
2733 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002734}
2735
2736PyDoc_STRVAR(permutations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +00002737"permutations(iterable[, r]) --> permutations object\n\
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002738\n\
2739Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl8550d692008-06-10 12:46:39 +00002740permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002741
2742static PyTypeObject permutations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002743 PyVarObject_HEAD_INIT(NULL, 0)
2744 "itertools.permutations", /* tp_name */
2745 sizeof(permutationsobject), /* tp_basicsize */
2746 0, /* tp_itemsize */
2747 /* methods */
2748 (destructor)permutations_dealloc, /* tp_dealloc */
2749 0, /* tp_print */
2750 0, /* tp_getattr */
2751 0, /* tp_setattr */
2752 0, /* tp_compare */
2753 0, /* tp_repr */
2754 0, /* tp_as_number */
2755 0, /* tp_as_sequence */
2756 0, /* tp_as_mapping */
2757 0, /* tp_hash */
2758 0, /* tp_call */
2759 0, /* tp_str */
2760 PyObject_GenericGetAttr, /* tp_getattro */
2761 0, /* tp_setattro */
2762 0, /* tp_as_buffer */
2763 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2764 Py_TPFLAGS_BASETYPE, /* tp_flags */
2765 permutations_doc, /* tp_doc */
2766 (traverseproc)permutations_traverse, /* tp_traverse */
2767 0, /* tp_clear */
2768 0, /* tp_richcompare */
2769 0, /* tp_weaklistoffset */
2770 PyObject_SelfIter, /* tp_iter */
2771 (iternextfunc)permutations_next, /* tp_iternext */
2772 0, /* tp_methods */
2773 0, /* tp_members */
2774 0, /* tp_getset */
2775 0, /* tp_base */
2776 0, /* tp_dict */
2777 0, /* tp_descr_get */
2778 0, /* tp_descr_set */
2779 0, /* tp_dictoffset */
2780 0, /* tp_init */
2781 0, /* tp_alloc */
2782 permutations_new, /* tp_new */
2783 PyObject_GC_Del, /* tp_free */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002784};
2785
2786
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002787/* compress object ************************************************************/
2788
2789/* Equivalent to:
2790
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002791 def compress(data, selectors):
2792 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2793 return (d for d, s in izip(data, selectors) if s)
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002794*/
2795
2796typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002797 PyObject_HEAD
2798 PyObject *data;
2799 PyObject *selectors;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002800} compressobject;
2801
2802static PyTypeObject compress_type;
2803
2804static PyObject *
2805compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2806{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002807 PyObject *seq1, *seq2;
2808 PyObject *data=NULL, *selectors=NULL;
2809 compressobject *lz;
2810 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002811
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002812 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2813 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002814
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002815 data = PyObject_GetIter(seq1);
2816 if (data == NULL)
2817 goto fail;
2818 selectors = PyObject_GetIter(seq2);
2819 if (selectors == NULL)
2820 goto fail;
2821
2822 /* create compressobject structure */
2823 lz = (compressobject *)type->tp_alloc(type, 0);
2824 if (lz == NULL)
2825 goto fail;
2826 lz->data = data;
2827 lz->selectors = selectors;
2828 return (PyObject *)lz;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002829
2830fail:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002831 Py_XDECREF(data);
2832 Py_XDECREF(selectors);
2833 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002834}
2835
2836static void
2837compress_dealloc(compressobject *lz)
2838{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002839 PyObject_GC_UnTrack(lz);
2840 Py_XDECREF(lz->data);
2841 Py_XDECREF(lz->selectors);
2842 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002843}
2844
2845static int
2846compress_traverse(compressobject *lz, visitproc visit, void *arg)
2847{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002848 Py_VISIT(lz->data);
2849 Py_VISIT(lz->selectors);
2850 return 0;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002851}
2852
2853static PyObject *
2854compress_next(compressobject *lz)
2855{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002856 PyObject *data = lz->data, *selectors = lz->selectors;
2857 PyObject *datum, *selector;
2858 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2859 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2860 int ok;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002861
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002862 while (1) {
2863 /* Steps: get datum, get selector, evaluate selector.
2864 Order is important (to match the pure python version
2865 in terms of which input gets a chance to raise an
2866 exception first).
2867 */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002868
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002869 datum = datanext(data);
2870 if (datum == NULL)
2871 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002872
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002873 selector = selectornext(selectors);
2874 if (selector == NULL) {
2875 Py_DECREF(datum);
2876 return NULL;
2877 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002878
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002879 ok = PyObject_IsTrue(selector);
2880 Py_DECREF(selector);
2881 if (ok == 1)
2882 return datum;
2883 Py_DECREF(datum);
2884 if (ok == -1)
2885 return NULL;
2886 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002887}
2888
2889PyDoc_STRVAR(compress_doc,
Raymond Hettinger2e2909f2009-02-19 02:15:14 +00002890"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002891\n\
2892Return data elements corresponding to true selector elements.\n\
2893Forms a shorter iterator from selected data elements using the\n\
2894selectors to choose the data elements.");
2895
2896static PyTypeObject compress_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002897 PyVarObject_HEAD_INIT(NULL, 0)
2898 "itertools.compress", /* tp_name */
2899 sizeof(compressobject), /* tp_basicsize */
2900 0, /* tp_itemsize */
2901 /* methods */
2902 (destructor)compress_dealloc, /* tp_dealloc */
2903 0, /* tp_print */
2904 0, /* tp_getattr */
2905 0, /* tp_setattr */
2906 0, /* tp_compare */
2907 0, /* tp_repr */
2908 0, /* tp_as_number */
2909 0, /* tp_as_sequence */
2910 0, /* tp_as_mapping */
2911 0, /* tp_hash */
2912 0, /* tp_call */
2913 0, /* tp_str */
2914 PyObject_GenericGetAttr, /* tp_getattro */
2915 0, /* tp_setattro */
2916 0, /* tp_as_buffer */
2917 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2918 Py_TPFLAGS_BASETYPE, /* tp_flags */
2919 compress_doc, /* tp_doc */
2920 (traverseproc)compress_traverse, /* tp_traverse */
2921 0, /* tp_clear */
2922 0, /* tp_richcompare */
2923 0, /* tp_weaklistoffset */
2924 PyObject_SelfIter, /* tp_iter */
2925 (iternextfunc)compress_next, /* tp_iternext */
2926 0, /* tp_methods */
2927 0, /* tp_members */
2928 0, /* tp_getset */
2929 0, /* tp_base */
2930 0, /* tp_dict */
2931 0, /* tp_descr_get */
2932 0, /* tp_descr_set */
2933 0, /* tp_dictoffset */
2934 0, /* tp_init */
2935 0, /* tp_alloc */
2936 compress_new, /* tp_new */
2937 PyObject_GC_Del, /* tp_free */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002938};
2939
2940
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002941/* ifilter object ************************************************************/
2942
2943typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002944 PyObject_HEAD
2945 PyObject *func;
2946 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002947} ifilterobject;
2948
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002949static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002950
2951static PyObject *
2952ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2953{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002954 PyObject *func, *seq;
2955 PyObject *it;
2956 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002957
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002958 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
2959 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002960
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002961 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
2962 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002963
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002964 /* Get iterator. */
2965 it = PyObject_GetIter(seq);
2966 if (it == NULL)
2967 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002968
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002969 /* create ifilterobject structure */
2970 lz = (ifilterobject *)type->tp_alloc(type, 0);
2971 if (lz == NULL) {
2972 Py_DECREF(it);
2973 return NULL;
2974 }
2975 Py_INCREF(func);
2976 lz->func = func;
2977 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002978
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002979 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002980}
2981
2982static void
2983ifilter_dealloc(ifilterobject *lz)
2984{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002985 PyObject_GC_UnTrack(lz);
2986 Py_XDECREF(lz->func);
2987 Py_XDECREF(lz->it);
2988 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002989}
2990
2991static int
2992ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
2993{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002994 Py_VISIT(lz->it);
2995 Py_VISIT(lz->func);
2996 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002997}
2998
2999static PyObject *
3000ifilter_next(ifilterobject *lz)
3001{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003002 PyObject *item;
3003 PyObject *it = lz->it;
3004 long ok;
3005 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003006
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003007 iternext = *Py_TYPE(it)->tp_iternext;
3008 for (;;) {
3009 item = iternext(it);
3010 if (item == NULL)
3011 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003012
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003013 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3014 ok = PyObject_IsTrue(item);
3015 } else {
3016 PyObject *good;
3017 good = PyObject_CallFunctionObjArgs(lz->func,
3018 item, NULL);
3019 if (good == NULL) {
3020 Py_DECREF(item);
3021 return NULL;
3022 }
3023 ok = PyObject_IsTrue(good);
3024 Py_DECREF(good);
3025 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003026 if (ok > 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003027 return item;
3028 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003029 if (ok < 0)
3030 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003031 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003032}
3033
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003034PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003035"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003036\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003037Return those items of sequence for which function(item) is true.\n\
3038If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003039
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003040static PyTypeObject ifilter_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003041 PyVarObject_HEAD_INIT(NULL, 0)
3042 "itertools.ifilter", /* tp_name */
3043 sizeof(ifilterobject), /* tp_basicsize */
3044 0, /* tp_itemsize */
3045 /* methods */
3046 (destructor)ifilter_dealloc, /* tp_dealloc */
3047 0, /* tp_print */
3048 0, /* tp_getattr */
3049 0, /* tp_setattr */
3050 0, /* tp_compare */
3051 0, /* tp_repr */
3052 0, /* tp_as_number */
3053 0, /* tp_as_sequence */
3054 0, /* tp_as_mapping */
3055 0, /* tp_hash */
3056 0, /* tp_call */
3057 0, /* tp_str */
3058 PyObject_GenericGetAttr, /* tp_getattro */
3059 0, /* tp_setattro */
3060 0, /* tp_as_buffer */
3061 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3062 Py_TPFLAGS_BASETYPE, /* tp_flags */
3063 ifilter_doc, /* tp_doc */
3064 (traverseproc)ifilter_traverse, /* tp_traverse */
3065 0, /* tp_clear */
3066 0, /* tp_richcompare */
3067 0, /* tp_weaklistoffset */
3068 PyObject_SelfIter, /* tp_iter */
3069 (iternextfunc)ifilter_next, /* tp_iternext */
3070 0, /* tp_methods */
3071 0, /* tp_members */
3072 0, /* tp_getset */
3073 0, /* tp_base */
3074 0, /* tp_dict */
3075 0, /* tp_descr_get */
3076 0, /* tp_descr_set */
3077 0, /* tp_dictoffset */
3078 0, /* tp_init */
3079 0, /* tp_alloc */
3080 ifilter_new, /* tp_new */
3081 PyObject_GC_Del, /* tp_free */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003082};
3083
3084
3085/* ifilterfalse object ************************************************************/
3086
3087typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003088 PyObject_HEAD
3089 PyObject *func;
3090 PyObject *it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003091} ifilterfalseobject;
3092
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003093static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003094
3095static PyObject *
3096ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3097{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003098 PyObject *func, *seq;
3099 PyObject *it;
3100 ifilterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003101
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003102 if (type == &ifilterfalse_type &&
3103 !_PyArg_NoKeywords("ifilterfalse()", kwds))
3104 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003105
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003106 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
3107 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003108
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003109 /* Get iterator. */
3110 it = PyObject_GetIter(seq);
3111 if (it == NULL)
3112 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003113
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003114 /* create ifilterfalseobject structure */
3115 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
3116 if (lz == NULL) {
3117 Py_DECREF(it);
3118 return NULL;
3119 }
3120 Py_INCREF(func);
3121 lz->func = func;
3122 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003123
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003124 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003125}
3126
3127static void
3128ifilterfalse_dealloc(ifilterfalseobject *lz)
3129{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003130 PyObject_GC_UnTrack(lz);
3131 Py_XDECREF(lz->func);
3132 Py_XDECREF(lz->it);
3133 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003134}
3135
3136static int
3137ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
3138{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003139 Py_VISIT(lz->it);
3140 Py_VISIT(lz->func);
3141 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003142}
3143
3144static PyObject *
3145ifilterfalse_next(ifilterfalseobject *lz)
3146{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003147 PyObject *item;
3148 PyObject *it = lz->it;
3149 long ok;
3150 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003151
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003152 iternext = *Py_TYPE(it)->tp_iternext;
3153 for (;;) {
3154 item = iternext(it);
3155 if (item == NULL)
3156 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003157
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003158 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3159 ok = PyObject_IsTrue(item);
3160 } else {
3161 PyObject *good;
3162 good = PyObject_CallFunctionObjArgs(lz->func,
3163 item, NULL);
3164 if (good == NULL) {
3165 Py_DECREF(item);
3166 return NULL;
3167 }
3168 ok = PyObject_IsTrue(good);
3169 Py_DECREF(good);
3170 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003171 if (ok == 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003172 return item;
3173 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003174 if (ok < 0)
3175 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003176 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003177}
3178
Raymond Hettinger60eca932003-02-09 06:40:58 +00003179PyDoc_STRVAR(ifilterfalse_doc,
3180"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
3181\n\
3182Return those items of sequence for which function(item) is false.\n\
3183If function is None, return the items that are false.");
3184
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003185static PyTypeObject ifilterfalse_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003186 PyVarObject_HEAD_INIT(NULL, 0)
3187 "itertools.ifilterfalse", /* tp_name */
3188 sizeof(ifilterfalseobject), /* tp_basicsize */
3189 0, /* tp_itemsize */
3190 /* methods */
3191 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
3192 0, /* tp_print */
3193 0, /* tp_getattr */
3194 0, /* tp_setattr */
3195 0, /* tp_compare */
3196 0, /* tp_repr */
3197 0, /* tp_as_number */
3198 0, /* tp_as_sequence */
3199 0, /* tp_as_mapping */
3200 0, /* tp_hash */
3201 0, /* tp_call */
3202 0, /* tp_str */
3203 PyObject_GenericGetAttr, /* tp_getattro */
3204 0, /* tp_setattro */
3205 0, /* tp_as_buffer */
3206 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3207 Py_TPFLAGS_BASETYPE, /* tp_flags */
3208 ifilterfalse_doc, /* tp_doc */
3209 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
3210 0, /* tp_clear */
3211 0, /* tp_richcompare */
3212 0, /* tp_weaklistoffset */
3213 PyObject_SelfIter, /* tp_iter */
3214 (iternextfunc)ifilterfalse_next, /* tp_iternext */
3215 0, /* tp_methods */
3216 0, /* tp_members */
3217 0, /* tp_getset */
3218 0, /* tp_base */
3219 0, /* tp_dict */
3220 0, /* tp_descr_get */
3221 0, /* tp_descr_set */
3222 0, /* tp_dictoffset */
3223 0, /* tp_init */
3224 0, /* tp_alloc */
3225 ifilterfalse_new, /* tp_new */
3226 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003227};
3228
3229
3230/* count object ************************************************************/
3231
3232typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003233 PyObject_HEAD
3234 Py_ssize_t cnt;
3235 PyObject *long_cnt;
3236 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003237} countobject;
3238
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003239/* Counting logic and invariants:
3240
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003241fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003242
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003243 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3244 Advances with: cnt += 1
3245 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003246
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003247slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003248
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003249 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3250 All counting is done with python objects (no overflows or underflows).
3251 Advances with: long_cnt += long_step
3252 Step may be zero -- effectively a slow version of repeat(cnt).
3253 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003254*/
3255
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003256static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003257
3258static PyObject *
3259count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3260{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003261 countobject *lz;
3262 int slow_mode = 0;
3263 Py_ssize_t cnt = 0;
3264 PyObject *long_cnt = NULL;
3265 PyObject *long_step = NULL;
3266 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003267
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003268 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3269 kwlist, &long_cnt, &long_step))
3270 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003271
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003272 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3273 (long_step != NULL && !PyNumber_Check(long_step))) {
3274 PyErr_SetString(PyExc_TypeError, "a number is required");
3275 return NULL;
3276 }
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003277
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003278 if (long_cnt != NULL) {
3279 cnt = PyInt_AsSsize_t(long_cnt);
3280 if ((cnt == -1 && PyErr_Occurred()) || !PyInt_Check(long_cnt)) {
3281 PyErr_Clear();
3282 slow_mode = 1;
3283 }
3284 Py_INCREF(long_cnt);
3285 } else {
3286 cnt = 0;
3287 long_cnt = PyInt_FromLong(0);
3288 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003289
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003290 /* If not specified, step defaults to 1 */
3291 if (long_step == NULL) {
3292 long_step = PyInt_FromLong(1);
3293 if (long_step == NULL) {
3294 Py_DECREF(long_cnt);
3295 return NULL;
3296 }
3297 } else
3298 Py_INCREF(long_step);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003299
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003300 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003301
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003302 /* Fast mode only works when the step is 1 */
3303 if (!PyInt_Check(long_step) ||
3304 PyInt_AS_LONG(long_step) != 1) {
3305 slow_mode = 1;
3306 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003307
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003308 if (slow_mode)
3309 cnt = PY_SSIZE_T_MAX;
3310 else
3311 Py_CLEAR(long_cnt);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003312
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003313 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3314 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3315 assert(slow_mode ||
3316 (PyInt_Check(long_step) && PyInt_AS_LONG(long_step) == 1));
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003317
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003318 /* create countobject structure */
3319 lz = (countobject *)type->tp_alloc(type, 0);
3320 if (lz == NULL) {
3321 Py_XDECREF(long_cnt);
3322 return NULL;
3323 }
3324 lz->cnt = cnt;
3325 lz->long_cnt = long_cnt;
3326 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003327
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003328 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003329}
3330
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003331static void
3332count_dealloc(countobject *lz)
3333{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003334 PyObject_GC_UnTrack(lz);
3335 Py_XDECREF(lz->long_cnt);
3336 Py_XDECREF(lz->long_step);
3337 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003338}
3339
Benjamin Peterson062a7c32009-02-16 21:07:52 +00003340static int
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003341count_traverse(countobject *lz, visitproc visit, void *arg)
3342{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003343 Py_VISIT(lz->long_cnt);
3344 Py_VISIT(lz->long_step);
3345 return 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003346}
3347
3348static PyObject *
3349count_nextlong(countobject *lz)
3350{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003351 PyObject *long_cnt;
3352 PyObject *stepped_up;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003353
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003354 long_cnt = lz->long_cnt;
3355 if (long_cnt == NULL) {
3356 /* Switch to slow_mode */
3357 long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
3358 if (long_cnt == NULL)
3359 return NULL;
3360 }
3361 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003362
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003363 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3364 if (stepped_up == NULL)
3365 return NULL;
3366 lz->long_cnt = stepped_up;
3367 return long_cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003368}
3369
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003370static PyObject *
3371count_next(countobject *lz)
3372{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003373 if (lz->cnt == PY_SSIZE_T_MAX)
3374 return count_nextlong(lz);
3375 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003376}
3377
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003378static PyObject *
3379count_repr(countobject *lz)
3380{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003381 PyObject *cnt_repr, *step_repr = NULL;
3382 PyObject *result = NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003383
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003384 if (lz->cnt != PY_SSIZE_T_MAX)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003385 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003386
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003387 cnt_repr = PyObject_Repr(lz->long_cnt);
3388 if (cnt_repr == NULL)
3389 return NULL;
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003390
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003391 if (PyInt_Check(lz->long_step) && PyInt_AS_LONG(lz->long_step) == 1) {
3392 /* Don't display step when it is an integer equal to 1 */
3393 result = PyString_FromFormat("count(%s)",
3394 PyString_AS_STRING(cnt_repr));
3395 } else {
3396 step_repr = PyObject_Repr(lz->long_step);
3397 if (step_repr != NULL)
3398 result = PyString_FromFormat("count(%s, %s)",
3399 PyString_AS_STRING(cnt_repr),
3400 PyString_AS_STRING(step_repr));
3401 }
3402 Py_DECREF(cnt_repr);
3403 Py_XDECREF(step_repr);
3404 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003405}
3406
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003407static PyObject *
3408count_reduce(countobject *lz)
3409{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003410 if (lz->cnt == PY_SSIZE_T_MAX)
3411 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3412 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003413}
3414
3415PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3416
3417static PyMethodDef count_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003418 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3419 count_reduce_doc},
3420 {NULL, NULL} /* sentinel */
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003421};
3422
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003423PyDoc_STRVAR(count_doc,
Ezio Melottifdc1e0d2010-06-11 02:21:25 +00003424 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003425\n\
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003426Return a count object whose .next() method returns consecutive values.\n\
3427Equivalent to:\n\n\
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003428 def count(firstval=0, step=1):\n\
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003429 x = firstval\n\
3430 while 1:\n\
3431 yield x\n\
3432 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003433
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003434static PyTypeObject count_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003435 PyVarObject_HEAD_INIT(NULL, 0)
3436 "itertools.count", /* tp_name */
3437 sizeof(countobject), /* tp_basicsize */
3438 0, /* tp_itemsize */
3439 /* methods */
3440 (destructor)count_dealloc, /* tp_dealloc */
3441 0, /* tp_print */
3442 0, /* tp_getattr */
3443 0, /* tp_setattr */
3444 0, /* tp_compare */
3445 (reprfunc)count_repr, /* tp_repr */
3446 0, /* tp_as_number */
3447 0, /* tp_as_sequence */
3448 0, /* tp_as_mapping */
3449 0, /* tp_hash */
3450 0, /* tp_call */
3451 0, /* tp_str */
3452 PyObject_GenericGetAttr, /* tp_getattro */
3453 0, /* tp_setattro */
3454 0, /* tp_as_buffer */
3455 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3456 Py_TPFLAGS_BASETYPE, /* tp_flags */
3457 count_doc, /* tp_doc */
3458 (traverseproc)count_traverse, /* tp_traverse */
3459 0, /* tp_clear */
3460 0, /* tp_richcompare */
3461 0, /* tp_weaklistoffset */
3462 PyObject_SelfIter, /* tp_iter */
3463 (iternextfunc)count_next, /* tp_iternext */
3464 count_methods, /* tp_methods */
3465 0, /* tp_members */
3466 0, /* tp_getset */
3467 0, /* tp_base */
3468 0, /* tp_dict */
3469 0, /* tp_descr_get */
3470 0, /* tp_descr_set */
3471 0, /* tp_dictoffset */
3472 0, /* tp_init */
3473 0, /* tp_alloc */
3474 count_new, /* tp_new */
3475 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003476};
3477
3478
3479/* izip object ************************************************************/
3480
3481#include "Python.h"
3482
3483typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003484 PyObject_HEAD
3485 Py_ssize_t tuplesize;
3486 PyObject *ittuple; /* tuple of iterators */
3487 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003488} izipobject;
3489
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003490static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003491
3492static PyObject *
3493izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3494{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003495 izipobject *lz;
3496 Py_ssize_t i;
3497 PyObject *ittuple; /* tuple of iterators */
3498 PyObject *result;
3499 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003500
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003501 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
3502 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003503
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003504 /* args must be a tuple */
3505 assert(PyTuple_Check(args));
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003506
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003507 /* obtain iterators */
3508 ittuple = PyTuple_New(tuplesize);
3509 if (ittuple == NULL)
3510 return NULL;
3511 for (i=0; i < tuplesize; ++i) {
3512 PyObject *item = PyTuple_GET_ITEM(args, i);
3513 PyObject *it = PyObject_GetIter(item);
3514 if (it == NULL) {
3515 if (PyErr_ExceptionMatches(PyExc_TypeError))
3516 PyErr_Format(PyExc_TypeError,
3517 "izip argument #%zd must support iteration",
3518 i+1);
3519 Py_DECREF(ittuple);
3520 return NULL;
3521 }
3522 PyTuple_SET_ITEM(ittuple, i, it);
3523 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003524
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003525 /* create a result holder */
3526 result = PyTuple_New(tuplesize);
3527 if (result == NULL) {
3528 Py_DECREF(ittuple);
3529 return NULL;
3530 }
3531 for (i=0 ; i < tuplesize ; i++) {
3532 Py_INCREF(Py_None);
3533 PyTuple_SET_ITEM(result, i, Py_None);
3534 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003535
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003536 /* create izipobject structure */
3537 lz = (izipobject *)type->tp_alloc(type, 0);
3538 if (lz == NULL) {
3539 Py_DECREF(ittuple);
3540 Py_DECREF(result);
3541 return NULL;
3542 }
3543 lz->ittuple = ittuple;
3544 lz->tuplesize = tuplesize;
3545 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003546
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003547 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003548}
3549
3550static void
3551izip_dealloc(izipobject *lz)
3552{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003553 PyObject_GC_UnTrack(lz);
3554 Py_XDECREF(lz->ittuple);
3555 Py_XDECREF(lz->result);
3556 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003557}
3558
3559static int
3560izip_traverse(izipobject *lz, visitproc visit, void *arg)
3561{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003562 Py_VISIT(lz->ittuple);
3563 Py_VISIT(lz->result);
3564 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003565}
3566
3567static PyObject *
3568izip_next(izipobject *lz)
3569{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003570 Py_ssize_t i;
3571 Py_ssize_t tuplesize = lz->tuplesize;
3572 PyObject *result = lz->result;
3573 PyObject *it;
3574 PyObject *item;
3575 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003576
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003577 if (tuplesize == 0)
3578 return NULL;
3579 if (Py_REFCNT(result) == 1) {
3580 Py_INCREF(result);
3581 for (i=0 ; i < tuplesize ; i++) {
3582 it = PyTuple_GET_ITEM(lz->ittuple, i);
3583 item = (*Py_TYPE(it)->tp_iternext)(it);
3584 if (item == NULL) {
3585 Py_DECREF(result);
3586 return NULL;
3587 }
3588 olditem = PyTuple_GET_ITEM(result, i);
3589 PyTuple_SET_ITEM(result, i, item);
3590 Py_DECREF(olditem);
3591 }
3592 } else {
3593 result = PyTuple_New(tuplesize);
3594 if (result == NULL)
3595 return NULL;
3596 for (i=0 ; i < tuplesize ; i++) {
3597 it = PyTuple_GET_ITEM(lz->ittuple, i);
3598 item = (*Py_TYPE(it)->tp_iternext)(it);
3599 if (item == NULL) {
3600 Py_DECREF(result);
3601 return NULL;
3602 }
3603 PyTuple_SET_ITEM(result, i, item);
3604 }
3605 }
3606 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003607}
3608
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003609PyDoc_STRVAR(izip_doc,
3610"izip(iter1 [,iter2 [...]]) --> izip object\n\
3611\n\
3612Return a izip object whose .next() method returns a tuple where\n\
3613the i-th element comes from the i-th iterable argument. The .next()\n\
3614method continues until the shortest iterable in the argument sequence\n\
3615is exhausted and then it raises StopIteration. Works like the zip()\n\
3616function but consumes less memory by returning an iterator instead of\n\
3617a list.");
3618
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003619static PyTypeObject izip_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003620 PyVarObject_HEAD_INIT(NULL, 0)
3621 "itertools.izip", /* tp_name */
3622 sizeof(izipobject), /* tp_basicsize */
3623 0, /* tp_itemsize */
3624 /* methods */
3625 (destructor)izip_dealloc, /* tp_dealloc */
3626 0, /* tp_print */
3627 0, /* tp_getattr */
3628 0, /* tp_setattr */
3629 0, /* tp_compare */
3630 0, /* tp_repr */
3631 0, /* tp_as_number */
3632 0, /* tp_as_sequence */
3633 0, /* tp_as_mapping */
3634 0, /* tp_hash */
3635 0, /* tp_call */
3636 0, /* tp_str */
3637 PyObject_GenericGetAttr, /* tp_getattro */
3638 0, /* tp_setattro */
3639 0, /* tp_as_buffer */
3640 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3641 Py_TPFLAGS_BASETYPE, /* tp_flags */
3642 izip_doc, /* tp_doc */
3643 (traverseproc)izip_traverse, /* tp_traverse */
3644 0, /* tp_clear */
3645 0, /* tp_richcompare */
3646 0, /* tp_weaklistoffset */
3647 PyObject_SelfIter, /* tp_iter */
3648 (iternextfunc)izip_next, /* tp_iternext */
3649 0, /* tp_methods */
3650 0, /* tp_members */
3651 0, /* tp_getset */
3652 0, /* tp_base */
3653 0, /* tp_dict */
3654 0, /* tp_descr_get */
3655 0, /* tp_descr_set */
3656 0, /* tp_dictoffset */
3657 0, /* tp_init */
3658 0, /* tp_alloc */
3659 izip_new, /* tp_new */
3660 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003661};
3662
3663
3664/* repeat object ************************************************************/
3665
3666typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003667 PyObject_HEAD
3668 PyObject *element;
3669 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003670} repeatobject;
3671
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003672static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003673
3674static PyObject *
3675repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3676{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003677 repeatobject *ro;
3678 PyObject *element;
3679 Py_ssize_t cnt = -1;
3680 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003681
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003682 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3683 &element, &cnt))
3684 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003685
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003686 if (PyTuple_Size(args) == 2 && cnt < 0)
3687 cnt = 0;
3688
3689 ro = (repeatobject *)type->tp_alloc(type, 0);
3690 if (ro == NULL)
3691 return NULL;
3692 Py_INCREF(element);
3693 ro->element = element;
3694 ro->cnt = cnt;
3695 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003696}
3697
3698static void
3699repeat_dealloc(repeatobject *ro)
3700{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003701 PyObject_GC_UnTrack(ro);
3702 Py_XDECREF(ro->element);
3703 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003704}
3705
3706static int
3707repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3708{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003709 Py_VISIT(ro->element);
3710 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003711}
3712
3713static PyObject *
3714repeat_next(repeatobject *ro)
3715{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003716 if (ro->cnt == 0)
3717 return NULL;
3718 if (ro->cnt > 0)
3719 ro->cnt--;
3720 Py_INCREF(ro->element);
3721 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003722}
3723
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003724static PyObject *
3725repeat_repr(repeatobject *ro)
3726{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003727 PyObject *result, *objrepr;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003728
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003729 objrepr = PyObject_Repr(ro->element);
3730 if (objrepr == NULL)
3731 return NULL;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003732
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003733 if (ro->cnt == -1)
3734 result = PyString_FromFormat("repeat(%s)",
3735 PyString_AS_STRING(objrepr));
3736 else
3737 result = PyString_FromFormat("repeat(%s, %zd)",
3738 PyString_AS_STRING(objrepr), ro->cnt);
3739 Py_DECREF(objrepr);
3740 return result;
3741}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003742
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003743static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003744repeat_len(repeatobject *ro)
3745{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003746 if (ro->cnt == -1) {
3747 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3748 return NULL;
3749 }
3750 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003751}
3752
Armin Rigof5b3e362006-02-11 21:32:43 +00003753PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003754
3755static PyMethodDef repeat_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003756 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3757 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003758};
3759
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003760PyDoc_STRVAR(repeat_doc,
Raymond Hettinger182edae2009-02-19 02:38:25 +00003761"repeat(object [,times]) -> create an iterator which returns the object\n\
3762for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003763endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003764
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003765static PyTypeObject repeat_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003766 PyVarObject_HEAD_INIT(NULL, 0)
3767 "itertools.repeat", /* tp_name */
3768 sizeof(repeatobject), /* tp_basicsize */
3769 0, /* tp_itemsize */
3770 /* methods */
3771 (destructor)repeat_dealloc, /* tp_dealloc */
3772 0, /* tp_print */
3773 0, /* tp_getattr */
3774 0, /* tp_setattr */
3775 0, /* tp_compare */
3776 (reprfunc)repeat_repr, /* tp_repr */
3777 0, /* tp_as_number */
3778 0, /* tp_as_sequence */
3779 0, /* tp_as_mapping */
3780 0, /* tp_hash */
3781 0, /* tp_call */
3782 0, /* tp_str */
3783 PyObject_GenericGetAttr, /* tp_getattro */
3784 0, /* tp_setattro */
3785 0, /* tp_as_buffer */
3786 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3787 Py_TPFLAGS_BASETYPE, /* tp_flags */
3788 repeat_doc, /* tp_doc */
3789 (traverseproc)repeat_traverse, /* tp_traverse */
3790 0, /* tp_clear */
3791 0, /* tp_richcompare */
3792 0, /* tp_weaklistoffset */
3793 PyObject_SelfIter, /* tp_iter */
3794 (iternextfunc)repeat_next, /* tp_iternext */
3795 repeat_methods, /* tp_methods */
3796 0, /* tp_members */
3797 0, /* tp_getset */
3798 0, /* tp_base */
3799 0, /* tp_dict */
3800 0, /* tp_descr_get */
3801 0, /* tp_descr_set */
3802 0, /* tp_dictoffset */
3803 0, /* tp_init */
3804 0, /* tp_alloc */
3805 repeat_new, /* tp_new */
3806 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003807};
3808
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003809/* iziplongest object ************************************************************/
3810
3811#include "Python.h"
3812
3813typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003814 PyObject_HEAD
3815 Py_ssize_t tuplesize;
3816 Py_ssize_t numactive;
3817 PyObject *ittuple; /* tuple of iterators */
3818 PyObject *result;
3819 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003820} iziplongestobject;
3821
3822static PyTypeObject iziplongest_type;
3823
3824static PyObject *
3825izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3826{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003827 iziplongestobject *lz;
3828 Py_ssize_t i;
3829 PyObject *ittuple; /* tuple of iterators */
3830 PyObject *result;
3831 PyObject *fillvalue = Py_None;
3832 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003833
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003834 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3835 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3836 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3837 PyErr_SetString(PyExc_TypeError,
3838 "izip_longest() got an unexpected keyword argument");
3839 return NULL;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003840 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003841 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003842
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003843 /* args must be a tuple */
3844 assert(PyTuple_Check(args));
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003845
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003846 /* obtain iterators */
3847 ittuple = PyTuple_New(tuplesize);
3848 if (ittuple == NULL)
3849 return NULL;
3850 for (i=0; i < tuplesize; ++i) {
3851 PyObject *item = PyTuple_GET_ITEM(args, i);
3852 PyObject *it = PyObject_GetIter(item);
3853 if (it == NULL) {
3854 if (PyErr_ExceptionMatches(PyExc_TypeError))
3855 PyErr_Format(PyExc_TypeError,
3856 "izip_longest argument #%zd must support iteration",
3857 i+1);
3858 Py_DECREF(ittuple);
3859 return NULL;
3860 }
3861 PyTuple_SET_ITEM(ittuple, i, it);
3862 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003863
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003864 /* create a result holder */
3865 result = PyTuple_New(tuplesize);
3866 if (result == NULL) {
3867 Py_DECREF(ittuple);
3868 return NULL;
3869 }
3870 for (i=0 ; i < tuplesize ; i++) {
3871 Py_INCREF(Py_None);
3872 PyTuple_SET_ITEM(result, i, Py_None);
3873 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003874
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003875 /* create iziplongestobject structure */
3876 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3877 if (lz == NULL) {
3878 Py_DECREF(ittuple);
3879 Py_DECREF(result);
3880 return NULL;
3881 }
3882 lz->ittuple = ittuple;
3883 lz->tuplesize = tuplesize;
3884 lz->numactive = tuplesize;
3885 lz->result = result;
3886 Py_INCREF(fillvalue);
3887 lz->fillvalue = fillvalue;
3888 return (PyObject *)lz;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003889}
3890
3891static void
3892izip_longest_dealloc(iziplongestobject *lz)
3893{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003894 PyObject_GC_UnTrack(lz);
3895 Py_XDECREF(lz->ittuple);
3896 Py_XDECREF(lz->result);
3897 Py_XDECREF(lz->fillvalue);
3898 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003899}
3900
3901static int
3902izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3903{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003904 Py_VISIT(lz->ittuple);
3905 Py_VISIT(lz->result);
3906 Py_VISIT(lz->fillvalue);
3907 return 0;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003908}
3909
3910static PyObject *
3911izip_longest_next(iziplongestobject *lz)
3912{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003913 Py_ssize_t i;
3914 Py_ssize_t tuplesize = lz->tuplesize;
3915 PyObject *result = lz->result;
3916 PyObject *it;
3917 PyObject *item;
3918 PyObject *olditem;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003919
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003920 if (tuplesize == 0)
3921 return NULL;
3922 if (lz->numactive == 0)
3923 return NULL;
3924 if (Py_REFCNT(result) == 1) {
3925 Py_INCREF(result);
3926 for (i=0 ; i < tuplesize ; i++) {
3927 it = PyTuple_GET_ITEM(lz->ittuple, i);
3928 if (it == NULL) {
3929 Py_INCREF(lz->fillvalue);
3930 item = lz->fillvalue;
3931 } else {
3932 item = PyIter_Next(it);
3933 if (item == NULL) {
3934 lz->numactive -= 1;
3935 if (lz->numactive == 0 || PyErr_Occurred()) {
3936 lz->numactive = 0;
3937 Py_DECREF(result);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003938 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003939 } else {
3940 Py_INCREF(lz->fillvalue);
3941 item = lz->fillvalue;
3942 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3943 Py_DECREF(it);
3944 }
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003945 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003946 }
3947 olditem = PyTuple_GET_ITEM(result, i);
3948 PyTuple_SET_ITEM(result, i, item);
3949 Py_DECREF(olditem);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003950 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003951 } else {
3952 result = PyTuple_New(tuplesize);
3953 if (result == NULL)
3954 return NULL;
3955 for (i=0 ; i < tuplesize ; i++) {
3956 it = PyTuple_GET_ITEM(lz->ittuple, i);
3957 if (it == NULL) {
3958 Py_INCREF(lz->fillvalue);
3959 item = lz->fillvalue;
3960 } else {
3961 item = PyIter_Next(it);
3962 if (item == NULL) {
3963 lz->numactive -= 1;
3964 if (lz->numactive == 0 || PyErr_Occurred()) {
3965 lz->numactive = 0;
3966 Py_DECREF(result);
3967 return NULL;
3968 } else {
3969 Py_INCREF(lz->fillvalue);
3970 item = lz->fillvalue;
3971 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3972 Py_DECREF(it);
3973 }
3974 }
3975 }
3976 PyTuple_SET_ITEM(result, i, item);
3977 }
3978 }
3979 return result;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003980}
3981
3982PyDoc_STRVAR(izip_longest_doc,
3983"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3984\n\
3985Return an izip_longest object whose .next() method returns a tuple where\n\
3986the i-th element comes from the i-th iterable argument. The .next()\n\
3987method continues until the longest iterable in the argument sequence\n\
3988is exhausted and then it raises StopIteration. When the shorter iterables\n\
3989are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3990defaults to None or can be specified by a keyword argument.\n\
3991");
3992
3993static PyTypeObject iziplongest_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003994 PyVarObject_HEAD_INIT(NULL, 0)
3995 "itertools.izip_longest", /* tp_name */
3996 sizeof(iziplongestobject), /* tp_basicsize */
3997 0, /* tp_itemsize */
3998 /* methods */
3999 (destructor)izip_longest_dealloc, /* tp_dealloc */
4000 0, /* tp_print */
4001 0, /* tp_getattr */
4002 0, /* tp_setattr */
4003 0, /* tp_compare */
4004 0, /* tp_repr */
4005 0, /* tp_as_number */
4006 0, /* tp_as_sequence */
4007 0, /* tp_as_mapping */
4008 0, /* tp_hash */
4009 0, /* tp_call */
4010 0, /* tp_str */
4011 PyObject_GenericGetAttr, /* tp_getattro */
4012 0, /* tp_setattro */
4013 0, /* tp_as_buffer */
4014 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4015 Py_TPFLAGS_BASETYPE, /* tp_flags */
4016 izip_longest_doc, /* tp_doc */
4017 (traverseproc)izip_longest_traverse, /* tp_traverse */
4018 0, /* tp_clear */
4019 0, /* tp_richcompare */
4020 0, /* tp_weaklistoffset */
4021 PyObject_SelfIter, /* tp_iter */
4022 (iternextfunc)izip_longest_next, /* tp_iternext */
4023 0, /* tp_methods */
4024 0, /* tp_members */
4025 0, /* tp_getset */
4026 0, /* tp_base */
4027 0, /* tp_dict */
4028 0, /* tp_descr_get */
4029 0, /* tp_descr_set */
4030 0, /* tp_dictoffset */
4031 0, /* tp_init */
4032 0, /* tp_alloc */
4033 izip_longest_new, /* tp_new */
4034 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd36862c2007-02-21 05:20:38 +00004035};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004036
4037/* module level code ********************************************************/
4038
4039PyDoc_STRVAR(module_doc,
4040"Functional tools for creating and using iterators.\n\
4041\n\
4042Infinite iterators:\n\
4043count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004044cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004045repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004046\n\
4047Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004048chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettingere76816b2009-01-29 03:43:44 +00004049compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004050dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4051groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00004052ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
4053ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004054islice(seq, [start,] stop [, step]) --> elements from\n\
4055 seq[start:stop:step]\n\
4056imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
4057starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004058tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004059takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004060izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4061izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4062\n\
4063Combinatoric generators:\n\
4064product(p, q, ... [repeat=1]) --> cartesian product\n\
4065permutations(p[, r])\n\
Raymond Hettinger9eac1192009-11-19 01:22:04 +00004066combinations(p, r)\n\
4067combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004068");
4069
4070
Raymond Hettingerad983e72003-11-12 14:32:26 +00004071static PyMethodDef module_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004072 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4073 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004074};
4075
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004076PyMODINIT_FUNC
4077inititertools(void)
4078{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004079 int i;
4080 PyObject *m;
4081 char *name;
4082 PyTypeObject *typelist[] = {
4083 &combinations_type,
4084 &cwr_type,
4085 &cycle_type,
4086 &dropwhile_type,
4087 &takewhile_type,
4088 &islice_type,
4089 &starmap_type,
4090 &imap_type,
4091 &chain_type,
4092 &compress_type,
4093 &ifilter_type,
4094 &ifilterfalse_type,
4095 &count_type,
4096 &izip_type,
4097 &iziplongest_type,
4098 &permutations_type,
4099 &product_type,
4100 &repeat_type,
4101 &groupby_type,
4102 NULL
4103 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004104
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004105 Py_TYPE(&teedataobject_type) = &PyType_Type;
4106 m = Py_InitModule3("itertools", module_methods, module_doc);
4107 if (m == NULL)
4108 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004109
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004110 for (i=0 ; typelist[i] != NULL ; i++) {
4111 if (PyType_Ready(typelist[i]) < 0)
4112 return;
4113 name = strchr(typelist[i]->tp_name, '.');
4114 assert (name != NULL);
4115 Py_INCREF(typelist[i]);
4116 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4117 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004118
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004119 if (PyType_Ready(&teedataobject_type) < 0)
4120 return;
4121 if (PyType_Ready(&tee_type) < 0)
4122 return;
4123 if (PyType_Ready(&_grouper_type) < 0)
4124 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004125}