blob: 04bfffc5b0db51095248ec9093ca005f5ea3a89e [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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +000098 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000101
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000114 tmp = gbo->currkey;
115 gbo->currkey = newkey;
116 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000117
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000118 tmp = gbo->currvalue;
119 gbo->currvalue = newvalue;
120 Py_XDECREF(tmp);
121 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000131
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000199 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000200
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000233 groupbyobject *gbo = (groupbyobject *)igo->parent;
234 PyObject *newvalue, *newkey, *r;
235 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000236
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000254 assert(gbo->currkey == NULL);
255 gbo->currkey = newkey;
256 gbo->currvalue = newvalue;
257 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000265 r = gbo->currvalue;
266 gbo->currvalue = NULL;
267 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000268
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000269 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000270}
271
272static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000349 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000350
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000351 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
352 if (tdo == NULL)
353 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000354
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000375 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000376
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000474 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000475
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000499 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000500
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000517 teeobject *to;
518 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000519
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000538 to->index = 0;
539 to->weakreflist = NULL;
540 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000541done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000549 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000550
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000627 Py_ssize_t i, n=2;
628 PyObject *it, *iterable, *copyable, *result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000629
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000685 PyObject *it;
686 PyObject *iterable;
687 PyObject *saved;
688 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000689
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000690 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
691 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000692
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000693 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
694 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000695
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000696 /* Get iterator. */
697 it = PyObject_GetIter(iterable);
698 if (it == NULL)
699 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000700
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000718 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000719}
720
721static void
722cycle_dealloc(cycleobject *lz)
723{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000741 PyObject *item;
742 PyObject *it;
743 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000744
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000837 PyObject *func, *seq;
838 PyObject *it;
839 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000840
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000841 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
842 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000843
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000844 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
845 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000846
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000847 /* Get iterator. */
848 it = PyObject_GetIter(seq);
849 if (it == NULL)
850 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000851
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000863 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000864}
865
866static void
867dropwhile_dealloc(dropwhileobject *lz)
868{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +0000979 PyObject *func, *seq;
980 PyObject *it;
981 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000982
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000983 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
984 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000985
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000986 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
987 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000988
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000989 /* Get iterator. */
990 it = PyObject_GetIter(seq);
991 if (it == NULL)
992 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000993
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001005 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001006}
1007
1008static void
1009takewhile_dealloc(takewhileobject *lz)
1010{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001028 PyObject *item, *good;
1029 PyObject *it = lz->it;
1030 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001031
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001032 if (lz->stop == 1)
1033 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001034
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001035 item = (*Py_TYPE(it)->tp_iternext)(it);
1036 if (item == NULL)
1037 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001038
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001126 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1127 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001128
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001129 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1130 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001131
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001178 /* Get iterator. */
1179 it = PyObject_GetIter(seq);
1180 if (it == NULL)
1181 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001182
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001195 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196}
1197
1198static void
1199islice_dealloc(isliceobject *lz)
1200{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001216 PyObject *item;
1217 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001218 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001219 Py_ssize_t oldnext;
1220 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001221
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Hettinger69b34bf2010-11-30 02:49:29 +00001230 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Hettinger69b34bf2010-11-30 02:49:29 +00001238 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1239 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001311 PyObject *func, *seq;
1312 PyObject *it;
1313 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001314
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001315 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1316 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001317
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001318 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1319 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001320
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001321 /* Get iterator. */
1322 it = PyObject_GetIter(seq);
1323 if (it == NULL)
1324 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001325
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001336 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001337}
1338
1339static void
1340starmap_dealloc(starmapobject *lz)
1341{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001359 PyObject *args;
1360 PyObject *result;
1361 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001362
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001439static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001440chain_new_internal(PyTypeObject *type, PyObject *source)
1441{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001442 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001443
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001458 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001460 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1461 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001462
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001473 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001474
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001502 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001503
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001504 if (lz->source == NULL)
1505 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001506
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001645 pools = PyTuple_New(npools);
1646 if (pools == NULL)
1647 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001648
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001674 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001675
1676error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001713 if (lz->stopped)
1714 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001715
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001781 Py_INCREF(result);
1782 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001783
1784empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001873 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
1874 &iterable, &r))
1875 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001876
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001892 for (i=0 ; i<r ; i++)
1893 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00001894
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00001946 if (co->stopped)
1947 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001948
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00002012 Py_INCREF(result);
2013 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002014
2015empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002016 co->stopped = 1;
2017 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002018}
2019
2020PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00002141 for (i=0 ; i<r ; i++)
2142 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002143
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00002195 if (co->stopped)
2196 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002197
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00002257 Py_INCREF(result);
2258 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002259
2260empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002261 co->stopped = 1;
2262 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002263}
2264
2265PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +00002459 if (po->stopped)
2460 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002461
Antoine Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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 Pitrouf95a1b32010-05-09 15:52:27 +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
Raymond Hettinger482ba772010-12-01 22:48:00 +00002587/* accumulate object ************************************************************/
2588
2589typedef struct {
2590 PyObject_HEAD
2591 PyObject *total;
2592 PyObject *it;
2593} accumulateobject;
2594
2595static PyTypeObject accumulate_type;
2596
2597static PyObject *
2598accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2599{
2600 static char *kwargs[] = {"iterable", "start", NULL};
2601 PyObject *iterable;
2602 PyObject *it;
2603 PyObject *start = NULL;
2604 accumulateobject *lz;
2605
2606 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
2607 kwargs, &iterable, &start))
2608 return NULL;
2609
2610 /* Get iterator. */
2611 it = PyObject_GetIter(iterable);
2612 if (it == NULL)
2613 return NULL;
2614
2615 /* Default start value */
2616 if (start == NULL) {
2617 start = PyLong_FromLong(0);
2618 if (start == NULL) {
2619 Py_DECREF(it);
2620 return NULL;
2621 }
2622 } else {
2623 Py_INCREF(start);
2624 }
2625
2626 /* create accumulateobject structure */
2627 lz = (accumulateobject *)type->tp_alloc(type, 0);
2628 if (lz == NULL) {
2629 Py_DECREF(it);
2630 Py_DECREF(start);
2631 return NULL;
2632 }
2633
2634 lz->total = start;
2635 lz->it = it;
2636 return (PyObject *)lz;
2637}
2638
2639static void
2640accumulate_dealloc(accumulateobject *lz)
2641{
2642 PyObject_GC_UnTrack(lz);
2643 Py_XDECREF(lz->total);
2644 Py_XDECREF(lz->it);
2645 Py_TYPE(lz)->tp_free(lz);
2646}
2647
2648static int
2649accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
2650{
2651 Py_VISIT(lz->it);
2652 Py_VISIT(lz->total);
2653 return 0;
2654}
2655
2656static PyObject *
2657accumulate_next(accumulateobject *lz)
2658{
2659 PyObject *val, *oldtotal, *newtotal;
2660
2661 val = PyIter_Next(lz->it);
2662 if (val == NULL)
2663 return NULL;
2664
2665 newtotal = PyNumber_Add(lz->total, val);
2666 Py_DECREF(val);
2667 if (newtotal == NULL)
2668 return NULL;
2669
2670 oldtotal = lz->total;
2671 lz->total = newtotal;
2672 Py_DECREF(oldtotal);
2673
2674 Py_INCREF(newtotal);
2675 return newtotal;
2676}
2677
2678PyDoc_STRVAR(accumulate_doc,
2679"accumulate(iterable, start=0) --> accumulate object\n\
2680\n\
2681Return series of accumulated sums.");
2682
2683static PyTypeObject accumulate_type = {
2684 PyVarObject_HEAD_INIT(NULL, 0)
2685 "itertools.accumulate", /* tp_name */
2686 sizeof(accumulateobject), /* tp_basicsize */
2687 0, /* tp_itemsize */
2688 /* methods */
2689 (destructor)accumulate_dealloc, /* tp_dealloc */
2690 0, /* tp_print */
2691 0, /* tp_getattr */
2692 0, /* tp_setattr */
2693 0, /* tp_reserved */
2694 0, /* tp_repr */
2695 0, /* tp_as_number */
2696 0, /* tp_as_sequence */
2697 0, /* tp_as_mapping */
2698 0, /* tp_hash */
2699 0, /* tp_call */
2700 0, /* tp_str */
2701 PyObject_GenericGetAttr, /* tp_getattro */
2702 0, /* tp_setattro */
2703 0, /* tp_as_buffer */
2704 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2705 Py_TPFLAGS_BASETYPE, /* tp_flags */
2706 accumulate_doc, /* tp_doc */
2707 (traverseproc)accumulate_traverse, /* tp_traverse */
2708 0, /* tp_clear */
2709 0, /* tp_richcompare */
2710 0, /* tp_weaklistoffset */
2711 PyObject_SelfIter, /* tp_iter */
2712 (iternextfunc)accumulate_next, /* tp_iternext */
2713 0, /* tp_methods */
2714 0, /* tp_members */
2715 0, /* tp_getset */
2716 0, /* tp_base */
2717 0, /* tp_dict */
2718 0, /* tp_descr_get */
2719 0, /* tp_descr_set */
2720 0, /* tp_dictoffset */
2721 0, /* tp_init */
2722 0, /* tp_alloc */
2723 accumulate_new, /* tp_new */
2724 PyObject_GC_Del, /* tp_free */
2725};
2726
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002727
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002728/* compress object ************************************************************/
2729
2730/* Equivalent to:
2731
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002732 def compress(data, selectors):
2733 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2734 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002735*/
2736
2737typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002738 PyObject_HEAD
2739 PyObject *data;
2740 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002741} compressobject;
2742
2743static PyTypeObject compress_type;
2744
2745static PyObject *
2746compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2747{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002748 PyObject *seq1, *seq2;
2749 PyObject *data=NULL, *selectors=NULL;
2750 compressobject *lz;
2751 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002752
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002753 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2754 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002755
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002756 data = PyObject_GetIter(seq1);
2757 if (data == NULL)
2758 goto fail;
2759 selectors = PyObject_GetIter(seq2);
2760 if (selectors == NULL)
2761 goto fail;
2762
2763 /* create compressobject structure */
2764 lz = (compressobject *)type->tp_alloc(type, 0);
2765 if (lz == NULL)
2766 goto fail;
2767 lz->data = data;
2768 lz->selectors = selectors;
2769 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002770
2771fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002772 Py_XDECREF(data);
2773 Py_XDECREF(selectors);
2774 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002775}
2776
2777static void
2778compress_dealloc(compressobject *lz)
2779{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002780 PyObject_GC_UnTrack(lz);
2781 Py_XDECREF(lz->data);
2782 Py_XDECREF(lz->selectors);
2783 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002784}
2785
2786static int
2787compress_traverse(compressobject *lz, visitproc visit, void *arg)
2788{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002789 Py_VISIT(lz->data);
2790 Py_VISIT(lz->selectors);
2791 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002792}
2793
2794static PyObject *
2795compress_next(compressobject *lz)
2796{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002797 PyObject *data = lz->data, *selectors = lz->selectors;
2798 PyObject *datum, *selector;
2799 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2800 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2801 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002802
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002803 while (1) {
2804 /* Steps: get datum, get selector, evaluate selector.
2805 Order is important (to match the pure python version
2806 in terms of which input gets a chance to raise an
2807 exception first).
2808 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002809
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002810 datum = datanext(data);
2811 if (datum == NULL)
2812 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002813
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002814 selector = selectornext(selectors);
2815 if (selector == NULL) {
2816 Py_DECREF(datum);
2817 return NULL;
2818 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002819
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002820 ok = PyObject_IsTrue(selector);
2821 Py_DECREF(selector);
2822 if (ok == 1)
2823 return datum;
2824 Py_DECREF(datum);
2825 if (ok == -1)
2826 return NULL;
2827 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002828}
2829
2830PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00002831"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002832\n\
2833Return data elements corresponding to true selector elements.\n\
2834Forms a shorter iterator from selected data elements using the\n\
2835selectors to choose the data elements.");
2836
2837static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002838 PyVarObject_HEAD_INIT(NULL, 0)
2839 "itertools.compress", /* tp_name */
2840 sizeof(compressobject), /* tp_basicsize */
2841 0, /* tp_itemsize */
2842 /* methods */
2843 (destructor)compress_dealloc, /* tp_dealloc */
2844 0, /* tp_print */
2845 0, /* tp_getattr */
2846 0, /* tp_setattr */
2847 0, /* tp_reserved */
2848 0, /* tp_repr */
2849 0, /* tp_as_number */
2850 0, /* tp_as_sequence */
2851 0, /* tp_as_mapping */
2852 0, /* tp_hash */
2853 0, /* tp_call */
2854 0, /* tp_str */
2855 PyObject_GenericGetAttr, /* tp_getattro */
2856 0, /* tp_setattro */
2857 0, /* tp_as_buffer */
2858 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2859 Py_TPFLAGS_BASETYPE, /* tp_flags */
2860 compress_doc, /* tp_doc */
2861 (traverseproc)compress_traverse, /* tp_traverse */
2862 0, /* tp_clear */
2863 0, /* tp_richcompare */
2864 0, /* tp_weaklistoffset */
2865 PyObject_SelfIter, /* tp_iter */
2866 (iternextfunc)compress_next, /* tp_iternext */
2867 0, /* tp_methods */
2868 0, /* tp_members */
2869 0, /* tp_getset */
2870 0, /* tp_base */
2871 0, /* tp_dict */
2872 0, /* tp_descr_get */
2873 0, /* tp_descr_set */
2874 0, /* tp_dictoffset */
2875 0, /* tp_init */
2876 0, /* tp_alloc */
2877 compress_new, /* tp_new */
2878 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002879};
2880
2881
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002882/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00002883
2884typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002885 PyObject_HEAD
2886 PyObject *func;
2887 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002888} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002889
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002890static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002891
2892static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002893filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002894{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002895 PyObject *func, *seq;
2896 PyObject *it;
2897 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002898
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002899 if (type == &filterfalse_type &&
2900 !_PyArg_NoKeywords("filterfalse()", kwds))
2901 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002902
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002903 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
2904 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002905
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002906 /* Get iterator. */
2907 it = PyObject_GetIter(seq);
2908 if (it == NULL)
2909 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002910
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002911 /* create filterfalseobject structure */
2912 lz = (filterfalseobject *)type->tp_alloc(type, 0);
2913 if (lz == NULL) {
2914 Py_DECREF(it);
2915 return NULL;
2916 }
2917 Py_INCREF(func);
2918 lz->func = func;
2919 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002920
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002921 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002922}
2923
2924static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002925filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002926{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002927 PyObject_GC_UnTrack(lz);
2928 Py_XDECREF(lz->func);
2929 Py_XDECREF(lz->it);
2930 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002931}
2932
2933static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002934filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002935{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002936 Py_VISIT(lz->it);
2937 Py_VISIT(lz->func);
2938 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002939}
2940
2941static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002942filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002943{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002944 PyObject *item;
2945 PyObject *it = lz->it;
2946 long ok;
2947 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002948
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002949 iternext = *Py_TYPE(it)->tp_iternext;
2950 for (;;) {
2951 item = iternext(it);
2952 if (item == NULL)
2953 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002954
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002955 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2956 ok = PyObject_IsTrue(item);
2957 } else {
2958 PyObject *good;
2959 good = PyObject_CallFunctionObjArgs(lz->func,
2960 item, NULL);
2961 if (good == NULL) {
2962 Py_DECREF(item);
2963 return NULL;
2964 }
2965 ok = PyObject_IsTrue(good);
2966 Py_DECREF(good);
2967 }
2968 if (!ok)
2969 return item;
2970 Py_DECREF(item);
2971 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00002972}
2973
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002974PyDoc_STRVAR(filterfalse_doc,
2975"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002976\n\
2977Return those items of sequence for which function(item) is false.\n\
2978If function is None, return the items that are false.");
2979
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002980static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002981 PyVarObject_HEAD_INIT(NULL, 0)
2982 "itertools.filterfalse", /* tp_name */
2983 sizeof(filterfalseobject), /* tp_basicsize */
2984 0, /* tp_itemsize */
2985 /* methods */
2986 (destructor)filterfalse_dealloc, /* tp_dealloc */
2987 0, /* tp_print */
2988 0, /* tp_getattr */
2989 0, /* tp_setattr */
2990 0, /* tp_reserved */
2991 0, /* tp_repr */
2992 0, /* tp_as_number */
2993 0, /* tp_as_sequence */
2994 0, /* tp_as_mapping */
2995 0, /* tp_hash */
2996 0, /* tp_call */
2997 0, /* tp_str */
2998 PyObject_GenericGetAttr, /* tp_getattro */
2999 0, /* tp_setattro */
3000 0, /* tp_as_buffer */
3001 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3002 Py_TPFLAGS_BASETYPE, /* tp_flags */
3003 filterfalse_doc, /* tp_doc */
3004 (traverseproc)filterfalse_traverse, /* tp_traverse */
3005 0, /* tp_clear */
3006 0, /* tp_richcompare */
3007 0, /* tp_weaklistoffset */
3008 PyObject_SelfIter, /* tp_iter */
3009 (iternextfunc)filterfalse_next, /* tp_iternext */
3010 0, /* tp_methods */
3011 0, /* tp_members */
3012 0, /* tp_getset */
3013 0, /* tp_base */
3014 0, /* tp_dict */
3015 0, /* tp_descr_get */
3016 0, /* tp_descr_set */
3017 0, /* tp_dictoffset */
3018 0, /* tp_init */
3019 0, /* tp_alloc */
3020 filterfalse_new, /* tp_new */
3021 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003022};
3023
3024
3025/* count object ************************************************************/
3026
3027typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003028 PyObject_HEAD
3029 Py_ssize_t cnt;
3030 PyObject *long_cnt;
3031 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003032} countobject;
3033
Raymond Hettinger30729212009-02-12 06:28:27 +00003034/* Counting logic and invariants:
3035
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003036fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003037
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003038 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3039 Advances with: cnt += 1
3040 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003041
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003042slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003043
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003044 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3045 All counting is done with python objects (no overflows or underflows).
3046 Advances with: long_cnt += long_step
3047 Step may be zero -- effectively a slow version of repeat(cnt).
3048 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003049*/
3050
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003051static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003052
3053static PyObject *
3054count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3055{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003056 countobject *lz;
3057 int slow_mode = 0;
3058 Py_ssize_t cnt = 0;
3059 PyObject *long_cnt = NULL;
3060 PyObject *long_step = NULL;
3061 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003062
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003063 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3064 kwlist, &long_cnt, &long_step))
3065 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003066
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003067 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3068 (long_step != NULL && !PyNumber_Check(long_step))) {
3069 PyErr_SetString(PyExc_TypeError, "a number is required");
3070 return NULL;
3071 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003072
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003073 if (long_cnt != NULL) {
3074 cnt = PyLong_AsSsize_t(long_cnt);
3075 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3076 PyErr_Clear();
3077 slow_mode = 1;
3078 }
3079 Py_INCREF(long_cnt);
3080 } else {
3081 cnt = 0;
3082 long_cnt = PyLong_FromLong(0);
3083 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003084
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003085 /* If not specified, step defaults to 1 */
3086 if (long_step == NULL) {
3087 long_step = PyLong_FromLong(1);
3088 if (long_step == NULL) {
3089 Py_DECREF(long_cnt);
3090 return NULL;
3091 }
3092 } else
3093 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003094
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003095 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003096
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003097 /* Fast mode only works when the step is 1 */
3098 if (!PyLong_Check(long_step) ||
3099 PyLong_AS_LONG(long_step) != 1) {
3100 slow_mode = 1;
3101 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003102
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003103 if (slow_mode)
3104 cnt = PY_SSIZE_T_MAX;
3105 else
3106 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003107
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003108 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3109 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3110 assert(slow_mode ||
3111 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003112
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003113 /* create countobject structure */
3114 lz = (countobject *)type->tp_alloc(type, 0);
3115 if (lz == NULL) {
3116 Py_XDECREF(long_cnt);
3117 return NULL;
3118 }
3119 lz->cnt = cnt;
3120 lz->long_cnt = long_cnt;
3121 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003122
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003123 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003124}
3125
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003126static void
3127count_dealloc(countobject *lz)
3128{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003129 PyObject_GC_UnTrack(lz);
3130 Py_XDECREF(lz->long_cnt);
3131 Py_XDECREF(lz->long_step);
3132 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003133}
3134
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003135static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003136count_traverse(countobject *lz, visitproc visit, void *arg)
3137{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003138 Py_VISIT(lz->long_cnt);
3139 Py_VISIT(lz->long_step);
3140 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003141}
3142
3143static PyObject *
3144count_nextlong(countobject *lz)
3145{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003146 PyObject *long_cnt;
3147 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003148
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003149 long_cnt = lz->long_cnt;
3150 if (long_cnt == NULL) {
3151 /* Switch to slow_mode */
3152 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3153 if (long_cnt == NULL)
3154 return NULL;
3155 }
3156 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00003157
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003158 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3159 if (stepped_up == NULL)
3160 return NULL;
3161 lz->long_cnt = stepped_up;
3162 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003163}
3164
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003165static PyObject *
3166count_next(countobject *lz)
3167{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003168 if (lz->cnt == PY_SSIZE_T_MAX)
3169 return count_nextlong(lz);
3170 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003171}
3172
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003173static PyObject *
3174count_repr(countobject *lz)
3175{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003176 if (lz->cnt != PY_SSIZE_T_MAX)
3177 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003178
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003179 if (PyLong_Check(lz->long_step)) {
3180 long step = PyLong_AsLong(lz->long_step);
3181 if (step == -1 && PyErr_Occurred()) {
3182 PyErr_Clear();
3183 }
3184 if (step == 1) {
3185 /* Don't display step when it is an integer equal to 1 */
3186 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
3187 }
3188 }
3189 return PyUnicode_FromFormat("count(%R, %R)",
3190 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003191}
3192
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003193static PyObject *
3194count_reduce(countobject *lz)
3195{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003196 if (lz->cnt == PY_SSIZE_T_MAX)
3197 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3198 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003199}
3200
3201PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3202
3203static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003204 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3205 count_reduce_doc},
3206 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003207};
3208
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003209PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00003210 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003211\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003212Return a count object whose .__next__() method returns consecutive values.\n\
3213Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00003214 def count(firstval=0, step=1):\n\
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003215 x = firstval\n\
3216 while 1:\n\
3217 yield x\n\
3218 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003219
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003220static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003221 PyVarObject_HEAD_INIT(NULL, 0)
3222 "itertools.count", /* tp_name */
3223 sizeof(countobject), /* tp_basicsize */
3224 0, /* tp_itemsize */
3225 /* methods */
3226 (destructor)count_dealloc, /* tp_dealloc */
3227 0, /* tp_print */
3228 0, /* tp_getattr */
3229 0, /* tp_setattr */
3230 0, /* tp_reserved */
3231 (reprfunc)count_repr, /* tp_repr */
3232 0, /* tp_as_number */
3233 0, /* tp_as_sequence */
3234 0, /* tp_as_mapping */
3235 0, /* tp_hash */
3236 0, /* tp_call */
3237 0, /* tp_str */
3238 PyObject_GenericGetAttr, /* tp_getattro */
3239 0, /* tp_setattro */
3240 0, /* tp_as_buffer */
3241 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3242 Py_TPFLAGS_BASETYPE, /* tp_flags */
3243 count_doc, /* tp_doc */
3244 (traverseproc)count_traverse, /* tp_traverse */
3245 0, /* tp_clear */
3246 0, /* tp_richcompare */
3247 0, /* tp_weaklistoffset */
3248 PyObject_SelfIter, /* tp_iter */
3249 (iternextfunc)count_next, /* tp_iternext */
3250 count_methods, /* tp_methods */
3251 0, /* tp_members */
3252 0, /* tp_getset */
3253 0, /* tp_base */
3254 0, /* tp_dict */
3255 0, /* tp_descr_get */
3256 0, /* tp_descr_set */
3257 0, /* tp_dictoffset */
3258 0, /* tp_init */
3259 0, /* tp_alloc */
3260 count_new, /* tp_new */
3261 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003262};
3263
3264
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003265/* repeat object ************************************************************/
3266
3267typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003268 PyObject_HEAD
3269 PyObject *element;
3270 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003271} repeatobject;
3272
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003273static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003274
3275static PyObject *
3276repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3277{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003278 repeatobject *ro;
3279 PyObject *element;
3280 Py_ssize_t cnt = -1;
3281 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003282
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003283 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3284 &element, &cnt))
3285 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003286
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003287 if (PyTuple_Size(args) == 2 && cnt < 0)
3288 cnt = 0;
3289
3290 ro = (repeatobject *)type->tp_alloc(type, 0);
3291 if (ro == NULL)
3292 return NULL;
3293 Py_INCREF(element);
3294 ro->element = element;
3295 ro->cnt = cnt;
3296 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003297}
3298
3299static void
3300repeat_dealloc(repeatobject *ro)
3301{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003302 PyObject_GC_UnTrack(ro);
3303 Py_XDECREF(ro->element);
3304 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003305}
3306
3307static int
3308repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3309{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003310 Py_VISIT(ro->element);
3311 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003312}
3313
3314static PyObject *
3315repeat_next(repeatobject *ro)
3316{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003317 if (ro->cnt == 0)
3318 return NULL;
3319 if (ro->cnt > 0)
3320 ro->cnt--;
3321 Py_INCREF(ro->element);
3322 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003323}
3324
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003325static PyObject *
3326repeat_repr(repeatobject *ro)
3327{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003328 if (ro->cnt == -1)
3329 return PyUnicode_FromFormat("repeat(%R)", ro->element);
3330 else
3331 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
3332}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003333
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003334static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003335repeat_len(repeatobject *ro)
3336{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003337 if (ro->cnt == -1) {
3338 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3339 return NULL;
3340 }
3341 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003342}
3343
Armin Rigof5b3e362006-02-11 21:32:43 +00003344PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003345
3346static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003347 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3348 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003349};
3350
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003351PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00003352"repeat(object [,times]) -> create an iterator which returns the object\n\
3353for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003354endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003355
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003356static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003357 PyVarObject_HEAD_INIT(NULL, 0)
3358 "itertools.repeat", /* tp_name */
3359 sizeof(repeatobject), /* tp_basicsize */
3360 0, /* tp_itemsize */
3361 /* methods */
3362 (destructor)repeat_dealloc, /* tp_dealloc */
3363 0, /* tp_print */
3364 0, /* tp_getattr */
3365 0, /* tp_setattr */
3366 0, /* tp_reserved */
3367 (reprfunc)repeat_repr, /* tp_repr */
3368 0, /* tp_as_number */
3369 0, /* tp_as_sequence */
3370 0, /* tp_as_mapping */
3371 0, /* tp_hash */
3372 0, /* tp_call */
3373 0, /* tp_str */
3374 PyObject_GenericGetAttr, /* tp_getattro */
3375 0, /* tp_setattro */
3376 0, /* tp_as_buffer */
3377 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3378 Py_TPFLAGS_BASETYPE, /* tp_flags */
3379 repeat_doc, /* tp_doc */
3380 (traverseproc)repeat_traverse, /* tp_traverse */
3381 0, /* tp_clear */
3382 0, /* tp_richcompare */
3383 0, /* tp_weaklistoffset */
3384 PyObject_SelfIter, /* tp_iter */
3385 (iternextfunc)repeat_next, /* tp_iternext */
3386 repeat_methods, /* tp_methods */
3387 0, /* tp_members */
3388 0, /* tp_getset */
3389 0, /* tp_base */
3390 0, /* tp_dict */
3391 0, /* tp_descr_get */
3392 0, /* tp_descr_set */
3393 0, /* tp_dictoffset */
3394 0, /* tp_init */
3395 0, /* tp_alloc */
3396 repeat_new, /* tp_new */
3397 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003398};
3399
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003400/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00003401
3402#include "Python.h"
3403
3404typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003405 PyObject_HEAD
3406 Py_ssize_t tuplesize;
3407 Py_ssize_t numactive;
3408 PyObject *ittuple; /* tuple of iterators */
3409 PyObject *result;
3410 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003411} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003412
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003413static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003414
3415static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003416zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003417{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003418 ziplongestobject *lz;
3419 Py_ssize_t i;
3420 PyObject *ittuple; /* tuple of iterators */
3421 PyObject *result;
3422 PyObject *fillvalue = Py_None;
3423 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003424
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003425 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3426 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3427 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3428 PyErr_SetString(PyExc_TypeError,
3429 "zip_longest() got an unexpected keyword argument");
3430 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003431 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003432 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003433
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003434 /* args must be a tuple */
3435 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00003436
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003437 /* obtain iterators */
3438 ittuple = PyTuple_New(tuplesize);
3439 if (ittuple == NULL)
3440 return NULL;
3441 for (i=0; i < tuplesize; ++i) {
3442 PyObject *item = PyTuple_GET_ITEM(args, i);
3443 PyObject *it = PyObject_GetIter(item);
3444 if (it == NULL) {
3445 if (PyErr_ExceptionMatches(PyExc_TypeError))
3446 PyErr_Format(PyExc_TypeError,
3447 "zip_longest argument #%zd must support iteration",
3448 i+1);
3449 Py_DECREF(ittuple);
3450 return NULL;
3451 }
3452 PyTuple_SET_ITEM(ittuple, i, it);
3453 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003454
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003455 /* create a result holder */
3456 result = PyTuple_New(tuplesize);
3457 if (result == NULL) {
3458 Py_DECREF(ittuple);
3459 return NULL;
3460 }
3461 for (i=0 ; i < tuplesize ; i++) {
3462 Py_INCREF(Py_None);
3463 PyTuple_SET_ITEM(result, i, Py_None);
3464 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003465
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003466 /* create ziplongestobject structure */
3467 lz = (ziplongestobject *)type->tp_alloc(type, 0);
3468 if (lz == NULL) {
3469 Py_DECREF(ittuple);
3470 Py_DECREF(result);
3471 return NULL;
3472 }
3473 lz->ittuple = ittuple;
3474 lz->tuplesize = tuplesize;
3475 lz->numactive = tuplesize;
3476 lz->result = result;
3477 Py_INCREF(fillvalue);
3478 lz->fillvalue = fillvalue;
3479 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003480}
3481
3482static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003483zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003484{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003485 PyObject_GC_UnTrack(lz);
3486 Py_XDECREF(lz->ittuple);
3487 Py_XDECREF(lz->result);
3488 Py_XDECREF(lz->fillvalue);
3489 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003490}
3491
3492static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003493zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003494{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003495 Py_VISIT(lz->ittuple);
3496 Py_VISIT(lz->result);
3497 Py_VISIT(lz->fillvalue);
3498 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003499}
3500
3501static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003502zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003503{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003504 Py_ssize_t i;
3505 Py_ssize_t tuplesize = lz->tuplesize;
3506 PyObject *result = lz->result;
3507 PyObject *it;
3508 PyObject *item;
3509 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003510
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003511 if (tuplesize == 0)
3512 return NULL;
3513 if (lz->numactive == 0)
3514 return NULL;
3515 if (Py_REFCNT(result) == 1) {
3516 Py_INCREF(result);
3517 for (i=0 ; i < tuplesize ; i++) {
3518 it = PyTuple_GET_ITEM(lz->ittuple, i);
3519 if (it == NULL) {
3520 Py_INCREF(lz->fillvalue);
3521 item = lz->fillvalue;
3522 } else {
3523 item = PyIter_Next(it);
3524 if (item == NULL) {
3525 lz->numactive -= 1;
3526 if (lz->numactive == 0 || PyErr_Occurred()) {
3527 lz->numactive = 0;
3528 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00003529 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003530 } else {
3531 Py_INCREF(lz->fillvalue);
3532 item = lz->fillvalue;
3533 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3534 Py_DECREF(it);
3535 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00003536 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003537 }
3538 olditem = PyTuple_GET_ITEM(result, i);
3539 PyTuple_SET_ITEM(result, i, item);
3540 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00003541 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003542 } else {
3543 result = PyTuple_New(tuplesize);
3544 if (result == NULL)
3545 return NULL;
3546 for (i=0 ; i < tuplesize ; i++) {
3547 it = PyTuple_GET_ITEM(lz->ittuple, i);
3548 if (it == NULL) {
3549 Py_INCREF(lz->fillvalue);
3550 item = lz->fillvalue;
3551 } else {
3552 item = PyIter_Next(it);
3553 if (item == NULL) {
3554 lz->numactive -= 1;
3555 if (lz->numactive == 0 || PyErr_Occurred()) {
3556 lz->numactive = 0;
3557 Py_DECREF(result);
3558 return NULL;
3559 } else {
3560 Py_INCREF(lz->fillvalue);
3561 item = lz->fillvalue;
3562 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3563 Py_DECREF(it);
3564 }
3565 }
3566 }
3567 PyTuple_SET_ITEM(result, i, item);
3568 }
3569 }
3570 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003571}
3572
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003573PyDoc_STRVAR(zip_longest_doc,
3574"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003575\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003576Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00003577the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003578method continues until the longest iterable in the argument sequence\n\
3579is exhausted and then it raises StopIteration. When the shorter iterables\n\
3580are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3581defaults to None or can be specified by a keyword argument.\n\
3582");
3583
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003584static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003585 PyVarObject_HEAD_INIT(NULL, 0)
3586 "itertools.zip_longest", /* tp_name */
3587 sizeof(ziplongestobject), /* tp_basicsize */
3588 0, /* tp_itemsize */
3589 /* methods */
3590 (destructor)zip_longest_dealloc, /* tp_dealloc */
3591 0, /* tp_print */
3592 0, /* tp_getattr */
3593 0, /* tp_setattr */
3594 0, /* tp_reserved */
3595 0, /* tp_repr */
3596 0, /* tp_as_number */
3597 0, /* tp_as_sequence */
3598 0, /* tp_as_mapping */
3599 0, /* tp_hash */
3600 0, /* tp_call */
3601 0, /* tp_str */
3602 PyObject_GenericGetAttr, /* tp_getattro */
3603 0, /* tp_setattro */
3604 0, /* tp_as_buffer */
3605 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3606 Py_TPFLAGS_BASETYPE, /* tp_flags */
3607 zip_longest_doc, /* tp_doc */
3608 (traverseproc)zip_longest_traverse, /* tp_traverse */
3609 0, /* tp_clear */
3610 0, /* tp_richcompare */
3611 0, /* tp_weaklistoffset */
3612 PyObject_SelfIter, /* tp_iter */
3613 (iternextfunc)zip_longest_next, /* tp_iternext */
3614 0, /* tp_methods */
3615 0, /* tp_members */
3616 0, /* tp_getset */
3617 0, /* tp_base */
3618 0, /* tp_dict */
3619 0, /* tp_descr_get */
3620 0, /* tp_descr_set */
3621 0, /* tp_dictoffset */
3622 0, /* tp_init */
3623 0, /* tp_alloc */
3624 zip_longest_new, /* tp_new */
3625 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003626};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003627
3628/* module level code ********************************************************/
3629
3630PyDoc_STRVAR(module_doc,
3631"Functional tools for creating and using iterators.\n\
3632\n\
3633Infinite iterators:\n\
3634count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003635cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003636repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003637\n\
3638Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003639accumulate(p, start=0) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003640chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
3641compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
3642dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
3643groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003644filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003645islice(seq, [start,] stop [, step]) --> elements from\n\
3646 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003647starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003648tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003649takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003650zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00003651\n\
3652Combinatoric generators:\n\
3653product(p, q, ... [repeat=1]) --> cartesian product\n\
3654permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00003655combinations(p, r)\n\
3656combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003657");
3658
3659
Raymond Hettingerad983e72003-11-12 14:32:26 +00003660static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003661 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3662 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00003663};
3664
Martin v. Löwis1a214512008-06-11 05:26:20 +00003665
3666static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003667 PyModuleDef_HEAD_INIT,
3668 "itertools",
3669 module_doc,
3670 -1,
3671 module_methods,
3672 NULL,
3673 NULL,
3674 NULL,
3675 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00003676};
3677
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003678PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00003679PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003680{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003681 int i;
3682 PyObject *m;
3683 char *name;
3684 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00003685 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003686 &combinations_type,
3687 &cwr_type,
3688 &cycle_type,
3689 &dropwhile_type,
3690 &takewhile_type,
3691 &islice_type,
3692 &starmap_type,
3693 &chain_type,
3694 &compress_type,
3695 &filterfalse_type,
3696 &count_type,
3697 &ziplongest_type,
3698 &permutations_type,
3699 &product_type,
3700 &repeat_type,
3701 &groupby_type,
3702 NULL
3703 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00003704
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003705 Py_TYPE(&teedataobject_type) = &PyType_Type;
3706 m = PyModule_Create(&itertoolsmodule);
3707 if (m == NULL)
3708 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003709
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003710 for (i=0 ; typelist[i] != NULL ; i++) {
3711 if (PyType_Ready(typelist[i]) < 0)
3712 return NULL;
3713 name = strchr(typelist[i]->tp_name, '.');
3714 assert (name != NULL);
3715 Py_INCREF(typelist[i]);
3716 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
3717 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003718
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003719 if (PyType_Ready(&teedataobject_type) < 0)
3720 return NULL;
3721 if (PyType_Ready(&tee_type) < 0)
3722 return NULL;
3723 if (PyType_Ready(&_grouper_type) < 0)
3724 return NULL;
3725 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003726}