blob: ad22ec790c84e4a6bda49f3847736bdc343da3d9 [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;
Raymond Hettinger5d446132011-03-27 18:52:10 -07002593 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002594} accumulateobject;
2595
2596static PyTypeObject accumulate_type;
2597
2598static PyObject *
2599accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2600{
Raymond Hettinger5d446132011-03-27 18:52:10 -07002601 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00002602 PyObject *iterable;
2603 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07002604 PyObject *binop = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002605 accumulateobject *lz;
2606
Raymond Hettinger5d446132011-03-27 18:52:10 -07002607 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
2608 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002609 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002610
2611 /* Get iterator. */
2612 it = PyObject_GetIter(iterable);
2613 if (it == NULL)
2614 return NULL;
2615
Raymond Hettinger482ba772010-12-01 22:48:00 +00002616 /* create accumulateobject structure */
2617 lz = (accumulateobject *)type->tp_alloc(type, 0);
2618 if (lz == NULL) {
2619 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002620 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002621 }
2622
Raymond Hettinger5d446132011-03-27 18:52:10 -07002623 Py_XINCREF(binop);
2624 lz->binop = binop;
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002625 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002626 lz->it = it;
2627 return (PyObject *)lz;
2628}
2629
2630static void
2631accumulate_dealloc(accumulateobject *lz)
2632{
2633 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07002634 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00002635 Py_XDECREF(lz->total);
2636 Py_XDECREF(lz->it);
2637 Py_TYPE(lz)->tp_free(lz);
2638}
2639
2640static int
2641accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
2642{
Raymond Hettinger5d446132011-03-27 18:52:10 -07002643 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00002644 Py_VISIT(lz->it);
2645 Py_VISIT(lz->total);
2646 return 0;
2647}
2648
2649static PyObject *
2650accumulate_next(accumulateobject *lz)
2651{
2652 PyObject *val, *oldtotal, *newtotal;
2653
2654 val = PyIter_Next(lz->it);
2655 if (val == NULL)
2656 return NULL;
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002657
2658 if (lz->total == NULL) {
2659 Py_INCREF(val);
2660 lz->total = val;
2661 return lz->total;
2662 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07002663
2664 if (lz->binop == NULL)
2665 newtotal = PyNumber_Add(lz->total, val);
2666 else
2667 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002668 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00002669 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002670 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002671
2672 oldtotal = lz->total;
2673 lz->total = newtotal;
2674 Py_DECREF(oldtotal);
2675
2676 Py_INCREF(newtotal);
2677 return newtotal;
2678}
2679
2680PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07002681"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00002682\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07002683Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00002684
2685static PyTypeObject accumulate_type = {
2686 PyVarObject_HEAD_INIT(NULL, 0)
2687 "itertools.accumulate", /* tp_name */
2688 sizeof(accumulateobject), /* tp_basicsize */
2689 0, /* tp_itemsize */
2690 /* methods */
2691 (destructor)accumulate_dealloc, /* tp_dealloc */
2692 0, /* tp_print */
2693 0, /* tp_getattr */
2694 0, /* tp_setattr */
2695 0, /* tp_reserved */
2696 0, /* tp_repr */
2697 0, /* tp_as_number */
2698 0, /* tp_as_sequence */
2699 0, /* tp_as_mapping */
2700 0, /* tp_hash */
2701 0, /* tp_call */
2702 0, /* tp_str */
2703 PyObject_GenericGetAttr, /* tp_getattro */
2704 0, /* tp_setattro */
2705 0, /* tp_as_buffer */
2706 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2707 Py_TPFLAGS_BASETYPE, /* tp_flags */
2708 accumulate_doc, /* tp_doc */
2709 (traverseproc)accumulate_traverse, /* tp_traverse */
2710 0, /* tp_clear */
2711 0, /* tp_richcompare */
2712 0, /* tp_weaklistoffset */
2713 PyObject_SelfIter, /* tp_iter */
2714 (iternextfunc)accumulate_next, /* tp_iternext */
2715 0, /* tp_methods */
2716 0, /* tp_members */
2717 0, /* tp_getset */
2718 0, /* tp_base */
2719 0, /* tp_dict */
2720 0, /* tp_descr_get */
2721 0, /* tp_descr_set */
2722 0, /* tp_dictoffset */
2723 0, /* tp_init */
2724 0, /* tp_alloc */
2725 accumulate_new, /* tp_new */
2726 PyObject_GC_Del, /* tp_free */
2727};
2728
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002729
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002730/* compress object ************************************************************/
2731
2732/* Equivalent to:
2733
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002734 def compress(data, selectors):
2735 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2736 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002737*/
2738
2739typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002740 PyObject_HEAD
2741 PyObject *data;
2742 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002743} compressobject;
2744
2745static PyTypeObject compress_type;
2746
2747static PyObject *
2748compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2749{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002750 PyObject *seq1, *seq2;
2751 PyObject *data=NULL, *selectors=NULL;
2752 compressobject *lz;
2753 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002754
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002755 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2756 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002757
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002758 data = PyObject_GetIter(seq1);
2759 if (data == NULL)
2760 goto fail;
2761 selectors = PyObject_GetIter(seq2);
2762 if (selectors == NULL)
2763 goto fail;
2764
2765 /* create compressobject structure */
2766 lz = (compressobject *)type->tp_alloc(type, 0);
2767 if (lz == NULL)
2768 goto fail;
2769 lz->data = data;
2770 lz->selectors = selectors;
2771 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002772
2773fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002774 Py_XDECREF(data);
2775 Py_XDECREF(selectors);
2776 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002777}
2778
2779static void
2780compress_dealloc(compressobject *lz)
2781{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002782 PyObject_GC_UnTrack(lz);
2783 Py_XDECREF(lz->data);
2784 Py_XDECREF(lz->selectors);
2785 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002786}
2787
2788static int
2789compress_traverse(compressobject *lz, visitproc visit, void *arg)
2790{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002791 Py_VISIT(lz->data);
2792 Py_VISIT(lz->selectors);
2793 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002794}
2795
2796static PyObject *
2797compress_next(compressobject *lz)
2798{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002799 PyObject *data = lz->data, *selectors = lz->selectors;
2800 PyObject *datum, *selector;
2801 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2802 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2803 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002804
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002805 while (1) {
2806 /* Steps: get datum, get selector, evaluate selector.
2807 Order is important (to match the pure python version
2808 in terms of which input gets a chance to raise an
2809 exception first).
2810 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002811
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002812 datum = datanext(data);
2813 if (datum == NULL)
2814 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002815
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002816 selector = selectornext(selectors);
2817 if (selector == NULL) {
2818 Py_DECREF(datum);
2819 return NULL;
2820 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002821
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002822 ok = PyObject_IsTrue(selector);
2823 Py_DECREF(selector);
2824 if (ok == 1)
2825 return datum;
2826 Py_DECREF(datum);
2827 if (ok == -1)
2828 return NULL;
2829 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002830}
2831
2832PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00002833"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002834\n\
2835Return data elements corresponding to true selector elements.\n\
2836Forms a shorter iterator from selected data elements using the\n\
2837selectors to choose the data elements.");
2838
2839static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002840 PyVarObject_HEAD_INIT(NULL, 0)
2841 "itertools.compress", /* tp_name */
2842 sizeof(compressobject), /* tp_basicsize */
2843 0, /* tp_itemsize */
2844 /* methods */
2845 (destructor)compress_dealloc, /* tp_dealloc */
2846 0, /* tp_print */
2847 0, /* tp_getattr */
2848 0, /* tp_setattr */
2849 0, /* tp_reserved */
2850 0, /* tp_repr */
2851 0, /* tp_as_number */
2852 0, /* tp_as_sequence */
2853 0, /* tp_as_mapping */
2854 0, /* tp_hash */
2855 0, /* tp_call */
2856 0, /* tp_str */
2857 PyObject_GenericGetAttr, /* tp_getattro */
2858 0, /* tp_setattro */
2859 0, /* tp_as_buffer */
2860 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2861 Py_TPFLAGS_BASETYPE, /* tp_flags */
2862 compress_doc, /* tp_doc */
2863 (traverseproc)compress_traverse, /* tp_traverse */
2864 0, /* tp_clear */
2865 0, /* tp_richcompare */
2866 0, /* tp_weaklistoffset */
2867 PyObject_SelfIter, /* tp_iter */
2868 (iternextfunc)compress_next, /* tp_iternext */
2869 0, /* tp_methods */
2870 0, /* tp_members */
2871 0, /* tp_getset */
2872 0, /* tp_base */
2873 0, /* tp_dict */
2874 0, /* tp_descr_get */
2875 0, /* tp_descr_set */
2876 0, /* tp_dictoffset */
2877 0, /* tp_init */
2878 0, /* tp_alloc */
2879 compress_new, /* tp_new */
2880 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002881};
2882
2883
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002884/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00002885
2886typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002887 PyObject_HEAD
2888 PyObject *func;
2889 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002890} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002891
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002892static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002893
2894static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002895filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002896{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002897 PyObject *func, *seq;
2898 PyObject *it;
2899 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002900
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002901 if (type == &filterfalse_type &&
2902 !_PyArg_NoKeywords("filterfalse()", kwds))
2903 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002904
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002905 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
2906 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002907
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002908 /* Get iterator. */
2909 it = PyObject_GetIter(seq);
2910 if (it == NULL)
2911 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002912
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002913 /* create filterfalseobject structure */
2914 lz = (filterfalseobject *)type->tp_alloc(type, 0);
2915 if (lz == NULL) {
2916 Py_DECREF(it);
2917 return NULL;
2918 }
2919 Py_INCREF(func);
2920 lz->func = func;
2921 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002922
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002923 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002924}
2925
2926static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002927filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002928{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002929 PyObject_GC_UnTrack(lz);
2930 Py_XDECREF(lz->func);
2931 Py_XDECREF(lz->it);
2932 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002933}
2934
2935static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002936filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002937{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002938 Py_VISIT(lz->it);
2939 Py_VISIT(lz->func);
2940 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002941}
2942
2943static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002944filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002945{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002946 PyObject *item;
2947 PyObject *it = lz->it;
2948 long ok;
2949 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002950
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002951 iternext = *Py_TYPE(it)->tp_iternext;
2952 for (;;) {
2953 item = iternext(it);
2954 if (item == NULL)
2955 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002956
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002957 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2958 ok = PyObject_IsTrue(item);
2959 } else {
2960 PyObject *good;
2961 good = PyObject_CallFunctionObjArgs(lz->func,
2962 item, NULL);
2963 if (good == NULL) {
2964 Py_DECREF(item);
2965 return NULL;
2966 }
2967 ok = PyObject_IsTrue(good);
2968 Py_DECREF(good);
2969 }
2970 if (!ok)
2971 return item;
2972 Py_DECREF(item);
2973 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00002974}
2975
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002976PyDoc_STRVAR(filterfalse_doc,
2977"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002978\n\
2979Return those items of sequence for which function(item) is false.\n\
2980If function is None, return the items that are false.");
2981
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002982static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002983 PyVarObject_HEAD_INIT(NULL, 0)
2984 "itertools.filterfalse", /* tp_name */
2985 sizeof(filterfalseobject), /* tp_basicsize */
2986 0, /* tp_itemsize */
2987 /* methods */
2988 (destructor)filterfalse_dealloc, /* tp_dealloc */
2989 0, /* tp_print */
2990 0, /* tp_getattr */
2991 0, /* tp_setattr */
2992 0, /* tp_reserved */
2993 0, /* tp_repr */
2994 0, /* tp_as_number */
2995 0, /* tp_as_sequence */
2996 0, /* tp_as_mapping */
2997 0, /* tp_hash */
2998 0, /* tp_call */
2999 0, /* tp_str */
3000 PyObject_GenericGetAttr, /* tp_getattro */
3001 0, /* tp_setattro */
3002 0, /* tp_as_buffer */
3003 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3004 Py_TPFLAGS_BASETYPE, /* tp_flags */
3005 filterfalse_doc, /* tp_doc */
3006 (traverseproc)filterfalse_traverse, /* tp_traverse */
3007 0, /* tp_clear */
3008 0, /* tp_richcompare */
3009 0, /* tp_weaklistoffset */
3010 PyObject_SelfIter, /* tp_iter */
3011 (iternextfunc)filterfalse_next, /* tp_iternext */
3012 0, /* tp_methods */
3013 0, /* tp_members */
3014 0, /* tp_getset */
3015 0, /* tp_base */
3016 0, /* tp_dict */
3017 0, /* tp_descr_get */
3018 0, /* tp_descr_set */
3019 0, /* tp_dictoffset */
3020 0, /* tp_init */
3021 0, /* tp_alloc */
3022 filterfalse_new, /* tp_new */
3023 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003024};
3025
3026
3027/* count object ************************************************************/
3028
3029typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003030 PyObject_HEAD
3031 Py_ssize_t cnt;
3032 PyObject *long_cnt;
3033 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003034} countobject;
3035
Raymond Hettinger30729212009-02-12 06:28:27 +00003036/* Counting logic and invariants:
3037
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003038fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003039
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003040 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3041 Advances with: cnt += 1
3042 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003043
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003044slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003045
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003046 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3047 All counting is done with python objects (no overflows or underflows).
3048 Advances with: long_cnt += long_step
3049 Step may be zero -- effectively a slow version of repeat(cnt).
3050 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003051*/
3052
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003053static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003054
3055static PyObject *
3056count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3057{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003058 countobject *lz;
3059 int slow_mode = 0;
3060 Py_ssize_t cnt = 0;
3061 PyObject *long_cnt = NULL;
3062 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003063 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003064 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003065
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003066 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3067 kwlist, &long_cnt, &long_step))
3068 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003069
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003070 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3071 (long_step != NULL && !PyNumber_Check(long_step))) {
3072 PyErr_SetString(PyExc_TypeError, "a number is required");
3073 return NULL;
3074 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003075
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003076 if (long_cnt != NULL) {
3077 cnt = PyLong_AsSsize_t(long_cnt);
3078 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3079 PyErr_Clear();
3080 slow_mode = 1;
3081 }
3082 Py_INCREF(long_cnt);
3083 } else {
3084 cnt = 0;
3085 long_cnt = PyLong_FromLong(0);
3086 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003087
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003088 /* If not specified, step defaults to 1 */
3089 if (long_step == NULL) {
3090 long_step = PyLong_FromLong(1);
3091 if (long_step == NULL) {
3092 Py_DECREF(long_cnt);
3093 return NULL;
3094 }
3095 } else
3096 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003097
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003098 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003099
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003100 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003101 step = PyLong_AsLong(long_step);
3102 if (step != 1) {
3103 slow_mode = 1;
3104 if (step == -1 && PyErr_Occurred())
3105 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003106 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003107
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003108 if (slow_mode)
3109 cnt = PY_SSIZE_T_MAX;
3110 else
3111 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003112
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003113 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3114 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3115 assert(slow_mode ||
3116 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003117
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003118 /* create countobject structure */
3119 lz = (countobject *)type->tp_alloc(type, 0);
3120 if (lz == NULL) {
3121 Py_XDECREF(long_cnt);
3122 return NULL;
3123 }
3124 lz->cnt = cnt;
3125 lz->long_cnt = long_cnt;
3126 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003127
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003128 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003129}
3130
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003131static void
3132count_dealloc(countobject *lz)
3133{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003134 PyObject_GC_UnTrack(lz);
3135 Py_XDECREF(lz->long_cnt);
3136 Py_XDECREF(lz->long_step);
3137 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003138}
3139
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003140static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003141count_traverse(countobject *lz, visitproc visit, void *arg)
3142{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003143 Py_VISIT(lz->long_cnt);
3144 Py_VISIT(lz->long_step);
3145 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003146}
3147
3148static PyObject *
3149count_nextlong(countobject *lz)
3150{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003151 PyObject *long_cnt;
3152 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003153
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003154 long_cnt = lz->long_cnt;
3155 if (long_cnt == NULL) {
3156 /* Switch to slow_mode */
3157 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3158 if (long_cnt == NULL)
3159 return NULL;
3160 }
3161 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00003162
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003163 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3164 if (stepped_up == NULL)
3165 return NULL;
3166 lz->long_cnt = stepped_up;
3167 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003168}
3169
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003170static PyObject *
3171count_next(countobject *lz)
3172{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003173 if (lz->cnt == PY_SSIZE_T_MAX)
3174 return count_nextlong(lz);
3175 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003176}
3177
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003178static PyObject *
3179count_repr(countobject *lz)
3180{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003181 if (lz->cnt != PY_SSIZE_T_MAX)
3182 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003183
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003184 if (PyLong_Check(lz->long_step)) {
3185 long step = PyLong_AsLong(lz->long_step);
3186 if (step == -1 && PyErr_Occurred()) {
3187 PyErr_Clear();
3188 }
3189 if (step == 1) {
3190 /* Don't display step when it is an integer equal to 1 */
3191 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
3192 }
3193 }
3194 return PyUnicode_FromFormat("count(%R, %R)",
3195 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003196}
3197
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003198static PyObject *
3199count_reduce(countobject *lz)
3200{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003201 if (lz->cnt == PY_SSIZE_T_MAX)
3202 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3203 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003204}
3205
3206PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3207
3208static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003209 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3210 count_reduce_doc},
3211 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003212};
3213
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003214PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00003215 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003216\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003217Return a count object whose .__next__() method returns consecutive values.\n\
3218Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00003219 def count(firstval=0, step=1):\n\
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003220 x = firstval\n\
3221 while 1:\n\
3222 yield x\n\
3223 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003224
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003225static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003226 PyVarObject_HEAD_INIT(NULL, 0)
3227 "itertools.count", /* tp_name */
3228 sizeof(countobject), /* tp_basicsize */
3229 0, /* tp_itemsize */
3230 /* methods */
3231 (destructor)count_dealloc, /* tp_dealloc */
3232 0, /* tp_print */
3233 0, /* tp_getattr */
3234 0, /* tp_setattr */
3235 0, /* tp_reserved */
3236 (reprfunc)count_repr, /* tp_repr */
3237 0, /* tp_as_number */
3238 0, /* tp_as_sequence */
3239 0, /* tp_as_mapping */
3240 0, /* tp_hash */
3241 0, /* tp_call */
3242 0, /* tp_str */
3243 PyObject_GenericGetAttr, /* tp_getattro */
3244 0, /* tp_setattro */
3245 0, /* tp_as_buffer */
3246 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3247 Py_TPFLAGS_BASETYPE, /* tp_flags */
3248 count_doc, /* tp_doc */
3249 (traverseproc)count_traverse, /* tp_traverse */
3250 0, /* tp_clear */
3251 0, /* tp_richcompare */
3252 0, /* tp_weaklistoffset */
3253 PyObject_SelfIter, /* tp_iter */
3254 (iternextfunc)count_next, /* tp_iternext */
3255 count_methods, /* tp_methods */
3256 0, /* tp_members */
3257 0, /* tp_getset */
3258 0, /* tp_base */
3259 0, /* tp_dict */
3260 0, /* tp_descr_get */
3261 0, /* tp_descr_set */
3262 0, /* tp_dictoffset */
3263 0, /* tp_init */
3264 0, /* tp_alloc */
3265 count_new, /* tp_new */
3266 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003267};
3268
3269
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003270/* repeat object ************************************************************/
3271
3272typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003273 PyObject_HEAD
3274 PyObject *element;
3275 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003276} repeatobject;
3277
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003278static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003279
3280static PyObject *
3281repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3282{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003283 repeatobject *ro;
3284 PyObject *element;
3285 Py_ssize_t cnt = -1;
3286 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003287
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003288 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3289 &element, &cnt))
3290 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003291
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003292 if (PyTuple_Size(args) == 2 && cnt < 0)
3293 cnt = 0;
3294
3295 ro = (repeatobject *)type->tp_alloc(type, 0);
3296 if (ro == NULL)
3297 return NULL;
3298 Py_INCREF(element);
3299 ro->element = element;
3300 ro->cnt = cnt;
3301 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003302}
3303
3304static void
3305repeat_dealloc(repeatobject *ro)
3306{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003307 PyObject_GC_UnTrack(ro);
3308 Py_XDECREF(ro->element);
3309 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003310}
3311
3312static int
3313repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3314{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003315 Py_VISIT(ro->element);
3316 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003317}
3318
3319static PyObject *
3320repeat_next(repeatobject *ro)
3321{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003322 if (ro->cnt == 0)
3323 return NULL;
3324 if (ro->cnt > 0)
3325 ro->cnt--;
3326 Py_INCREF(ro->element);
3327 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003328}
3329
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003330static PyObject *
3331repeat_repr(repeatobject *ro)
3332{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003333 if (ro->cnt == -1)
3334 return PyUnicode_FromFormat("repeat(%R)", ro->element);
3335 else
3336 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
3337}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003338
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003339static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003340repeat_len(repeatobject *ro)
3341{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003342 if (ro->cnt == -1) {
3343 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3344 return NULL;
3345 }
3346 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003347}
3348
Armin Rigof5b3e362006-02-11 21:32:43 +00003349PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003350
3351static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003352 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3353 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003354};
3355
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003356PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00003357"repeat(object [,times]) -> create an iterator which returns the object\n\
3358for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003359endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003360
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003361static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003362 PyVarObject_HEAD_INIT(NULL, 0)
3363 "itertools.repeat", /* tp_name */
3364 sizeof(repeatobject), /* tp_basicsize */
3365 0, /* tp_itemsize */
3366 /* methods */
3367 (destructor)repeat_dealloc, /* tp_dealloc */
3368 0, /* tp_print */
3369 0, /* tp_getattr */
3370 0, /* tp_setattr */
3371 0, /* tp_reserved */
3372 (reprfunc)repeat_repr, /* tp_repr */
3373 0, /* tp_as_number */
3374 0, /* tp_as_sequence */
3375 0, /* tp_as_mapping */
3376 0, /* tp_hash */
3377 0, /* tp_call */
3378 0, /* tp_str */
3379 PyObject_GenericGetAttr, /* tp_getattro */
3380 0, /* tp_setattro */
3381 0, /* tp_as_buffer */
3382 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3383 Py_TPFLAGS_BASETYPE, /* tp_flags */
3384 repeat_doc, /* tp_doc */
3385 (traverseproc)repeat_traverse, /* tp_traverse */
3386 0, /* tp_clear */
3387 0, /* tp_richcompare */
3388 0, /* tp_weaklistoffset */
3389 PyObject_SelfIter, /* tp_iter */
3390 (iternextfunc)repeat_next, /* tp_iternext */
3391 repeat_methods, /* tp_methods */
3392 0, /* tp_members */
3393 0, /* tp_getset */
3394 0, /* tp_base */
3395 0, /* tp_dict */
3396 0, /* tp_descr_get */
3397 0, /* tp_descr_set */
3398 0, /* tp_dictoffset */
3399 0, /* tp_init */
3400 0, /* tp_alloc */
3401 repeat_new, /* tp_new */
3402 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003403};
3404
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003405/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00003406
3407#include "Python.h"
3408
3409typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003410 PyObject_HEAD
3411 Py_ssize_t tuplesize;
3412 Py_ssize_t numactive;
3413 PyObject *ittuple; /* tuple of iterators */
3414 PyObject *result;
3415 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003416} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003417
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003418static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003419
3420static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003421zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003422{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003423 ziplongestobject *lz;
3424 Py_ssize_t i;
3425 PyObject *ittuple; /* tuple of iterators */
3426 PyObject *result;
3427 PyObject *fillvalue = Py_None;
3428 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003429
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003430 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3431 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3432 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3433 PyErr_SetString(PyExc_TypeError,
3434 "zip_longest() got an unexpected keyword argument");
3435 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003436 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003437 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003438
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003439 /* args must be a tuple */
3440 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00003441
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003442 /* obtain iterators */
3443 ittuple = PyTuple_New(tuplesize);
3444 if (ittuple == NULL)
3445 return NULL;
3446 for (i=0; i < tuplesize; ++i) {
3447 PyObject *item = PyTuple_GET_ITEM(args, i);
3448 PyObject *it = PyObject_GetIter(item);
3449 if (it == NULL) {
3450 if (PyErr_ExceptionMatches(PyExc_TypeError))
3451 PyErr_Format(PyExc_TypeError,
3452 "zip_longest argument #%zd must support iteration",
3453 i+1);
3454 Py_DECREF(ittuple);
3455 return NULL;
3456 }
3457 PyTuple_SET_ITEM(ittuple, i, it);
3458 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003459
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003460 /* create a result holder */
3461 result = PyTuple_New(tuplesize);
3462 if (result == NULL) {
3463 Py_DECREF(ittuple);
3464 return NULL;
3465 }
3466 for (i=0 ; i < tuplesize ; i++) {
3467 Py_INCREF(Py_None);
3468 PyTuple_SET_ITEM(result, i, Py_None);
3469 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003470
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003471 /* create ziplongestobject structure */
3472 lz = (ziplongestobject *)type->tp_alloc(type, 0);
3473 if (lz == NULL) {
3474 Py_DECREF(ittuple);
3475 Py_DECREF(result);
3476 return NULL;
3477 }
3478 lz->ittuple = ittuple;
3479 lz->tuplesize = tuplesize;
3480 lz->numactive = tuplesize;
3481 lz->result = result;
3482 Py_INCREF(fillvalue);
3483 lz->fillvalue = fillvalue;
3484 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003485}
3486
3487static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003488zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003489{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003490 PyObject_GC_UnTrack(lz);
3491 Py_XDECREF(lz->ittuple);
3492 Py_XDECREF(lz->result);
3493 Py_XDECREF(lz->fillvalue);
3494 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003495}
3496
3497static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003498zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003499{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003500 Py_VISIT(lz->ittuple);
3501 Py_VISIT(lz->result);
3502 Py_VISIT(lz->fillvalue);
3503 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003504}
3505
3506static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003507zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003508{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003509 Py_ssize_t i;
3510 Py_ssize_t tuplesize = lz->tuplesize;
3511 PyObject *result = lz->result;
3512 PyObject *it;
3513 PyObject *item;
3514 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003515
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003516 if (tuplesize == 0)
3517 return NULL;
3518 if (lz->numactive == 0)
3519 return NULL;
3520 if (Py_REFCNT(result) == 1) {
3521 Py_INCREF(result);
3522 for (i=0 ; i < tuplesize ; i++) {
3523 it = PyTuple_GET_ITEM(lz->ittuple, i);
3524 if (it == NULL) {
3525 Py_INCREF(lz->fillvalue);
3526 item = lz->fillvalue;
3527 } else {
3528 item = PyIter_Next(it);
3529 if (item == NULL) {
3530 lz->numactive -= 1;
3531 if (lz->numactive == 0 || PyErr_Occurred()) {
3532 lz->numactive = 0;
3533 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00003534 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003535 } else {
3536 Py_INCREF(lz->fillvalue);
3537 item = lz->fillvalue;
3538 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3539 Py_DECREF(it);
3540 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00003541 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003542 }
3543 olditem = PyTuple_GET_ITEM(result, i);
3544 PyTuple_SET_ITEM(result, i, item);
3545 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00003546 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003547 } else {
3548 result = PyTuple_New(tuplesize);
3549 if (result == NULL)
3550 return NULL;
3551 for (i=0 ; i < tuplesize ; i++) {
3552 it = PyTuple_GET_ITEM(lz->ittuple, i);
3553 if (it == NULL) {
3554 Py_INCREF(lz->fillvalue);
3555 item = lz->fillvalue;
3556 } else {
3557 item = PyIter_Next(it);
3558 if (item == NULL) {
3559 lz->numactive -= 1;
3560 if (lz->numactive == 0 || PyErr_Occurred()) {
3561 lz->numactive = 0;
3562 Py_DECREF(result);
3563 return NULL;
3564 } else {
3565 Py_INCREF(lz->fillvalue);
3566 item = lz->fillvalue;
3567 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3568 Py_DECREF(it);
3569 }
3570 }
3571 }
3572 PyTuple_SET_ITEM(result, i, item);
3573 }
3574 }
3575 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003576}
3577
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003578PyDoc_STRVAR(zip_longest_doc,
3579"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003580\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003581Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00003582the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003583method continues until the longest iterable in the argument sequence\n\
3584is exhausted and then it raises StopIteration. When the shorter iterables\n\
3585are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3586defaults to None or can be specified by a keyword argument.\n\
3587");
3588
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003589static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003590 PyVarObject_HEAD_INIT(NULL, 0)
3591 "itertools.zip_longest", /* tp_name */
3592 sizeof(ziplongestobject), /* tp_basicsize */
3593 0, /* tp_itemsize */
3594 /* methods */
3595 (destructor)zip_longest_dealloc, /* tp_dealloc */
3596 0, /* tp_print */
3597 0, /* tp_getattr */
3598 0, /* tp_setattr */
3599 0, /* tp_reserved */
3600 0, /* tp_repr */
3601 0, /* tp_as_number */
3602 0, /* tp_as_sequence */
3603 0, /* tp_as_mapping */
3604 0, /* tp_hash */
3605 0, /* tp_call */
3606 0, /* tp_str */
3607 PyObject_GenericGetAttr, /* tp_getattro */
3608 0, /* tp_setattro */
3609 0, /* tp_as_buffer */
3610 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3611 Py_TPFLAGS_BASETYPE, /* tp_flags */
3612 zip_longest_doc, /* tp_doc */
3613 (traverseproc)zip_longest_traverse, /* tp_traverse */
3614 0, /* tp_clear */
3615 0, /* tp_richcompare */
3616 0, /* tp_weaklistoffset */
3617 PyObject_SelfIter, /* tp_iter */
3618 (iternextfunc)zip_longest_next, /* tp_iternext */
3619 0, /* tp_methods */
3620 0, /* tp_members */
3621 0, /* tp_getset */
3622 0, /* tp_base */
3623 0, /* tp_dict */
3624 0, /* tp_descr_get */
3625 0, /* tp_descr_set */
3626 0, /* tp_dictoffset */
3627 0, /* tp_init */
3628 0, /* tp_alloc */
3629 zip_longest_new, /* tp_new */
3630 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003631};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003632
3633/* module level code ********************************************************/
3634
3635PyDoc_STRVAR(module_doc,
3636"Functional tools for creating and using iterators.\n\
3637\n\
3638Infinite iterators:\n\
3639count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003640cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003641repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003642\n\
3643Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003644accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003645chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
3646compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
3647dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
3648groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003649filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003650islice(seq, [start,] stop [, step]) --> elements from\n\
3651 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003652starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003653tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003654takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003655zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00003656\n\
3657Combinatoric generators:\n\
3658product(p, q, ... [repeat=1]) --> cartesian product\n\
3659permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00003660combinations(p, r)\n\
3661combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003662");
3663
3664
Raymond Hettingerad983e72003-11-12 14:32:26 +00003665static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003666 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3667 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00003668};
3669
Martin v. Löwis1a214512008-06-11 05:26:20 +00003670
3671static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003672 PyModuleDef_HEAD_INIT,
3673 "itertools",
3674 module_doc,
3675 -1,
3676 module_methods,
3677 NULL,
3678 NULL,
3679 NULL,
3680 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00003681};
3682
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003683PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00003684PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003685{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003686 int i;
3687 PyObject *m;
3688 char *name;
3689 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00003690 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003691 &combinations_type,
3692 &cwr_type,
3693 &cycle_type,
3694 &dropwhile_type,
3695 &takewhile_type,
3696 &islice_type,
3697 &starmap_type,
3698 &chain_type,
3699 &compress_type,
3700 &filterfalse_type,
3701 &count_type,
3702 &ziplongest_type,
3703 &permutations_type,
3704 &product_type,
3705 &repeat_type,
3706 &groupby_type,
3707 NULL
3708 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00003709
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003710 Py_TYPE(&teedataobject_type) = &PyType_Type;
3711 m = PyModule_Create(&itertoolsmodule);
3712 if (m == NULL)
3713 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003714
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003715 for (i=0 ; typelist[i] != NULL ; i++) {
3716 if (PyType_Ready(typelist[i]) < 0)
3717 return NULL;
3718 name = strchr(typelist[i]->tp_name, '.');
3719 assert (name != NULL);
3720 Py_INCREF(typelist[i]);
3721 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
3722 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003723
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003724 if (PyType_Ready(&teedataobject_type) < 0)
3725 return NULL;
3726 if (PyType_Ready(&tee_type) < 0)
3727 return NULL;
3728 if (PyType_Ready(&_grouper_type) < 0)
3729 return NULL;
3730 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003731}