blob: d5336f24294df40c8f08d63f22cd5ff02f5628eb [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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +000098 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000101
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000114 tmp = gbo->currkey;
115 gbo->currkey = newkey;
116 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000117
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000118 tmp = gbo->currvalue;
119 gbo->currvalue = newvalue;
120 Py_XDECREF(tmp);
121 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000131
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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_reserved */
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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000199 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000200
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000216 PyObject_GC_UnTrack(igo);
217 Py_DECREF(igo->parent);
218 Py_DECREF(igo->tgtkey);
219 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000220}
221
222static int
223_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
224{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000233 groupbyobject *gbo = (groupbyobject *)igo->parent;
234 PyObject *newvalue, *newkey, *r;
235 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000236
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000254 assert(gbo->currkey == NULL);
255 gbo->currkey = newkey;
256 gbo->currvalue = newvalue;
257 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000265 r = gbo->currvalue;
266 gbo->currvalue = NULL;
267 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000268
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000269 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000270}
271
272static PyTypeObject _grouper_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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_reserved */
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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000349 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000350
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000351 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
352 if (tdo == NULL)
353 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000354
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000375 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000376
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Wouters49fd7fa2006-04-21 10:40:58 +0000393static int
394teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
395{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Wouters49fd7fa2006-04-21 10:40:58 +0000402}
403
404static int
405teedataobject_clear(teedataobject *tdo)
406{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000407 int i;
408 Py_CLEAR(tdo->it);
409 for (i=0 ; i<tdo->numread ; i++)
410 Py_CLEAR(tdo->values[i]);
411 Py_CLEAR(tdo->nextlink);
412 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000413}
414
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000416teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000418 PyObject_GC_UnTrack(tdo);
419 teedataobject_clear(tdo);
420 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000421}
422
Raymond Hettingerad983e72003-11-12 14:32:26 +0000423PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000424
Raymond Hettingerad983e72003-11-12 14:32:26 +0000425static PyTypeObject teedataobject_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000426 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
427 "itertools.tee_dataobject", /* tp_name */
428 sizeof(teedataobject), /* tp_basicsize */
429 0, /* tp_itemsize */
430 /* methods */
431 (destructor)teedataobject_dealloc, /* tp_dealloc */
432 0, /* tp_print */
433 0, /* tp_getattr */
434 0, /* tp_setattr */
435 0, /* tp_reserved */
436 0, /* tp_repr */
437 0, /* tp_as_number */
438 0, /* tp_as_sequence */
439 0, /* tp_as_mapping */
440 0, /* tp_hash */
441 0, /* tp_call */
442 0, /* tp_str */
443 PyObject_GenericGetAttr, /* tp_getattro */
444 0, /* tp_setattro */
445 0, /* tp_as_buffer */
446 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
447 teedataobject_doc, /* tp_doc */
448 (traverseproc)teedataobject_traverse, /* tp_traverse */
449 (inquiry)teedataobject_clear, /* tp_clear */
450 0, /* tp_richcompare */
451 0, /* tp_weaklistoffset */
452 0, /* tp_iter */
453 0, /* tp_iternext */
454 0, /* tp_methods */
455 0, /* tp_members */
456 0, /* tp_getset */
457 0, /* tp_base */
458 0, /* tp_dict */
459 0, /* tp_descr_get */
460 0, /* tp_descr_set */
461 0, /* tp_dictoffset */
462 0, /* tp_init */
463 0, /* tp_alloc */
464 0, /* tp_new */
465 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000466};
467
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000468
469static PyTypeObject tee_type;
470
471static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000472tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000473{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000474 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000475
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000476 if (to->index >= LINKCELLS) {
477 link = teedataobject_jumplink(to->dataobj);
478 Py_DECREF(to->dataobj);
479 to->dataobj = (teedataobject *)link;
480 to->index = 0;
481 }
482 value = teedataobject_getitem(to->dataobj, to->index);
483 if (value == NULL)
484 return NULL;
485 to->index++;
486 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000487}
488
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000489static int
490tee_traverse(teeobject *to, visitproc visit, void *arg)
491{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000492 Py_VISIT((PyObject *)to->dataobj);
493 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000494}
495
Raymond Hettingerad983e72003-11-12 14:32:26 +0000496static PyObject *
497tee_copy(teeobject *to)
498{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000499 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000500
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000501 newto = PyObject_GC_New(teeobject, &tee_type);
502 if (newto == NULL)
503 return NULL;
504 Py_INCREF(to->dataobj);
505 newto->dataobj = to->dataobj;
506 newto->index = to->index;
507 newto->weakreflist = NULL;
508 PyObject_GC_Track(newto);
509 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000510}
511
512PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
513
514static PyObject *
515tee_fromiterable(PyObject *iterable)
516{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000517 teeobject *to;
518 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000519
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000520 it = PyObject_GetIter(iterable);
521 if (it == NULL)
522 return NULL;
523 if (PyObject_TypeCheck(it, &tee_type)) {
524 to = (teeobject *)tee_copy((teeobject *)it);
525 goto done;
526 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000527
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000528 to = PyObject_GC_New(teeobject, &tee_type);
529 if (to == NULL)
530 goto done;
531 to->dataobj = (teedataobject *)teedataobject_new(it);
532 if (!to->dataobj) {
533 PyObject_GC_Del(to);
534 to = NULL;
535 goto done;
536 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000537
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000538 to->index = 0;
539 to->weakreflist = NULL;
540 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000541done:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000542 Py_XDECREF(it);
543 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000544}
545
546static PyObject *
547tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
548{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000549 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000550
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000551 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
552 return NULL;
553 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000554}
555
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000556static int
557tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000558{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000559 if (to->weakreflist != NULL)
560 PyObject_ClearWeakRefs((PyObject *) to);
561 Py_CLEAR(to->dataobj);
562 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000563}
564
565static void
566tee_dealloc(teeobject *to)
567{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000568 PyObject_GC_UnTrack(to);
569 tee_clear(to);
570 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000571}
572
Raymond Hettingerad983e72003-11-12 14:32:26 +0000573PyDoc_STRVAR(teeobject_doc,
574"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000575
Raymond Hettingerad983e72003-11-12 14:32:26 +0000576static PyMethodDef tee_methods[] = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000577 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
578 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000579};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000580
581static PyTypeObject tee_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000582 PyVarObject_HEAD_INIT(NULL, 0)
583 "itertools.tee", /* tp_name */
584 sizeof(teeobject), /* tp_basicsize */
585 0, /* tp_itemsize */
586 /* methods */
587 (destructor)tee_dealloc, /* tp_dealloc */
588 0, /* tp_print */
589 0, /* tp_getattr */
590 0, /* tp_setattr */
591 0, /* tp_reserved */
592 0, /* tp_repr */
593 0, /* tp_as_number */
594 0, /* tp_as_sequence */
595 0, /* tp_as_mapping */
596 0, /* tp_hash */
597 0, /* tp_call */
598 0, /* tp_str */
599 0, /* tp_getattro */
600 0, /* tp_setattro */
601 0, /* tp_as_buffer */
602 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
603 teeobject_doc, /* tp_doc */
604 (traverseproc)tee_traverse, /* tp_traverse */
605 (inquiry)tee_clear, /* tp_clear */
606 0, /* tp_richcompare */
607 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
608 PyObject_SelfIter, /* tp_iter */
609 (iternextfunc)tee_next, /* tp_iternext */
610 tee_methods, /* tp_methods */
611 0, /* tp_members */
612 0, /* tp_getset */
613 0, /* tp_base */
614 0, /* tp_dict */
615 0, /* tp_descr_get */
616 0, /* tp_descr_set */
617 0, /* tp_dictoffset */
618 0, /* tp_init */
619 0, /* tp_alloc */
620 tee_new, /* tp_new */
621 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000622};
623
Raymond Hettingerad983e72003-11-12 14:32:26 +0000624static PyObject *
625tee(PyObject *self, PyObject *args)
626{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000627 Py_ssize_t i, n=2;
628 PyObject *it, *iterable, *copyable, *result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000629
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000630 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
631 return NULL;
632 if (n < 0) {
633 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
634 return NULL;
635 }
636 result = PyTuple_New(n);
637 if (result == NULL)
638 return NULL;
639 if (n == 0)
640 return result;
641 it = PyObject_GetIter(iterable);
642 if (it == NULL) {
643 Py_DECREF(result);
644 return NULL;
645 }
646 if (!PyObject_HasAttrString(it, "__copy__")) {
647 copyable = tee_fromiterable(it);
648 Py_DECREF(it);
649 if (copyable == NULL) {
650 Py_DECREF(result);
651 return NULL;
652 }
653 } else
654 copyable = it;
655 PyTuple_SET_ITEM(result, 0, copyable);
656 for (i=1 ; i<n ; i++) {
657 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
658 if (copyable == NULL) {
659 Py_DECREF(result);
660 return NULL;
661 }
662 PyTuple_SET_ITEM(result, i, copyable);
663 }
664 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000665}
666
667PyDoc_STRVAR(tee_doc,
668"tee(iterable, n=2) --> tuple of n independent iterators.");
669
670
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000671/* cycle object **********************************************************/
672
673typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000674 PyObject_HEAD
675 PyObject *it;
676 PyObject *saved;
677 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000678} cycleobject;
679
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000680static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000681
682static PyObject *
683cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
684{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000685 PyObject *it;
686 PyObject *iterable;
687 PyObject *saved;
688 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000689
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000690 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
691 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000692
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000693 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
694 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000695
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000696 /* Get iterator. */
697 it = PyObject_GetIter(iterable);
698 if (it == NULL)
699 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000700
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000701 saved = PyList_New(0);
702 if (saved == NULL) {
703 Py_DECREF(it);
704 return NULL;
705 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000706
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000707 /* create cycleobject structure */
708 lz = (cycleobject *)type->tp_alloc(type, 0);
709 if (lz == NULL) {
710 Py_DECREF(it);
711 Py_DECREF(saved);
712 return NULL;
713 }
714 lz->it = it;
715 lz->saved = saved;
716 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000717
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000718 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000719}
720
721static void
722cycle_dealloc(cycleobject *lz)
723{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000724 PyObject_GC_UnTrack(lz);
725 Py_XDECREF(lz->saved);
726 Py_XDECREF(lz->it);
727 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000728}
729
730static int
731cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
732{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000733 Py_VISIT(lz->it);
734 Py_VISIT(lz->saved);
735 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000736}
737
738static PyObject *
739cycle_next(cycleobject *lz)
740{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000741 PyObject *item;
742 PyObject *it;
743 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000744
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000745 while (1) {
746 item = PyIter_Next(lz->it);
747 if (item != NULL) {
748 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
749 Py_DECREF(item);
750 return NULL;
751 }
752 return item;
753 }
754 if (PyErr_Occurred()) {
755 if (PyErr_ExceptionMatches(PyExc_StopIteration))
756 PyErr_Clear();
757 else
758 return NULL;
759 }
760 if (PyList_Size(lz->saved) == 0)
761 return NULL;
762 it = PyObject_GetIter(lz->saved);
763 if (it == NULL)
764 return NULL;
765 tmp = lz->it;
766 lz->it = it;
767 lz->firstpass = 1;
768 Py_DECREF(tmp);
769 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000770}
771
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000772PyDoc_STRVAR(cycle_doc,
773"cycle(iterable) --> cycle object\n\
774\n\
775Return elements from the iterable until it is exhausted.\n\
776Then repeat the sequence indefinitely.");
777
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000778static PyTypeObject cycle_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000779 PyVarObject_HEAD_INIT(NULL, 0)
780 "itertools.cycle", /* tp_name */
781 sizeof(cycleobject), /* tp_basicsize */
782 0, /* tp_itemsize */
783 /* methods */
784 (destructor)cycle_dealloc, /* tp_dealloc */
785 0, /* tp_print */
786 0, /* tp_getattr */
787 0, /* tp_setattr */
788 0, /* tp_reserved */
789 0, /* tp_repr */
790 0, /* tp_as_number */
791 0, /* tp_as_sequence */
792 0, /* tp_as_mapping */
793 0, /* tp_hash */
794 0, /* tp_call */
795 0, /* tp_str */
796 PyObject_GenericGetAttr, /* tp_getattro */
797 0, /* tp_setattro */
798 0, /* tp_as_buffer */
799 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
800 Py_TPFLAGS_BASETYPE, /* tp_flags */
801 cycle_doc, /* tp_doc */
802 (traverseproc)cycle_traverse, /* tp_traverse */
803 0, /* tp_clear */
804 0, /* tp_richcompare */
805 0, /* tp_weaklistoffset */
806 PyObject_SelfIter, /* tp_iter */
807 (iternextfunc)cycle_next, /* tp_iternext */
808 0, /* tp_methods */
809 0, /* tp_members */
810 0, /* tp_getset */
811 0, /* tp_base */
812 0, /* tp_dict */
813 0, /* tp_descr_get */
814 0, /* tp_descr_set */
815 0, /* tp_dictoffset */
816 0, /* tp_init */
817 0, /* tp_alloc */
818 cycle_new, /* tp_new */
819 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000820};
821
822
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000823/* dropwhile object **********************************************************/
824
825typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000826 PyObject_HEAD
827 PyObject *func;
828 PyObject *it;
829 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000830} dropwhileobject;
831
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000832static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000833
834static PyObject *
835dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
836{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000837 PyObject *func, *seq;
838 PyObject *it;
839 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000840
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000841 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
842 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000843
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000844 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
845 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000846
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000847 /* Get iterator. */
848 it = PyObject_GetIter(seq);
849 if (it == NULL)
850 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000851
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000852 /* create dropwhileobject structure */
853 lz = (dropwhileobject *)type->tp_alloc(type, 0);
854 if (lz == NULL) {
855 Py_DECREF(it);
856 return NULL;
857 }
858 Py_INCREF(func);
859 lz->func = func;
860 lz->it = it;
861 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000862
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000863 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000864}
865
866static void
867dropwhile_dealloc(dropwhileobject *lz)
868{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000869 PyObject_GC_UnTrack(lz);
870 Py_XDECREF(lz->func);
871 Py_XDECREF(lz->it);
872 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000873}
874
875static int
876dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
877{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000878 Py_VISIT(lz->it);
879 Py_VISIT(lz->func);
880 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000881}
882
883static PyObject *
884dropwhile_next(dropwhileobject *lz)
885{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000886 PyObject *item, *good;
887 PyObject *it = lz->it;
888 long ok;
889 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000890
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000891 iternext = *Py_TYPE(it)->tp_iternext;
892 for (;;) {
893 item = iternext(it);
894 if (item == NULL)
895 return NULL;
896 if (lz->start == 1)
897 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000898
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000899 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
900 if (good == NULL) {
901 Py_DECREF(item);
902 return NULL;
903 }
904 ok = PyObject_IsTrue(good);
905 Py_DECREF(good);
906 if (!ok) {
907 lz->start = 1;
908 return item;
909 }
910 Py_DECREF(item);
911 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000912}
913
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000914PyDoc_STRVAR(dropwhile_doc,
915"dropwhile(predicate, iterable) --> dropwhile object\n\
916\n\
917Drop items from the iterable while predicate(item) is true.\n\
918Afterwards, return every element until the iterable is exhausted.");
919
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000920static PyTypeObject dropwhile_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000921 PyVarObject_HEAD_INIT(NULL, 0)
922 "itertools.dropwhile", /* tp_name */
923 sizeof(dropwhileobject), /* tp_basicsize */
924 0, /* tp_itemsize */
925 /* methods */
926 (destructor)dropwhile_dealloc, /* tp_dealloc */
927 0, /* tp_print */
928 0, /* tp_getattr */
929 0, /* tp_setattr */
930 0, /* tp_reserved */
931 0, /* tp_repr */
932 0, /* tp_as_number */
933 0, /* tp_as_sequence */
934 0, /* tp_as_mapping */
935 0, /* tp_hash */
936 0, /* tp_call */
937 0, /* tp_str */
938 PyObject_GenericGetAttr, /* tp_getattro */
939 0, /* tp_setattro */
940 0, /* tp_as_buffer */
941 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
942 Py_TPFLAGS_BASETYPE, /* tp_flags */
943 dropwhile_doc, /* tp_doc */
944 (traverseproc)dropwhile_traverse, /* tp_traverse */
945 0, /* tp_clear */
946 0, /* tp_richcompare */
947 0, /* tp_weaklistoffset */
948 PyObject_SelfIter, /* tp_iter */
949 (iternextfunc)dropwhile_next, /* tp_iternext */
950 0, /* tp_methods */
951 0, /* tp_members */
952 0, /* tp_getset */
953 0, /* tp_base */
954 0, /* tp_dict */
955 0, /* tp_descr_get */
956 0, /* tp_descr_set */
957 0, /* tp_dictoffset */
958 0, /* tp_init */
959 0, /* tp_alloc */
960 dropwhile_new, /* tp_new */
961 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000962};
963
964
965/* takewhile object **********************************************************/
966
967typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000968 PyObject_HEAD
969 PyObject *func;
970 PyObject *it;
971 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000972} takewhileobject;
973
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000974static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000975
976static PyObject *
977takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
978{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000979 PyObject *func, *seq;
980 PyObject *it;
981 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000982
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000983 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
984 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000985
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000986 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
987 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000988
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000989 /* Get iterator. */
990 it = PyObject_GetIter(seq);
991 if (it == NULL)
992 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000993
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000994 /* create takewhileobject structure */
995 lz = (takewhileobject *)type->tp_alloc(type, 0);
996 if (lz == NULL) {
997 Py_DECREF(it);
998 return NULL;
999 }
1000 Py_INCREF(func);
1001 lz->func = func;
1002 lz->it = it;
1003 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001004
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001005 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001006}
1007
1008static void
1009takewhile_dealloc(takewhileobject *lz)
1010{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001011 PyObject_GC_UnTrack(lz);
1012 Py_XDECREF(lz->func);
1013 Py_XDECREF(lz->it);
1014 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001015}
1016
1017static int
1018takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1019{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001020 Py_VISIT(lz->it);
1021 Py_VISIT(lz->func);
1022 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001023}
1024
1025static PyObject *
1026takewhile_next(takewhileobject *lz)
1027{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001028 PyObject *item, *good;
1029 PyObject *it = lz->it;
1030 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001031
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001032 if (lz->stop == 1)
1033 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001034
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001035 item = (*Py_TYPE(it)->tp_iternext)(it);
1036 if (item == NULL)
1037 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001038
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001039 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1040 if (good == NULL) {
1041 Py_DECREF(item);
1042 return NULL;
1043 }
1044 ok = PyObject_IsTrue(good);
1045 Py_DECREF(good);
1046 if (ok)
1047 return item;
1048 Py_DECREF(item);
1049 lz->stop = 1;
1050 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001051}
1052
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053PyDoc_STRVAR(takewhile_doc,
1054"takewhile(predicate, iterable) --> takewhile object\n\
1055\n\
1056Return successive entries from an iterable as long as the \n\
1057predicate evaluates to true for each entry.");
1058
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001059static PyTypeObject takewhile_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001060 PyVarObject_HEAD_INIT(NULL, 0)
1061 "itertools.takewhile", /* tp_name */
1062 sizeof(takewhileobject), /* tp_basicsize */
1063 0, /* tp_itemsize */
1064 /* methods */
1065 (destructor)takewhile_dealloc, /* tp_dealloc */
1066 0, /* tp_print */
1067 0, /* tp_getattr */
1068 0, /* tp_setattr */
1069 0, /* tp_reserved */
1070 0, /* tp_repr */
1071 0, /* tp_as_number */
1072 0, /* tp_as_sequence */
1073 0, /* tp_as_mapping */
1074 0, /* tp_hash */
1075 0, /* tp_call */
1076 0, /* tp_str */
1077 PyObject_GenericGetAttr, /* tp_getattro */
1078 0, /* tp_setattro */
1079 0, /* tp_as_buffer */
1080 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1081 Py_TPFLAGS_BASETYPE, /* tp_flags */
1082 takewhile_doc, /* tp_doc */
1083 (traverseproc)takewhile_traverse, /* tp_traverse */
1084 0, /* tp_clear */
1085 0, /* tp_richcompare */
1086 0, /* tp_weaklistoffset */
1087 PyObject_SelfIter, /* tp_iter */
1088 (iternextfunc)takewhile_next, /* tp_iternext */
1089 0, /* tp_methods */
1090 0, /* tp_members */
1091 0, /* tp_getset */
1092 0, /* tp_base */
1093 0, /* tp_dict */
1094 0, /* tp_descr_get */
1095 0, /* tp_descr_set */
1096 0, /* tp_dictoffset */
1097 0, /* tp_init */
1098 0, /* tp_alloc */
1099 takewhile_new, /* tp_new */
1100 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001101};
1102
1103
1104/* islice object ************************************************************/
1105
1106typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001107 PyObject_HEAD
1108 PyObject *it;
1109 Py_ssize_t next;
1110 Py_ssize_t stop;
1111 Py_ssize_t step;
1112 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001113} isliceobject;
1114
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001115static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116
1117static PyObject *
1118islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1119{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001120 PyObject *seq;
1121 Py_ssize_t start=0, stop=-1, step=1;
1122 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1123 Py_ssize_t numargs;
1124 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001125
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001126 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1127 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001128
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001129 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1130 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001131
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001132 numargs = PyTuple_Size(args);
1133 if (numargs == 2) {
1134 if (a1 != Py_None) {
1135 stop = PyLong_AsSsize_t(a1);
1136 if (stop == -1) {
1137 if (PyErr_Occurred())
1138 PyErr_Clear();
1139 PyErr_SetString(PyExc_ValueError,
1140 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1141 return NULL;
1142 }
1143 }
1144 } else {
1145 if (a1 != Py_None)
1146 start = PyLong_AsSsize_t(a1);
1147 if (start == -1 && PyErr_Occurred())
1148 PyErr_Clear();
1149 if (a2 != Py_None) {
1150 stop = PyLong_AsSsize_t(a2);
1151 if (stop == -1) {
1152 if (PyErr_Occurred())
1153 PyErr_Clear();
1154 PyErr_SetString(PyExc_ValueError,
1155 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1156 return NULL;
1157 }
1158 }
1159 }
1160 if (start<0 || stop<-1) {
1161 PyErr_SetString(PyExc_ValueError,
1162 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1163 return NULL;
1164 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001165
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001166 if (a3 != NULL) {
1167 if (a3 != Py_None)
1168 step = PyLong_AsSsize_t(a3);
1169 if (step == -1 && PyErr_Occurred())
1170 PyErr_Clear();
1171 }
1172 if (step<1) {
1173 PyErr_SetString(PyExc_ValueError,
1174 "Step for islice() must be a positive integer or None.");
1175 return NULL;
1176 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001177
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001178 /* Get iterator. */
1179 it = PyObject_GetIter(seq);
1180 if (it == NULL)
1181 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001182
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001183 /* create isliceobject structure */
1184 lz = (isliceobject *)type->tp_alloc(type, 0);
1185 if (lz == NULL) {
1186 Py_DECREF(it);
1187 return NULL;
1188 }
1189 lz->it = it;
1190 lz->next = start;
1191 lz->stop = stop;
1192 lz->step = step;
1193 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001194
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001195 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196}
1197
1198static void
1199islice_dealloc(isliceobject *lz)
1200{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001201 PyObject_GC_UnTrack(lz);
1202 Py_XDECREF(lz->it);
1203 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204}
1205
1206static int
1207islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1208{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001209 Py_VISIT(lz->it);
1210 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001211}
1212
1213static PyObject *
1214islice_next(isliceobject *lz)
1215{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001216 PyObject *item;
1217 PyObject *it = lz->it;
Raymond Hettinger101f09e2010-11-30 03:09:05 +00001218 Py_ssize_t stop = lz->stop;
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001219 Py_ssize_t oldnext;
1220 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001221
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001222 iternext = *Py_TYPE(it)->tp_iternext;
1223 while (lz->cnt < lz->next) {
1224 item = iternext(it);
1225 if (item == NULL)
1226 return NULL;
1227 Py_DECREF(item);
1228 lz->cnt++;
1229 }
Raymond Hettinger101f09e2010-11-30 03:09:05 +00001230 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001231 return NULL;
1232 item = iternext(it);
1233 if (item == NULL)
1234 return NULL;
1235 lz->cnt++;
1236 oldnext = lz->next;
1237 lz->next += lz->step;
Raymond Hettinger101f09e2010-11-30 03:09:05 +00001238 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1239 lz->next = stop;
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001240 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001241}
1242
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001243PyDoc_STRVAR(islice_doc,
1244"islice(iterable, [start,] stop [, step]) --> islice object\n\
1245\n\
1246Return an iterator whose next() method returns selected values from an\n\
1247iterable. If start is specified, will skip all preceding elements;\n\
1248otherwise, start defaults to zero. Step defaults to one. If\n\
1249specified as another value, step determines how many values are \n\
1250skipped between successive calls. Works like a slice() on a list\n\
1251but returns an iterator.");
1252
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001253static PyTypeObject islice_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001254 PyVarObject_HEAD_INIT(NULL, 0)
1255 "itertools.islice", /* tp_name */
1256 sizeof(isliceobject), /* tp_basicsize */
1257 0, /* tp_itemsize */
1258 /* methods */
1259 (destructor)islice_dealloc, /* tp_dealloc */
1260 0, /* tp_print */
1261 0, /* tp_getattr */
1262 0, /* tp_setattr */
1263 0, /* tp_reserved */
1264 0, /* tp_repr */
1265 0, /* tp_as_number */
1266 0, /* tp_as_sequence */
1267 0, /* tp_as_mapping */
1268 0, /* tp_hash */
1269 0, /* tp_call */
1270 0, /* tp_str */
1271 PyObject_GenericGetAttr, /* tp_getattro */
1272 0, /* tp_setattro */
1273 0, /* tp_as_buffer */
1274 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1275 Py_TPFLAGS_BASETYPE, /* tp_flags */
1276 islice_doc, /* tp_doc */
1277 (traverseproc)islice_traverse, /* tp_traverse */
1278 0, /* tp_clear */
1279 0, /* tp_richcompare */
1280 0, /* tp_weaklistoffset */
1281 PyObject_SelfIter, /* tp_iter */
1282 (iternextfunc)islice_next, /* tp_iternext */
1283 0, /* tp_methods */
1284 0, /* tp_members */
1285 0, /* tp_getset */
1286 0, /* tp_base */
1287 0, /* tp_dict */
1288 0, /* tp_descr_get */
1289 0, /* tp_descr_set */
1290 0, /* tp_dictoffset */
1291 0, /* tp_init */
1292 0, /* tp_alloc */
1293 islice_new, /* tp_new */
1294 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001295};
1296
1297
1298/* starmap object ************************************************************/
1299
1300typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001301 PyObject_HEAD
1302 PyObject *func;
1303 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001304} starmapobject;
1305
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001306static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001307
1308static PyObject *
1309starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1310{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001311 PyObject *func, *seq;
1312 PyObject *it;
1313 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001314
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001315 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1316 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001317
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001318 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1319 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001320
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001321 /* Get iterator. */
1322 it = PyObject_GetIter(seq);
1323 if (it == NULL)
1324 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001325
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001326 /* create starmapobject structure */
1327 lz = (starmapobject *)type->tp_alloc(type, 0);
1328 if (lz == NULL) {
1329 Py_DECREF(it);
1330 return NULL;
1331 }
1332 Py_INCREF(func);
1333 lz->func = func;
1334 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001335
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001336 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001337}
1338
1339static void
1340starmap_dealloc(starmapobject *lz)
1341{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001342 PyObject_GC_UnTrack(lz);
1343 Py_XDECREF(lz->func);
1344 Py_XDECREF(lz->it);
1345 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001346}
1347
1348static int
1349starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1350{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001351 Py_VISIT(lz->it);
1352 Py_VISIT(lz->func);
1353 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001354}
1355
1356static PyObject *
1357starmap_next(starmapobject *lz)
1358{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001359 PyObject *args;
1360 PyObject *result;
1361 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001362
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001363 args = (*Py_TYPE(it)->tp_iternext)(it);
1364 if (args == NULL)
1365 return NULL;
1366 if (!PyTuple_CheckExact(args)) {
1367 PyObject *newargs = PySequence_Tuple(args);
1368 Py_DECREF(args);
1369 if (newargs == NULL)
1370 return NULL;
1371 args = newargs;
1372 }
1373 result = PyObject_Call(lz->func, args, NULL);
1374 Py_DECREF(args);
1375 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001376}
1377
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001378PyDoc_STRVAR(starmap_doc,
1379"starmap(function, sequence) --> starmap object\n\
1380\n\
1381Return an iterator whose values are returned from the function evaluated\n\
1382with a argument tuple taken from the given sequence.");
1383
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001384static PyTypeObject starmap_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001385 PyVarObject_HEAD_INIT(NULL, 0)
1386 "itertools.starmap", /* tp_name */
1387 sizeof(starmapobject), /* tp_basicsize */
1388 0, /* tp_itemsize */
1389 /* methods */
1390 (destructor)starmap_dealloc, /* tp_dealloc */
1391 0, /* tp_print */
1392 0, /* tp_getattr */
1393 0, /* tp_setattr */
1394 0, /* tp_reserved */
1395 0, /* tp_repr */
1396 0, /* tp_as_number */
1397 0, /* tp_as_sequence */
1398 0, /* tp_as_mapping */
1399 0, /* tp_hash */
1400 0, /* tp_call */
1401 0, /* tp_str */
1402 PyObject_GenericGetAttr, /* tp_getattro */
1403 0, /* tp_setattro */
1404 0, /* tp_as_buffer */
1405 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1406 Py_TPFLAGS_BASETYPE, /* tp_flags */
1407 starmap_doc, /* tp_doc */
1408 (traverseproc)starmap_traverse, /* tp_traverse */
1409 0, /* tp_clear */
1410 0, /* tp_richcompare */
1411 0, /* tp_weaklistoffset */
1412 PyObject_SelfIter, /* tp_iter */
1413 (iternextfunc)starmap_next, /* tp_iternext */
1414 0, /* tp_methods */
1415 0, /* tp_members */
1416 0, /* tp_getset */
1417 0, /* tp_base */
1418 0, /* tp_dict */
1419 0, /* tp_descr_get */
1420 0, /* tp_descr_set */
1421 0, /* tp_dictoffset */
1422 0, /* tp_init */
1423 0, /* tp_alloc */
1424 starmap_new, /* tp_new */
1425 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001426};
1427
1428
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001429/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001430
1431typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001432 PyObject_HEAD
1433 PyObject *source; /* Iterator over input iterables */
1434 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001435} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001436
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001437static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001438
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001439static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001440chain_new_internal(PyTypeObject *type, PyObject *source)
1441{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001442 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001443
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001444 lz = (chainobject *)type->tp_alloc(type, 0);
1445 if (lz == NULL) {
1446 Py_DECREF(source);
1447 return NULL;
1448 }
1449
1450 lz->source = source;
1451 lz->active = NULL;
1452 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001453}
1454
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001455static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001456chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001458 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001460 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1461 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001462
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001463 source = PyObject_GetIter(args);
1464 if (source == NULL)
1465 return NULL;
1466
1467 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001468}
1469
1470static PyObject *
1471chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1472{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001473 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001474
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001475 source = PyObject_GetIter(arg);
1476 if (source == NULL)
1477 return NULL;
1478
1479 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001480}
1481
1482static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001483chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001484{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001485 PyObject_GC_UnTrack(lz);
1486 Py_XDECREF(lz->active);
1487 Py_XDECREF(lz->source);
1488 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001489}
1490
Raymond Hettinger2012f172003-02-07 05:32:58 +00001491static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001492chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001493{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001494 Py_VISIT(lz->source);
1495 Py_VISIT(lz->active);
1496 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001497}
1498
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001499static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001500chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001501{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001502 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001503
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001504 if (lz->source == NULL)
1505 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001506
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001507 if (lz->active == NULL) {
1508 PyObject *iterable = PyIter_Next(lz->source);
1509 if (iterable == NULL) {
1510 Py_CLEAR(lz->source);
1511 return NULL; /* no more input sources */
1512 }
1513 lz->active = PyObject_GetIter(iterable);
1514 Py_DECREF(iterable);
1515 if (lz->active == NULL) {
1516 Py_CLEAR(lz->source);
1517 return NULL; /* input not iterable */
1518 }
1519 }
1520 item = PyIter_Next(lz->active);
1521 if (item != NULL)
1522 return item;
1523 if (PyErr_Occurred()) {
1524 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1525 PyErr_Clear();
1526 else
1527 return NULL; /* input raised an exception */
1528 }
1529 Py_CLEAR(lz->active);
1530 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001531}
1532
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001533PyDoc_STRVAR(chain_doc,
1534"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001535\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001536Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001537first iterable until it is exhausted, then elements from the next\n\
1538iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001539
Christian Heimesf16baeb2008-02-29 14:57:44 +00001540PyDoc_STRVAR(chain_from_iterable_doc,
1541"chain.from_iterable(iterable) --> chain object\n\
1542\n\
1543Alternate chain() contructor taking a single iterable argument\n\
1544that evaluates lazily.");
1545
1546static PyMethodDef chain_methods[] = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001547 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1548 chain_from_iterable_doc},
1549 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001550};
1551
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001552static PyTypeObject chain_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001553 PyVarObject_HEAD_INIT(NULL, 0)
1554 "itertools.chain", /* tp_name */
1555 sizeof(chainobject), /* tp_basicsize */
1556 0, /* tp_itemsize */
1557 /* methods */
1558 (destructor)chain_dealloc, /* tp_dealloc */
1559 0, /* tp_print */
1560 0, /* tp_getattr */
1561 0, /* tp_setattr */
1562 0, /* tp_reserved */
1563 0, /* tp_repr */
1564 0, /* tp_as_number */
1565 0, /* tp_as_sequence */
1566 0, /* tp_as_mapping */
1567 0, /* tp_hash */
1568 0, /* tp_call */
1569 0, /* tp_str */
1570 PyObject_GenericGetAttr, /* tp_getattro */
1571 0, /* tp_setattro */
1572 0, /* tp_as_buffer */
1573 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1574 Py_TPFLAGS_BASETYPE, /* tp_flags */
1575 chain_doc, /* tp_doc */
1576 (traverseproc)chain_traverse, /* tp_traverse */
1577 0, /* tp_clear */
1578 0, /* tp_richcompare */
1579 0, /* tp_weaklistoffset */
1580 PyObject_SelfIter, /* tp_iter */
1581 (iternextfunc)chain_next, /* tp_iternext */
1582 chain_methods, /* tp_methods */
1583 0, /* tp_members */
1584 0, /* tp_getset */
1585 0, /* tp_base */
1586 0, /* tp_dict */
1587 0, /* tp_descr_get */
1588 0, /* tp_descr_set */
1589 0, /* tp_dictoffset */
1590 0, /* tp_init */
1591 0, /* tp_alloc */
1592 chain_new, /* tp_new */
1593 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001594};
1595
1596
Christian Heimesc3f30c42008-02-22 16:37:40 +00001597/* product object ************************************************************/
1598
1599typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001600 PyObject_HEAD
1601 PyObject *pools; /* tuple of pool tuples */
1602 Py_ssize_t *indices; /* one index per pool */
1603 PyObject *result; /* most recently returned result tuple */
1604 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001605} productobject;
1606
1607static PyTypeObject product_type;
1608
1609static PyObject *
1610product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1611{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001612 productobject *lz;
1613 Py_ssize_t nargs, npools, repeat=1;
1614 PyObject *pools = NULL;
1615 Py_ssize_t *indices = NULL;
1616 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001617
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001618 if (kwds != NULL) {
1619 char *kwlist[] = {"repeat", 0};
1620 PyObject *tmpargs = PyTuple_New(0);
1621 if (tmpargs == NULL)
1622 return NULL;
1623 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1624 Py_DECREF(tmpargs);
1625 return NULL;
1626 }
1627 Py_DECREF(tmpargs);
1628 if (repeat < 0) {
1629 PyErr_SetString(PyExc_ValueError,
1630 "repeat argument cannot be negative");
1631 return NULL;
1632 }
1633 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001634
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001635 assert(PyTuple_Check(args));
1636 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1637 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001638
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001639 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1640 if (indices == NULL) {
1641 PyErr_NoMemory();
1642 goto error;
1643 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001644
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001645 pools = PyTuple_New(npools);
1646 if (pools == NULL)
1647 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001648
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001649 for (i=0; i < nargs ; ++i) {
1650 PyObject *item = PyTuple_GET_ITEM(args, i);
1651 PyObject *pool = PySequence_Tuple(item);
1652 if (pool == NULL)
1653 goto error;
1654 PyTuple_SET_ITEM(pools, i, pool);
1655 indices[i] = 0;
1656 }
1657 for ( ; i < npools; ++i) {
1658 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1659 Py_INCREF(pool);
1660 PyTuple_SET_ITEM(pools, i, pool);
1661 indices[i] = 0;
1662 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001663
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001664 /* create productobject structure */
1665 lz = (productobject *)type->tp_alloc(type, 0);
1666 if (lz == NULL)
1667 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001668
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001669 lz->pools = pools;
1670 lz->indices = indices;
1671 lz->result = NULL;
1672 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001673
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001674 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001675
1676error:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001677 if (indices != NULL)
1678 PyMem_Free(indices);
1679 Py_XDECREF(pools);
1680 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001681}
1682
1683static void
1684product_dealloc(productobject *lz)
1685{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001686 PyObject_GC_UnTrack(lz);
1687 Py_XDECREF(lz->pools);
1688 Py_XDECREF(lz->result);
1689 if (lz->indices != NULL)
1690 PyMem_Free(lz->indices);
1691 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001692}
1693
1694static int
1695product_traverse(productobject *lz, visitproc visit, void *arg)
1696{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001697 Py_VISIT(lz->pools);
1698 Py_VISIT(lz->result);
1699 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001700}
1701
1702static PyObject *
1703product_next(productobject *lz)
1704{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001705 PyObject *pool;
1706 PyObject *elem;
1707 PyObject *oldelem;
1708 PyObject *pools = lz->pools;
1709 PyObject *result = lz->result;
1710 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1711 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001712
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001713 if (lz->stopped)
1714 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001715
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001716 if (result == NULL) {
1717 /* On the first pass, return an initial tuple filled with the
1718 first element from each pool. */
1719 result = PyTuple_New(npools);
1720 if (result == NULL)
1721 goto empty;
1722 lz->result = result;
1723 for (i=0; i < npools; i++) {
1724 pool = PyTuple_GET_ITEM(pools, i);
1725 if (PyTuple_GET_SIZE(pool) == 0)
1726 goto empty;
1727 elem = PyTuple_GET_ITEM(pool, 0);
1728 Py_INCREF(elem);
1729 PyTuple_SET_ITEM(result, i, elem);
1730 }
1731 } else {
1732 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001733
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001734 /* Copy the previous result tuple or re-use it if available */
1735 if (Py_REFCNT(result) > 1) {
1736 PyObject *old_result = result;
1737 result = PyTuple_New(npools);
1738 if (result == NULL)
1739 goto empty;
1740 lz->result = result;
1741 for (i=0; i < npools; i++) {
1742 elem = PyTuple_GET_ITEM(old_result, i);
1743 Py_INCREF(elem);
1744 PyTuple_SET_ITEM(result, i, elem);
1745 }
1746 Py_DECREF(old_result);
1747 }
1748 /* Now, we've got the only copy so we can update it in-place */
1749 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001750
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001751 /* Update the pool indices right-to-left. Only advance to the
1752 next pool when the previous one rolls-over */
1753 for (i=npools-1 ; i >= 0 ; i--) {
1754 pool = PyTuple_GET_ITEM(pools, i);
1755 indices[i]++;
1756 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1757 /* Roll-over and advance to next pool */
1758 indices[i] = 0;
1759 elem = PyTuple_GET_ITEM(pool, 0);
1760 Py_INCREF(elem);
1761 oldelem = PyTuple_GET_ITEM(result, i);
1762 PyTuple_SET_ITEM(result, i, elem);
1763 Py_DECREF(oldelem);
1764 } else {
1765 /* No rollover. Just increment and stop here. */
1766 elem = PyTuple_GET_ITEM(pool, indices[i]);
1767 Py_INCREF(elem);
1768 oldelem = PyTuple_GET_ITEM(result, i);
1769 PyTuple_SET_ITEM(result, i, elem);
1770 Py_DECREF(oldelem);
1771 break;
1772 }
1773 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001774
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001775 /* If i is negative, then the indices have all rolled-over
1776 and we're done. */
1777 if (i < 0)
1778 goto empty;
1779 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001780
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001781 Py_INCREF(result);
1782 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001783
1784empty:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001785 lz->stopped = 1;
1786 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001787}
1788
1789PyDoc_STRVAR(product_doc,
1790"product(*iterables) --> product object\n\
1791\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001792Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001793For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1794The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1795cycle in a manner similar to an odometer (with the rightmost element changing\n\
1796on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00001797To compute the product of an iterable with itself, specify the number\n\
1798of repetitions with the optional repeat keyword argument. For example,\n\
1799product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001800product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1801product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1802
1803static PyTypeObject product_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001804 PyVarObject_HEAD_INIT(NULL, 0)
1805 "itertools.product", /* tp_name */
1806 sizeof(productobject), /* tp_basicsize */
1807 0, /* tp_itemsize */
1808 /* methods */
1809 (destructor)product_dealloc, /* tp_dealloc */
1810 0, /* tp_print */
1811 0, /* tp_getattr */
1812 0, /* tp_setattr */
1813 0, /* tp_reserved */
1814 0, /* tp_repr */
1815 0, /* tp_as_number */
1816 0, /* tp_as_sequence */
1817 0, /* tp_as_mapping */
1818 0, /* tp_hash */
1819 0, /* tp_call */
1820 0, /* tp_str */
1821 PyObject_GenericGetAttr, /* tp_getattro */
1822 0, /* tp_setattro */
1823 0, /* tp_as_buffer */
1824 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1825 Py_TPFLAGS_BASETYPE, /* tp_flags */
1826 product_doc, /* tp_doc */
1827 (traverseproc)product_traverse, /* tp_traverse */
1828 0, /* tp_clear */
1829 0, /* tp_richcompare */
1830 0, /* tp_weaklistoffset */
1831 PyObject_SelfIter, /* tp_iter */
1832 (iternextfunc)product_next, /* tp_iternext */
1833 0, /* tp_methods */
1834 0, /* tp_members */
1835 0, /* tp_getset */
1836 0, /* tp_base */
1837 0, /* tp_dict */
1838 0, /* tp_descr_get */
1839 0, /* tp_descr_set */
1840 0, /* tp_dictoffset */
1841 0, /* tp_init */
1842 0, /* tp_alloc */
1843 product_new, /* tp_new */
1844 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001845};
1846
1847
Christian Heimes380f7f22008-02-28 11:19:05 +00001848/* combinations object ************************************************************/
1849
1850typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001851 PyObject_HEAD
1852 PyObject *pool; /* input converted to a tuple */
1853 Py_ssize_t *indices; /* one index per result element */
1854 PyObject *result; /* most recently returned result tuple */
1855 Py_ssize_t r; /* size of result tuple */
1856 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00001857} combinationsobject;
1858
1859static PyTypeObject combinations_type;
1860
1861static PyObject *
1862combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1863{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001864 combinationsobject *co;
1865 Py_ssize_t n;
1866 Py_ssize_t r;
1867 PyObject *pool = NULL;
1868 PyObject *iterable = NULL;
1869 Py_ssize_t *indices = NULL;
1870 Py_ssize_t i;
1871 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00001872
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001873 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
1874 &iterable, &r))
1875 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001876
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001877 pool = PySequence_Tuple(iterable);
1878 if (pool == NULL)
1879 goto error;
1880 n = PyTuple_GET_SIZE(pool);
1881 if (r < 0) {
1882 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
1883 goto error;
1884 }
Christian Heimes380f7f22008-02-28 11:19:05 +00001885
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001886 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
1887 if (indices == NULL) {
1888 PyErr_NoMemory();
1889 goto error;
1890 }
Christian Heimes380f7f22008-02-28 11:19:05 +00001891
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001892 for (i=0 ; i<r ; i++)
1893 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00001894
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001895 /* create combinationsobject structure */
1896 co = (combinationsobject *)type->tp_alloc(type, 0);
1897 if (co == NULL)
1898 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00001899
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001900 co->pool = pool;
1901 co->indices = indices;
1902 co->result = NULL;
1903 co->r = r;
1904 co->stopped = r > n ? 1 : 0;
1905
1906 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00001907
1908error:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001909 if (indices != NULL)
1910 PyMem_Free(indices);
1911 Py_XDECREF(pool);
1912 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001913}
1914
1915static void
1916combinations_dealloc(combinationsobject *co)
1917{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001918 PyObject_GC_UnTrack(co);
1919 Py_XDECREF(co->pool);
1920 Py_XDECREF(co->result);
1921 if (co->indices != NULL)
1922 PyMem_Free(co->indices);
1923 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00001924}
1925
1926static int
1927combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
1928{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001929 Py_VISIT(co->pool);
1930 Py_VISIT(co->result);
1931 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00001932}
1933
1934static PyObject *
1935combinations_next(combinationsobject *co)
1936{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001937 PyObject *elem;
1938 PyObject *oldelem;
1939 PyObject *pool = co->pool;
1940 Py_ssize_t *indices = co->indices;
1941 PyObject *result = co->result;
1942 Py_ssize_t n = PyTuple_GET_SIZE(pool);
1943 Py_ssize_t r = co->r;
1944 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00001945
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001946 if (co->stopped)
1947 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001948
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001949 if (result == NULL) {
1950 /* On the first pass, initialize result tuple using the indices */
1951 result = PyTuple_New(r);
1952 if (result == NULL)
1953 goto empty;
1954 co->result = result;
1955 for (i=0; i<r ; i++) {
1956 index = indices[i];
1957 elem = PyTuple_GET_ITEM(pool, index);
1958 Py_INCREF(elem);
1959 PyTuple_SET_ITEM(result, i, elem);
1960 }
1961 } else {
1962 /* Copy the previous result tuple or re-use it if available */
1963 if (Py_REFCNT(result) > 1) {
1964 PyObject *old_result = result;
1965 result = PyTuple_New(r);
1966 if (result == NULL)
1967 goto empty;
1968 co->result = result;
1969 for (i=0; i<r ; i++) {
1970 elem = PyTuple_GET_ITEM(old_result, i);
1971 Py_INCREF(elem);
1972 PyTuple_SET_ITEM(result, i, elem);
1973 }
1974 Py_DECREF(old_result);
1975 }
1976 /* Now, we've got the only copy so we can update it in-place
1977 * CPython's empty tuple is a singleton and cached in
1978 * PyTuple's freelist.
1979 */
1980 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00001981
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001982 /* Scan indices right-to-left until finding one that is not
1983 at its maximum (i + n - r). */
1984 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
1985 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00001986
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001987 /* If i is negative, then the indices are all at
1988 their maximum value and we're done. */
1989 if (i < 0)
1990 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00001991
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001992 /* Increment the current index which we know is not at its
1993 maximum. Then move back to the right setting each index
1994 to its lowest possible value (one higher than the index
1995 to its left -- this maintains the sort order invariant). */
1996 indices[i]++;
1997 for (j=i+1 ; j<r ; j++)
1998 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00001999
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002000 /* Update the result tuple for the new indices
2001 starting with i, the leftmost index that changed */
2002 for ( ; i<r ; i++) {
2003 index = indices[i];
2004 elem = PyTuple_GET_ITEM(pool, index);
2005 Py_INCREF(elem);
2006 oldelem = PyTuple_GET_ITEM(result, i);
2007 PyTuple_SET_ITEM(result, i, elem);
2008 Py_DECREF(oldelem);
2009 }
2010 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002011
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002012 Py_INCREF(result);
2013 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002014
2015empty:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002016 co->stopped = 1;
2017 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002018}
2019
2020PyDoc_STRVAR(combinations_doc,
Raymond Hettingerd3a77c02009-11-19 01:23:41 +00002021"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002022\n\
2023Return successive r-length combinations of elements in the iterable.\n\n\
2024combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2025
2026static PyTypeObject combinations_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002027 PyVarObject_HEAD_INIT(NULL, 0)
2028 "itertools.combinations", /* tp_name */
2029 sizeof(combinationsobject), /* tp_basicsize */
2030 0, /* tp_itemsize */
2031 /* methods */
2032 (destructor)combinations_dealloc, /* tp_dealloc */
2033 0, /* tp_print */
2034 0, /* tp_getattr */
2035 0, /* tp_setattr */
2036 0, /* tp_reserved */
2037 0, /* tp_repr */
2038 0, /* tp_as_number */
2039 0, /* tp_as_sequence */
2040 0, /* tp_as_mapping */
2041 0, /* tp_hash */
2042 0, /* tp_call */
2043 0, /* tp_str */
2044 PyObject_GenericGetAttr, /* tp_getattro */
2045 0, /* tp_setattro */
2046 0, /* tp_as_buffer */
2047 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2048 Py_TPFLAGS_BASETYPE, /* tp_flags */
2049 combinations_doc, /* tp_doc */
2050 (traverseproc)combinations_traverse, /* tp_traverse */
2051 0, /* tp_clear */
2052 0, /* tp_richcompare */
2053 0, /* tp_weaklistoffset */
2054 PyObject_SelfIter, /* tp_iter */
2055 (iternextfunc)combinations_next, /* tp_iternext */
2056 0, /* tp_methods */
2057 0, /* tp_members */
2058 0, /* tp_getset */
2059 0, /* tp_base */
2060 0, /* tp_dict */
2061 0, /* tp_descr_get */
2062 0, /* tp_descr_set */
2063 0, /* tp_dictoffset */
2064 0, /* tp_init */
2065 0, /* tp_alloc */
2066 combinations_new, /* tp_new */
2067 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002068};
2069
2070
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002071/* combinations with replacement object *******************************************/
2072
2073/* Equivalent to:
2074
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002075 def combinations_with_replacement(iterable, r):
2076 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2077 # number items returned: (n+r-1)! / r! / (n-1)!
2078 pool = tuple(iterable)
2079 n = len(pool)
2080 indices = [0] * r
2081 yield tuple(pool[i] for i in indices)
2082 while 1:
2083 for i in reversed(range(r)):
2084 if indices[i] != n - 1:
2085 break
2086 else:
2087 return
2088 indices[i:] = [indices[i] + 1] * (r - i)
2089 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002090
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002091 def combinations_with_replacement2(iterable, r):
2092 'Alternate version that filters from product()'
2093 pool = tuple(iterable)
2094 n = len(pool)
2095 for indices in product(range(n), repeat=r):
2096 if sorted(indices) == list(indices):
2097 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002098*/
2099typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002100 PyObject_HEAD
2101 PyObject *pool; /* input converted to a tuple */
2102 Py_ssize_t *indices; /* one index per result element */
2103 PyObject *result; /* most recently returned result tuple */
2104 Py_ssize_t r; /* size of result tuple */
2105 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002106} cwrobject;
2107
2108static PyTypeObject cwr_type;
2109
2110static PyObject *
2111cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2112{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002113 cwrobject *co;
2114 Py_ssize_t n;
2115 Py_ssize_t r;
2116 PyObject *pool = NULL;
2117 PyObject *iterable = NULL;
2118 Py_ssize_t *indices = NULL;
2119 Py_ssize_t i;
2120 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002121
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002122 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2123 &iterable, &r))
2124 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002125
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002126 pool = PySequence_Tuple(iterable);
2127 if (pool == NULL)
2128 goto error;
2129 n = PyTuple_GET_SIZE(pool);
2130 if (r < 0) {
2131 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2132 goto error;
2133 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002134
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002135 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2136 if (indices == NULL) {
2137 PyErr_NoMemory();
2138 goto error;
2139 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002140
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002141 for (i=0 ; i<r ; i++)
2142 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002143
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002144 /* create cwrobject structure */
2145 co = (cwrobject *)type->tp_alloc(type, 0);
2146 if (co == NULL)
2147 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002148
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002149 co->pool = pool;
2150 co->indices = indices;
2151 co->result = NULL;
2152 co->r = r;
2153 co->stopped = !n && r;
2154
2155 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002156
2157error:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002158 if (indices != NULL)
2159 PyMem_Free(indices);
2160 Py_XDECREF(pool);
2161 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002162}
2163
2164static void
2165cwr_dealloc(cwrobject *co)
2166{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002167 PyObject_GC_UnTrack(co);
2168 Py_XDECREF(co->pool);
2169 Py_XDECREF(co->result);
2170 if (co->indices != NULL)
2171 PyMem_Free(co->indices);
2172 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002173}
2174
2175static int
2176cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2177{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002178 Py_VISIT(co->pool);
2179 Py_VISIT(co->result);
2180 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002181}
2182
2183static PyObject *
2184cwr_next(cwrobject *co)
2185{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002186 PyObject *elem;
2187 PyObject *oldelem;
2188 PyObject *pool = co->pool;
2189 Py_ssize_t *indices = co->indices;
2190 PyObject *result = co->result;
2191 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2192 Py_ssize_t r = co->r;
2193 Py_ssize_t i, j, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002194
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002195 if (co->stopped)
2196 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002197
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002198 if (result == NULL) {
2199 /* On the first pass, initialize result tuple using the indices */
2200 result = PyTuple_New(r);
2201 if (result == NULL)
2202 goto empty;
2203 co->result = result;
2204 for (i=0; i<r ; i++) {
2205 index = indices[i];
2206 elem = PyTuple_GET_ITEM(pool, index);
2207 Py_INCREF(elem);
2208 PyTuple_SET_ITEM(result, i, elem);
2209 }
2210 } else {
2211 /* Copy the previous result tuple or re-use it if available */
2212 if (Py_REFCNT(result) > 1) {
2213 PyObject *old_result = result;
2214 result = PyTuple_New(r);
2215 if (result == NULL)
2216 goto empty;
2217 co->result = result;
2218 for (i=0; i<r ; i++) {
2219 elem = PyTuple_GET_ITEM(old_result, i);
2220 Py_INCREF(elem);
2221 PyTuple_SET_ITEM(result, i, elem);
2222 }
2223 Py_DECREF(old_result);
2224 }
2225 /* Now, we've got the only copy so we can update it in-place CPython's
2226 empty tuple is a singleton and cached in PyTuple's freelist. */
2227 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002228
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002229 /* Scan indices right-to-left until finding one that is not
2230 * at its maximum (n-1). */
2231 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2232 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002233
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002234 /* If i is negative, then the indices are all at
2235 their maximum value and we're done. */
2236 if (i < 0)
2237 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002238
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002239 /* Increment the current index which we know is not at its
2240 maximum. Then set all to the right to the same value. */
2241 indices[i]++;
2242 for (j=i+1 ; j<r ; j++)
2243 indices[j] = indices[j-1];
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002244
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002245 /* Update the result tuple for the new indices
2246 starting with i, the leftmost index that changed */
2247 for ( ; i<r ; i++) {
2248 index = indices[i];
2249 elem = PyTuple_GET_ITEM(pool, index);
2250 Py_INCREF(elem);
2251 oldelem = PyTuple_GET_ITEM(result, i);
2252 PyTuple_SET_ITEM(result, i, elem);
2253 Py_DECREF(oldelem);
2254 }
2255 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002256
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002257 Py_INCREF(result);
2258 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002259
2260empty:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002261 co->stopped = 1;
2262 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002263}
2264
2265PyDoc_STRVAR(cwr_doc,
Raymond Hettingerd3a77c02009-11-19 01:23:41 +00002266"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002267\n\
2268Return successive r-length combinations of elements in the iterable\n\
2269allowing individual elements to have successive repeats.\n\
2270combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2271
2272static PyTypeObject cwr_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002273 PyVarObject_HEAD_INIT(NULL, 0)
2274 "itertools.combinations_with_replacement", /* tp_name */
2275 sizeof(cwrobject), /* tp_basicsize */
2276 0, /* tp_itemsize */
2277 /* methods */
2278 (destructor)cwr_dealloc, /* tp_dealloc */
2279 0, /* tp_print */
2280 0, /* tp_getattr */
2281 0, /* tp_setattr */
2282 0, /* tp_reserved */
2283 0, /* tp_repr */
2284 0, /* tp_as_number */
2285 0, /* tp_as_sequence */
2286 0, /* tp_as_mapping */
2287 0, /* tp_hash */
2288 0, /* tp_call */
2289 0, /* tp_str */
2290 PyObject_GenericGetAttr, /* tp_getattro */
2291 0, /* tp_setattro */
2292 0, /* tp_as_buffer */
2293 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2294 Py_TPFLAGS_BASETYPE, /* tp_flags */
2295 cwr_doc, /* tp_doc */
2296 (traverseproc)cwr_traverse, /* tp_traverse */
2297 0, /* tp_clear */
2298 0, /* tp_richcompare */
2299 0, /* tp_weaklistoffset */
2300 PyObject_SelfIter, /* tp_iter */
2301 (iternextfunc)cwr_next, /* tp_iternext */
2302 0, /* tp_methods */
2303 0, /* tp_members */
2304 0, /* tp_getset */
2305 0, /* tp_base */
2306 0, /* tp_dict */
2307 0, /* tp_descr_get */
2308 0, /* tp_descr_set */
2309 0, /* tp_dictoffset */
2310 0, /* tp_init */
2311 0, /* tp_alloc */
2312 cwr_new, /* tp_new */
2313 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002314};
2315
2316
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002317/* permutations object ************************************************************
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002318
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002319def permutations(iterable, r=None):
2320 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2321 pool = tuple(iterable)
2322 n = len(pool)
2323 r = n if r is None else r
2324 indices = range(n)
2325 cycles = range(n-r+1, n+1)[::-1]
2326 yield tuple(pool[i] for i in indices[:r])
2327 while n:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002328 for i in reversed(range(r)):
2329 cycles[i] -= 1
2330 if cycles[i] == 0:
2331 indices[i:] = indices[i+1:] + indices[i:i+1]
2332 cycles[i] = n - i
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002333 else:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002334 j = cycles[i]
2335 indices[i], indices[-j] = indices[-j], indices[i]
2336 yield tuple(pool[i] for i in indices[:r])
2337 break
2338 else:
2339 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002340*/
2341
2342typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002343 PyObject_HEAD
2344 PyObject *pool; /* input converted to a tuple */
2345 Py_ssize_t *indices; /* one index per element in the pool */
2346 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2347 PyObject *result; /* most recently returned result tuple */
2348 Py_ssize_t r; /* size of result tuple */
2349 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002350} permutationsobject;
2351
2352static PyTypeObject permutations_type;
2353
2354static PyObject *
2355permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2356{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002357 permutationsobject *po;
2358 Py_ssize_t n;
2359 Py_ssize_t r;
2360 PyObject *robj = Py_None;
2361 PyObject *pool = NULL;
2362 PyObject *iterable = NULL;
2363 Py_ssize_t *indices = NULL;
2364 Py_ssize_t *cycles = NULL;
2365 Py_ssize_t i;
2366 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002367
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002368 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2369 &iterable, &robj))
2370 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002371
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002372 pool = PySequence_Tuple(iterable);
2373 if (pool == NULL)
2374 goto error;
2375 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002376
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002377 r = n;
2378 if (robj != Py_None) {
2379 if (!PyLong_Check(robj)) {
2380 PyErr_SetString(PyExc_TypeError, "Expected int as r");
2381 goto error;
2382 }
2383 r = PyLong_AsSsize_t(robj);
2384 if (r == -1 && PyErr_Occurred())
2385 goto error;
2386 }
2387 if (r < 0) {
2388 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2389 goto error;
2390 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002391
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002392 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2393 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2394 if (indices == NULL || cycles == NULL) {
2395 PyErr_NoMemory();
2396 goto error;
2397 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002398
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002399 for (i=0 ; i<n ; i++)
2400 indices[i] = i;
2401 for (i=0 ; i<r ; i++)
2402 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002403
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002404 /* create permutationsobject structure */
2405 po = (permutationsobject *)type->tp_alloc(type, 0);
2406 if (po == NULL)
2407 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002408
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002409 po->pool = pool;
2410 po->indices = indices;
2411 po->cycles = cycles;
2412 po->result = NULL;
2413 po->r = r;
2414 po->stopped = r > n ? 1 : 0;
2415
2416 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002417
2418error:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002419 if (indices != NULL)
2420 PyMem_Free(indices);
2421 if (cycles != NULL)
2422 PyMem_Free(cycles);
2423 Py_XDECREF(pool);
2424 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002425}
2426
2427static void
2428permutations_dealloc(permutationsobject *po)
2429{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002430 PyObject_GC_UnTrack(po);
2431 Py_XDECREF(po->pool);
2432 Py_XDECREF(po->result);
2433 PyMem_Free(po->indices);
2434 PyMem_Free(po->cycles);
2435 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002436}
2437
2438static int
2439permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2440{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002441 Py_VISIT(po->pool);
2442 Py_VISIT(po->result);
2443 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002444}
2445
2446static PyObject *
2447permutations_next(permutationsobject *po)
2448{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002449 PyObject *elem;
2450 PyObject *oldelem;
2451 PyObject *pool = po->pool;
2452 Py_ssize_t *indices = po->indices;
2453 Py_ssize_t *cycles = po->cycles;
2454 PyObject *result = po->result;
2455 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2456 Py_ssize_t r = po->r;
2457 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002458
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002459 if (po->stopped)
2460 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002461
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002462 if (result == NULL) {
2463 /* On the first pass, initialize result tuple using the indices */
2464 result = PyTuple_New(r);
2465 if (result == NULL)
2466 goto empty;
2467 po->result = result;
2468 for (i=0; i<r ; i++) {
2469 index = indices[i];
2470 elem = PyTuple_GET_ITEM(pool, index);
2471 Py_INCREF(elem);
2472 PyTuple_SET_ITEM(result, i, elem);
2473 }
2474 } else {
2475 if (n == 0)
2476 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002477
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002478 /* Copy the previous result tuple or re-use it if available */
2479 if (Py_REFCNT(result) > 1) {
2480 PyObject *old_result = result;
2481 result = PyTuple_New(r);
2482 if (result == NULL)
2483 goto empty;
2484 po->result = result;
2485 for (i=0; i<r ; i++) {
2486 elem = PyTuple_GET_ITEM(old_result, i);
2487 Py_INCREF(elem);
2488 PyTuple_SET_ITEM(result, i, elem);
2489 }
2490 Py_DECREF(old_result);
2491 }
2492 /* Now, we've got the only copy so we can update it in-place */
2493 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002494
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002495 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2496 for (i=r-1 ; i>=0 ; i--) {
2497 cycles[i] -= 1;
2498 if (cycles[i] == 0) {
2499 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2500 index = indices[i];
2501 for (j=i ; j<n-1 ; j++)
2502 indices[j] = indices[j+1];
2503 indices[n-1] = index;
2504 cycles[i] = n - i;
2505 } else {
2506 j = cycles[i];
2507 index = indices[i];
2508 indices[i] = indices[n-j];
2509 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002510
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002511 for (k=i; k<r ; k++) {
2512 /* start with i, the leftmost element that changed */
2513 /* yield tuple(pool[k] for k in indices[:r]) */
2514 index = indices[k];
2515 elem = PyTuple_GET_ITEM(pool, index);
2516 Py_INCREF(elem);
2517 oldelem = PyTuple_GET_ITEM(result, k);
2518 PyTuple_SET_ITEM(result, k, elem);
2519 Py_DECREF(oldelem);
2520 }
2521 break;
2522 }
2523 }
2524 /* If i is negative, then the cycles have all
2525 rolled-over and we're done. */
2526 if (i < 0)
2527 goto empty;
2528 }
2529 Py_INCREF(result);
2530 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002531
2532empty:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002533 po->stopped = 1;
2534 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002535}
2536
2537PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00002538"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002539\n\
2540Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00002541permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002542
2543static PyTypeObject permutations_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002544 PyVarObject_HEAD_INIT(NULL, 0)
2545 "itertools.permutations", /* tp_name */
2546 sizeof(permutationsobject), /* tp_basicsize */
2547 0, /* tp_itemsize */
2548 /* methods */
2549 (destructor)permutations_dealloc, /* tp_dealloc */
2550 0, /* tp_print */
2551 0, /* tp_getattr */
2552 0, /* tp_setattr */
2553 0, /* tp_reserved */
2554 0, /* tp_repr */
2555 0, /* tp_as_number */
2556 0, /* tp_as_sequence */
2557 0, /* tp_as_mapping */
2558 0, /* tp_hash */
2559 0, /* tp_call */
2560 0, /* tp_str */
2561 PyObject_GenericGetAttr, /* tp_getattro */
2562 0, /* tp_setattro */
2563 0, /* tp_as_buffer */
2564 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2565 Py_TPFLAGS_BASETYPE, /* tp_flags */
2566 permutations_doc, /* tp_doc */
2567 (traverseproc)permutations_traverse, /* tp_traverse */
2568 0, /* tp_clear */
2569 0, /* tp_richcompare */
2570 0, /* tp_weaklistoffset */
2571 PyObject_SelfIter, /* tp_iter */
2572 (iternextfunc)permutations_next, /* tp_iternext */
2573 0, /* tp_methods */
2574 0, /* tp_members */
2575 0, /* tp_getset */
2576 0, /* tp_base */
2577 0, /* tp_dict */
2578 0, /* tp_descr_get */
2579 0, /* tp_descr_set */
2580 0, /* tp_dictoffset */
2581 0, /* tp_init */
2582 0, /* tp_alloc */
2583 permutations_new, /* tp_new */
2584 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002585};
2586
2587
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002588/* compress object ************************************************************/
2589
2590/* Equivalent to:
2591
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002592 def compress(data, selectors):
2593 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2594 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002595*/
2596
2597typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002598 PyObject_HEAD
2599 PyObject *data;
2600 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002601} compressobject;
2602
2603static PyTypeObject compress_type;
2604
2605static PyObject *
2606compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2607{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002608 PyObject *seq1, *seq2;
2609 PyObject *data=NULL, *selectors=NULL;
2610 compressobject *lz;
2611 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002612
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002613 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2614 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002615
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002616 data = PyObject_GetIter(seq1);
2617 if (data == NULL)
2618 goto fail;
2619 selectors = PyObject_GetIter(seq2);
2620 if (selectors == NULL)
2621 goto fail;
2622
2623 /* create compressobject structure */
2624 lz = (compressobject *)type->tp_alloc(type, 0);
2625 if (lz == NULL)
2626 goto fail;
2627 lz->data = data;
2628 lz->selectors = selectors;
2629 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002630
2631fail:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002632 Py_XDECREF(data);
2633 Py_XDECREF(selectors);
2634 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002635}
2636
2637static void
2638compress_dealloc(compressobject *lz)
2639{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002640 PyObject_GC_UnTrack(lz);
2641 Py_XDECREF(lz->data);
2642 Py_XDECREF(lz->selectors);
2643 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002644}
2645
2646static int
2647compress_traverse(compressobject *lz, visitproc visit, void *arg)
2648{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002649 Py_VISIT(lz->data);
2650 Py_VISIT(lz->selectors);
2651 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002652}
2653
2654static PyObject *
2655compress_next(compressobject *lz)
2656{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002657 PyObject *data = lz->data, *selectors = lz->selectors;
2658 PyObject *datum, *selector;
2659 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2660 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2661 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002662
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002663 while (1) {
2664 /* Steps: get datum, get selector, evaluate selector.
2665 Order is important (to match the pure python version
2666 in terms of which input gets a chance to raise an
2667 exception first).
2668 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002669
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002670 datum = datanext(data);
2671 if (datum == NULL)
2672 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002673
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002674 selector = selectornext(selectors);
2675 if (selector == NULL) {
2676 Py_DECREF(datum);
2677 return NULL;
2678 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002679
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002680 ok = PyObject_IsTrue(selector);
2681 Py_DECREF(selector);
2682 if (ok == 1)
2683 return datum;
2684 Py_DECREF(datum);
2685 if (ok == -1)
2686 return NULL;
2687 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002688}
2689
2690PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00002691"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002692\n\
2693Return data elements corresponding to true selector elements.\n\
2694Forms a shorter iterator from selected data elements using the\n\
2695selectors to choose the data elements.");
2696
2697static PyTypeObject compress_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002698 PyVarObject_HEAD_INIT(NULL, 0)
2699 "itertools.compress", /* tp_name */
2700 sizeof(compressobject), /* tp_basicsize */
2701 0, /* tp_itemsize */
2702 /* methods */
2703 (destructor)compress_dealloc, /* tp_dealloc */
2704 0, /* tp_print */
2705 0, /* tp_getattr */
2706 0, /* tp_setattr */
2707 0, /* tp_reserved */
2708 0, /* tp_repr */
2709 0, /* tp_as_number */
2710 0, /* tp_as_sequence */
2711 0, /* tp_as_mapping */
2712 0, /* tp_hash */
2713 0, /* tp_call */
2714 0, /* tp_str */
2715 PyObject_GenericGetAttr, /* tp_getattro */
2716 0, /* tp_setattro */
2717 0, /* tp_as_buffer */
2718 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2719 Py_TPFLAGS_BASETYPE, /* tp_flags */
2720 compress_doc, /* tp_doc */
2721 (traverseproc)compress_traverse, /* tp_traverse */
2722 0, /* tp_clear */
2723 0, /* tp_richcompare */
2724 0, /* tp_weaklistoffset */
2725 PyObject_SelfIter, /* tp_iter */
2726 (iternextfunc)compress_next, /* tp_iternext */
2727 0, /* tp_methods */
2728 0, /* tp_members */
2729 0, /* tp_getset */
2730 0, /* tp_base */
2731 0, /* tp_dict */
2732 0, /* tp_descr_get */
2733 0, /* tp_descr_set */
2734 0, /* tp_dictoffset */
2735 0, /* tp_init */
2736 0, /* tp_alloc */
2737 compress_new, /* tp_new */
2738 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002739};
2740
2741
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002742/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00002743
2744typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002745 PyObject_HEAD
2746 PyObject *func;
2747 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002748} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002749
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002750static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002751
2752static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002753filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002754{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002755 PyObject *func, *seq;
2756 PyObject *it;
2757 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002758
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002759 if (type == &filterfalse_type &&
2760 !_PyArg_NoKeywords("filterfalse()", kwds))
2761 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002762
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002763 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
2764 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002765
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002766 /* Get iterator. */
2767 it = PyObject_GetIter(seq);
2768 if (it == NULL)
2769 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002770
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002771 /* create filterfalseobject structure */
2772 lz = (filterfalseobject *)type->tp_alloc(type, 0);
2773 if (lz == NULL) {
2774 Py_DECREF(it);
2775 return NULL;
2776 }
2777 Py_INCREF(func);
2778 lz->func = func;
2779 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002780
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002781 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002782}
2783
2784static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002785filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002786{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002787 PyObject_GC_UnTrack(lz);
2788 Py_XDECREF(lz->func);
2789 Py_XDECREF(lz->it);
2790 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002791}
2792
2793static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002794filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002795{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002796 Py_VISIT(lz->it);
2797 Py_VISIT(lz->func);
2798 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002799}
2800
2801static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002802filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002803{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002804 PyObject *item;
2805 PyObject *it = lz->it;
2806 long ok;
2807 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002808
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002809 iternext = *Py_TYPE(it)->tp_iternext;
2810 for (;;) {
2811 item = iternext(it);
2812 if (item == NULL)
2813 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002814
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002815 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2816 ok = PyObject_IsTrue(item);
2817 } else {
2818 PyObject *good;
2819 good = PyObject_CallFunctionObjArgs(lz->func,
2820 item, NULL);
2821 if (good == NULL) {
2822 Py_DECREF(item);
2823 return NULL;
2824 }
2825 ok = PyObject_IsTrue(good);
2826 Py_DECREF(good);
2827 }
2828 if (!ok)
2829 return item;
2830 Py_DECREF(item);
2831 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00002832}
2833
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002834PyDoc_STRVAR(filterfalse_doc,
2835"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002836\n\
2837Return those items of sequence for which function(item) is false.\n\
2838If function is None, return the items that are false.");
2839
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002840static PyTypeObject filterfalse_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002841 PyVarObject_HEAD_INIT(NULL, 0)
2842 "itertools.filterfalse", /* tp_name */
2843 sizeof(filterfalseobject), /* tp_basicsize */
2844 0, /* tp_itemsize */
2845 /* methods */
2846 (destructor)filterfalse_dealloc, /* tp_dealloc */
2847 0, /* tp_print */
2848 0, /* tp_getattr */
2849 0, /* tp_setattr */
2850 0, /* tp_reserved */
2851 0, /* tp_repr */
2852 0, /* tp_as_number */
2853 0, /* tp_as_sequence */
2854 0, /* tp_as_mapping */
2855 0, /* tp_hash */
2856 0, /* tp_call */
2857 0, /* tp_str */
2858 PyObject_GenericGetAttr, /* tp_getattro */
2859 0, /* tp_setattro */
2860 0, /* tp_as_buffer */
2861 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2862 Py_TPFLAGS_BASETYPE, /* tp_flags */
2863 filterfalse_doc, /* tp_doc */
2864 (traverseproc)filterfalse_traverse, /* tp_traverse */
2865 0, /* tp_clear */
2866 0, /* tp_richcompare */
2867 0, /* tp_weaklistoffset */
2868 PyObject_SelfIter, /* tp_iter */
2869 (iternextfunc)filterfalse_next, /* tp_iternext */
2870 0, /* tp_methods */
2871 0, /* tp_members */
2872 0, /* tp_getset */
2873 0, /* tp_base */
2874 0, /* tp_dict */
2875 0, /* tp_descr_get */
2876 0, /* tp_descr_set */
2877 0, /* tp_dictoffset */
2878 0, /* tp_init */
2879 0, /* tp_alloc */
2880 filterfalse_new, /* tp_new */
2881 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002882};
2883
2884
2885/* count object ************************************************************/
2886
2887typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002888 PyObject_HEAD
2889 Py_ssize_t cnt;
2890 PyObject *long_cnt;
2891 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002892} countobject;
2893
Raymond Hettinger30729212009-02-12 06:28:27 +00002894/* Counting logic and invariants:
2895
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002896fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00002897
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002898 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
2899 Advances with: cnt += 1
2900 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00002901
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002902slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00002903
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002904 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
2905 All counting is done with python objects (no overflows or underflows).
2906 Advances with: long_cnt += long_step
2907 Step may be zero -- effectively a slow version of repeat(cnt).
2908 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00002909*/
2910
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002911static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002912
2913static PyObject *
2914count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2915{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002916 countobject *lz;
2917 int slow_mode = 0;
2918 Py_ssize_t cnt = 0;
2919 PyObject *long_cnt = NULL;
2920 PyObject *long_step = NULL;
2921 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002922
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002923 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
2924 kwlist, &long_cnt, &long_step))
2925 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002926
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002927 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
2928 (long_step != NULL && !PyNumber_Check(long_step))) {
2929 PyErr_SetString(PyExc_TypeError, "a number is required");
2930 return NULL;
2931 }
Raymond Hettinger30729212009-02-12 06:28:27 +00002932
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002933 if (long_cnt != NULL) {
2934 cnt = PyLong_AsSsize_t(long_cnt);
2935 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
2936 PyErr_Clear();
2937 slow_mode = 1;
2938 }
2939 Py_INCREF(long_cnt);
2940 } else {
2941 cnt = 0;
2942 long_cnt = PyLong_FromLong(0);
2943 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002944
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002945 /* If not specified, step defaults to 1 */
2946 if (long_step == NULL) {
2947 long_step = PyLong_FromLong(1);
2948 if (long_step == NULL) {
2949 Py_DECREF(long_cnt);
2950 return NULL;
2951 }
2952 } else
2953 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002954
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002955 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002956
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002957 /* Fast mode only works when the step is 1 */
2958 if (!PyLong_Check(long_step) ||
2959 PyLong_AS_LONG(long_step) != 1) {
2960 slow_mode = 1;
2961 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002962
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002963 if (slow_mode)
2964 cnt = PY_SSIZE_T_MAX;
2965 else
2966 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002967
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002968 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
2969 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
2970 assert(slow_mode ||
2971 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002972
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002973 /* create countobject structure */
2974 lz = (countobject *)type->tp_alloc(type, 0);
2975 if (lz == NULL) {
2976 Py_XDECREF(long_cnt);
2977 return NULL;
2978 }
2979 lz->cnt = cnt;
2980 lz->long_cnt = long_cnt;
2981 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002982
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002983 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002984}
2985
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002986static void
2987count_dealloc(countobject *lz)
2988{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002989 PyObject_GC_UnTrack(lz);
2990 Py_XDECREF(lz->long_cnt);
2991 Py_XDECREF(lz->long_step);
2992 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00002993}
2994
Benjamin Peterson25e26a02009-02-16 21:28:29 +00002995static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00002996count_traverse(countobject *lz, visitproc visit, void *arg)
2997{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002998 Py_VISIT(lz->long_cnt);
2999 Py_VISIT(lz->long_step);
3000 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003001}
3002
3003static PyObject *
3004count_nextlong(countobject *lz)
3005{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003006 PyObject *long_cnt;
3007 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003008
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003009 long_cnt = lz->long_cnt;
3010 if (long_cnt == NULL) {
3011 /* Switch to slow_mode */
3012 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3013 if (long_cnt == NULL)
3014 return NULL;
3015 }
3016 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00003017
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003018 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3019 if (stepped_up == NULL)
3020 return NULL;
3021 lz->long_cnt = stepped_up;
3022 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003023}
3024
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003025static PyObject *
3026count_next(countobject *lz)
3027{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003028 if (lz->cnt == PY_SSIZE_T_MAX)
3029 return count_nextlong(lz);
3030 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003031}
3032
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003033static PyObject *
3034count_repr(countobject *lz)
3035{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003036 if (lz->cnt != PY_SSIZE_T_MAX)
3037 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003038
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003039 if (PyLong_Check(lz->long_step)) {
3040 long step = PyLong_AsLong(lz->long_step);
3041 if (step == -1 && PyErr_Occurred()) {
3042 PyErr_Clear();
3043 }
3044 if (step == 1) {
3045 /* Don't display step when it is an integer equal to 1 */
3046 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
3047 }
3048 }
3049 return PyUnicode_FromFormat("count(%R, %R)",
3050 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003051}
3052
Raymond Hettinger6c8ee7a2009-11-30 21:55:17 +00003053static PyObject *
3054count_reduce(countobject *lz)
3055{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003056 if (lz->cnt == PY_SSIZE_T_MAX)
3057 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3058 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettinger6c8ee7a2009-11-30 21:55:17 +00003059}
3060
3061PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3062
3063static PyMethodDef count_methods[] = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003064 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3065 count_reduce_doc},
3066 {NULL, NULL} /* sentinel */
Raymond Hettinger6c8ee7a2009-11-30 21:55:17 +00003067};
3068
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003069PyDoc_STRVAR(count_doc,
Ezio Melotti7a61e3c2010-06-11 02:28:37 +00003070 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003071\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003072Return a count object whose .__next__() method returns consecutive values.\n\
3073Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00003074 def count(firstval=0, step=1):\n\
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003075 x = firstval\n\
3076 while 1:\n\
3077 yield x\n\
3078 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003079
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003080static PyTypeObject count_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003081 PyVarObject_HEAD_INIT(NULL, 0)
3082 "itertools.count", /* tp_name */
3083 sizeof(countobject), /* tp_basicsize */
3084 0, /* tp_itemsize */
3085 /* methods */
3086 (destructor)count_dealloc, /* tp_dealloc */
3087 0, /* tp_print */
3088 0, /* tp_getattr */
3089 0, /* tp_setattr */
3090 0, /* tp_reserved */
3091 (reprfunc)count_repr, /* tp_repr */
3092 0, /* tp_as_number */
3093 0, /* tp_as_sequence */
3094 0, /* tp_as_mapping */
3095 0, /* tp_hash */
3096 0, /* tp_call */
3097 0, /* tp_str */
3098 PyObject_GenericGetAttr, /* tp_getattro */
3099 0, /* tp_setattro */
3100 0, /* tp_as_buffer */
3101 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3102 Py_TPFLAGS_BASETYPE, /* tp_flags */
3103 count_doc, /* tp_doc */
3104 (traverseproc)count_traverse, /* tp_traverse */
3105 0, /* tp_clear */
3106 0, /* tp_richcompare */
3107 0, /* tp_weaklistoffset */
3108 PyObject_SelfIter, /* tp_iter */
3109 (iternextfunc)count_next, /* tp_iternext */
3110 count_methods, /* tp_methods */
3111 0, /* tp_members */
3112 0, /* tp_getset */
3113 0, /* tp_base */
3114 0, /* tp_dict */
3115 0, /* tp_descr_get */
3116 0, /* tp_descr_set */
3117 0, /* tp_dictoffset */
3118 0, /* tp_init */
3119 0, /* tp_alloc */
3120 count_new, /* tp_new */
3121 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003122};
3123
3124
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003125/* repeat object ************************************************************/
3126
3127typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003128 PyObject_HEAD
3129 PyObject *element;
3130 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003131} repeatobject;
3132
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003133static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003134
3135static PyObject *
3136repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3137{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003138 repeatobject *ro;
3139 PyObject *element;
3140 Py_ssize_t cnt = -1;
3141 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003142
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003143 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3144 &element, &cnt))
3145 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003146
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003147 if (PyTuple_Size(args) == 2 && cnt < 0)
3148 cnt = 0;
3149
3150 ro = (repeatobject *)type->tp_alloc(type, 0);
3151 if (ro == NULL)
3152 return NULL;
3153 Py_INCREF(element);
3154 ro->element = element;
3155 ro->cnt = cnt;
3156 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003157}
3158
3159static void
3160repeat_dealloc(repeatobject *ro)
3161{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003162 PyObject_GC_UnTrack(ro);
3163 Py_XDECREF(ro->element);
3164 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003165}
3166
3167static int
3168repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3169{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003170 Py_VISIT(ro->element);
3171 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003172}
3173
3174static PyObject *
3175repeat_next(repeatobject *ro)
3176{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003177 if (ro->cnt == 0)
3178 return NULL;
3179 if (ro->cnt > 0)
3180 ro->cnt--;
3181 Py_INCREF(ro->element);
3182 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003183}
3184
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003185static PyObject *
3186repeat_repr(repeatobject *ro)
3187{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003188 if (ro->cnt == -1)
3189 return PyUnicode_FromFormat("repeat(%R)", ro->element);
3190 else
3191 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
3192}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003193
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003194static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003195repeat_len(repeatobject *ro)
3196{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003197 if (ro->cnt == -1) {
3198 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3199 return NULL;
3200 }
3201 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003202}
3203
Armin Rigof5b3e362006-02-11 21:32:43 +00003204PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003205
3206static PyMethodDef repeat_methods[] = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003207 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3208 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003209};
3210
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003211PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00003212"repeat(object [,times]) -> create an iterator which returns the object\n\
3213for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003214endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003215
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003216static PyTypeObject repeat_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003217 PyVarObject_HEAD_INIT(NULL, 0)
3218 "itertools.repeat", /* tp_name */
3219 sizeof(repeatobject), /* tp_basicsize */
3220 0, /* tp_itemsize */
3221 /* methods */
3222 (destructor)repeat_dealloc, /* tp_dealloc */
3223 0, /* tp_print */
3224 0, /* tp_getattr */
3225 0, /* tp_setattr */
3226 0, /* tp_reserved */
3227 (reprfunc)repeat_repr, /* tp_repr */
3228 0, /* tp_as_number */
3229 0, /* tp_as_sequence */
3230 0, /* tp_as_mapping */
3231 0, /* tp_hash */
3232 0, /* tp_call */
3233 0, /* tp_str */
3234 PyObject_GenericGetAttr, /* tp_getattro */
3235 0, /* tp_setattro */
3236 0, /* tp_as_buffer */
3237 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3238 Py_TPFLAGS_BASETYPE, /* tp_flags */
3239 repeat_doc, /* tp_doc */
3240 (traverseproc)repeat_traverse, /* tp_traverse */
3241 0, /* tp_clear */
3242 0, /* tp_richcompare */
3243 0, /* tp_weaklistoffset */
3244 PyObject_SelfIter, /* tp_iter */
3245 (iternextfunc)repeat_next, /* tp_iternext */
3246 repeat_methods, /* tp_methods */
3247 0, /* tp_members */
3248 0, /* tp_getset */
3249 0, /* tp_base */
3250 0, /* tp_dict */
3251 0, /* tp_descr_get */
3252 0, /* tp_descr_set */
3253 0, /* tp_dictoffset */
3254 0, /* tp_init */
3255 0, /* tp_alloc */
3256 repeat_new, /* tp_new */
3257 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003258};
3259
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003260/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00003261
3262#include "Python.h"
3263
3264typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003265 PyObject_HEAD
3266 Py_ssize_t tuplesize;
3267 Py_ssize_t numactive;
3268 PyObject *ittuple; /* tuple of iterators */
3269 PyObject *result;
3270 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003271} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003272
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003273static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003274
3275static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003276zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003277{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003278 ziplongestobject *lz;
3279 Py_ssize_t i;
3280 PyObject *ittuple; /* tuple of iterators */
3281 PyObject *result;
3282 PyObject *fillvalue = Py_None;
3283 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003284
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003285 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3286 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3287 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3288 PyErr_SetString(PyExc_TypeError,
3289 "zip_longest() got an unexpected keyword argument");
3290 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003291 }
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003292 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003293
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003294 /* args must be a tuple */
3295 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00003296
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003297 /* obtain iterators */
3298 ittuple = PyTuple_New(tuplesize);
3299 if (ittuple == NULL)
3300 return NULL;
3301 for (i=0; i < tuplesize; ++i) {
3302 PyObject *item = PyTuple_GET_ITEM(args, i);
3303 PyObject *it = PyObject_GetIter(item);
3304 if (it == NULL) {
3305 if (PyErr_ExceptionMatches(PyExc_TypeError))
3306 PyErr_Format(PyExc_TypeError,
3307 "zip_longest argument #%zd must support iteration",
3308 i+1);
3309 Py_DECREF(ittuple);
3310 return NULL;
3311 }
3312 PyTuple_SET_ITEM(ittuple, i, it);
3313 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003314
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003315 /* create a result holder */
3316 result = PyTuple_New(tuplesize);
3317 if (result == NULL) {
3318 Py_DECREF(ittuple);
3319 return NULL;
3320 }
3321 for (i=0 ; i < tuplesize ; i++) {
3322 Py_INCREF(Py_None);
3323 PyTuple_SET_ITEM(result, i, Py_None);
3324 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003325
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003326 /* create ziplongestobject structure */
3327 lz = (ziplongestobject *)type->tp_alloc(type, 0);
3328 if (lz == NULL) {
3329 Py_DECREF(ittuple);
3330 Py_DECREF(result);
3331 return NULL;
3332 }
3333 lz->ittuple = ittuple;
3334 lz->tuplesize = tuplesize;
3335 lz->numactive = tuplesize;
3336 lz->result = result;
3337 Py_INCREF(fillvalue);
3338 lz->fillvalue = fillvalue;
3339 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003340}
3341
3342static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003343zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003344{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003345 PyObject_GC_UnTrack(lz);
3346 Py_XDECREF(lz->ittuple);
3347 Py_XDECREF(lz->result);
3348 Py_XDECREF(lz->fillvalue);
3349 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003350}
3351
3352static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003353zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003354{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003355 Py_VISIT(lz->ittuple);
3356 Py_VISIT(lz->result);
3357 Py_VISIT(lz->fillvalue);
3358 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003359}
3360
3361static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003362zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003363{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003364 Py_ssize_t i;
3365 Py_ssize_t tuplesize = lz->tuplesize;
3366 PyObject *result = lz->result;
3367 PyObject *it;
3368 PyObject *item;
3369 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003370
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003371 if (tuplesize == 0)
3372 return NULL;
3373 if (lz->numactive == 0)
3374 return NULL;
3375 if (Py_REFCNT(result) == 1) {
3376 Py_INCREF(result);
3377 for (i=0 ; i < tuplesize ; i++) {
3378 it = PyTuple_GET_ITEM(lz->ittuple, i);
3379 if (it == NULL) {
3380 Py_INCREF(lz->fillvalue);
3381 item = lz->fillvalue;
3382 } else {
3383 item = PyIter_Next(it);
3384 if (item == NULL) {
3385 lz->numactive -= 1;
3386 if (lz->numactive == 0 || PyErr_Occurred()) {
3387 lz->numactive = 0;
3388 Py_DECREF(result);
Raymond Hettingera9311a32009-11-01 21:02:38 +00003389 return NULL;
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003390 } else {
3391 Py_INCREF(lz->fillvalue);
3392 item = lz->fillvalue;
3393 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3394 Py_DECREF(it);
3395 }
Raymond Hettingera9311a32009-11-01 21:02:38 +00003396 }
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003397 }
3398 olditem = PyTuple_GET_ITEM(result, i);
3399 PyTuple_SET_ITEM(result, i, item);
3400 Py_DECREF(olditem);
Raymond Hettingera9311a32009-11-01 21:02:38 +00003401 }
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003402 } else {
3403 result = PyTuple_New(tuplesize);
3404 if (result == NULL)
3405 return NULL;
3406 for (i=0 ; i < tuplesize ; i++) {
3407 it = PyTuple_GET_ITEM(lz->ittuple, i);
3408 if (it == NULL) {
3409 Py_INCREF(lz->fillvalue);
3410 item = lz->fillvalue;
3411 } else {
3412 item = PyIter_Next(it);
3413 if (item == NULL) {
3414 lz->numactive -= 1;
3415 if (lz->numactive == 0 || PyErr_Occurred()) {
3416 lz->numactive = 0;
3417 Py_DECREF(result);
3418 return NULL;
3419 } else {
3420 Py_INCREF(lz->fillvalue);
3421 item = lz->fillvalue;
3422 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3423 Py_DECREF(it);
3424 }
3425 }
3426 }
3427 PyTuple_SET_ITEM(result, i, item);
3428 }
3429 }
3430 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003431}
3432
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003433PyDoc_STRVAR(zip_longest_doc,
3434"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003435\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003436Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00003437the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003438method continues until the longest iterable in the argument sequence\n\
3439is exhausted and then it raises StopIteration. When the shorter iterables\n\
3440are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3441defaults to None or can be specified by a keyword argument.\n\
3442");
3443
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003444static PyTypeObject ziplongest_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003445 PyVarObject_HEAD_INIT(NULL, 0)
3446 "itertools.zip_longest", /* tp_name */
3447 sizeof(ziplongestobject), /* tp_basicsize */
3448 0, /* tp_itemsize */
3449 /* methods */
3450 (destructor)zip_longest_dealloc, /* tp_dealloc */
3451 0, /* tp_print */
3452 0, /* tp_getattr */
3453 0, /* tp_setattr */
3454 0, /* tp_reserved */
3455 0, /* tp_repr */
3456 0, /* tp_as_number */
3457 0, /* tp_as_sequence */
3458 0, /* tp_as_mapping */
3459 0, /* tp_hash */
3460 0, /* tp_call */
3461 0, /* tp_str */
3462 PyObject_GenericGetAttr, /* tp_getattro */
3463 0, /* tp_setattro */
3464 0, /* tp_as_buffer */
3465 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3466 Py_TPFLAGS_BASETYPE, /* tp_flags */
3467 zip_longest_doc, /* tp_doc */
3468 (traverseproc)zip_longest_traverse, /* tp_traverse */
3469 0, /* tp_clear */
3470 0, /* tp_richcompare */
3471 0, /* tp_weaklistoffset */
3472 PyObject_SelfIter, /* tp_iter */
3473 (iternextfunc)zip_longest_next, /* tp_iternext */
3474 0, /* tp_methods */
3475 0, /* tp_members */
3476 0, /* tp_getset */
3477 0, /* tp_base */
3478 0, /* tp_dict */
3479 0, /* tp_descr_get */
3480 0, /* tp_descr_set */
3481 0, /* tp_dictoffset */
3482 0, /* tp_init */
3483 0, /* tp_alloc */
3484 zip_longest_new, /* tp_new */
3485 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003486};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003487
3488/* module level code ********************************************************/
3489
3490PyDoc_STRVAR(module_doc,
3491"Functional tools for creating and using iterators.\n\
3492\n\
3493Infinite iterators:\n\
3494count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003495cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003496repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003497\n\
3498Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003499chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
3500compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
3501dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
3502groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003503filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003504islice(seq, [start,] stop [, step]) --> elements from\n\
3505 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003506starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003507tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003508takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003509zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00003510\n\
3511Combinatoric generators:\n\
3512product(p, q, ... [repeat=1]) --> cartesian product\n\
3513permutations(p[, r])\n\
Raymond Hettingerd3a77c02009-11-19 01:23:41 +00003514combinations(p, r)\n\
3515combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003516");
3517
3518
Raymond Hettingerad983e72003-11-12 14:32:26 +00003519static PyMethodDef module_methods[] = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003520 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3521 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00003522};
3523
Martin v. Löwis1a214512008-06-11 05:26:20 +00003524
3525static struct PyModuleDef itertoolsmodule = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003526 PyModuleDef_HEAD_INIT,
3527 "itertools",
3528 module_doc,
3529 -1,
3530 module_methods,
3531 NULL,
3532 NULL,
3533 NULL,
3534 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00003535};
3536
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003537PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00003538PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003539{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003540 int i;
3541 PyObject *m;
3542 char *name;
3543 PyTypeObject *typelist[] = {
3544 &combinations_type,
3545 &cwr_type,
3546 &cycle_type,
3547 &dropwhile_type,
3548 &takewhile_type,
3549 &islice_type,
3550 &starmap_type,
3551 &chain_type,
3552 &compress_type,
3553 &filterfalse_type,
3554 &count_type,
3555 &ziplongest_type,
3556 &permutations_type,
3557 &product_type,
3558 &repeat_type,
3559 &groupby_type,
3560 NULL
3561 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00003562
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003563 Py_TYPE(&teedataobject_type) = &PyType_Type;
3564 m = PyModule_Create(&itertoolsmodule);
3565 if (m == NULL)
3566 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003567
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003568 for (i=0 ; typelist[i] != NULL ; i++) {
3569 if (PyType_Ready(typelist[i]) < 0)
3570 return NULL;
3571 name = strchr(typelist[i]->tp_name, '.');
3572 assert (name != NULL);
3573 Py_INCREF(typelist[i]);
3574 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
3575 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003576
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003577 if (PyType_Ready(&teedataobject_type) < 0)
3578 return NULL;
3579 if (PyType_Ready(&tee_type) < 0)
3580 return NULL;
3581 if (PyType_Ready(&_grouper_type) < 0)
3582 return NULL;
3583 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003584}