blob: b202e5262bab463e6abe240c0974b6548265b3d4 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
2#include "Python.h"
Fred Drake08ebfec2004-10-17 19:36:57 +00003#include "structmember.h"
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00005/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00006 by Raymond D. Hettinger <python@rcn.com>
7 Copyright (c) 2003 Python Software Foundation.
8 All rights reserved.
9*/
10
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000011
12/* groupby object ***********************************************************/
13
14typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000015 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000021} groupbyobject;
22
23static PyTypeObject groupby_type;
24static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26static PyObject *
27groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000029 static char *kwargs[] = {"iterable", "key", NULL};
30 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000032
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000033 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
34 &it, &keyfunc))
35 return NULL;
36
37 gbo = (groupbyobject *)type->tp_alloc(type, 0);
38 if (gbo == NULL)
39 return NULL;
40 gbo->tgtkey = NULL;
41 gbo->currkey = NULL;
42 gbo->currvalue = NULL;
43 gbo->keyfunc = keyfunc;
44 Py_INCREF(keyfunc);
45 gbo->it = PyObject_GetIter(it);
46 if (gbo->it == NULL) {
47 Py_DECREF(gbo);
48 return NULL;
49 }
50 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000051}
52
53static void
54groupby_dealloc(groupbyobject *gbo)
55{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000056 PyObject_GC_UnTrack(gbo);
57 Py_XDECREF(gbo->it);
58 Py_XDECREF(gbo->keyfunc);
59 Py_XDECREF(gbo->tgtkey);
60 Py_XDECREF(gbo->currkey);
61 Py_XDECREF(gbo->currvalue);
62 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000063}
64
65static int
66groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
67{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000068 Py_VISIT(gbo->it);
69 Py_VISIT(gbo->keyfunc);
70 Py_VISIT(gbo->tgtkey);
71 Py_VISIT(gbo->currkey);
72 Py_VISIT(gbo->currvalue);
73 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000074}
75
76static PyObject *
77groupby_next(groupbyobject *gbo)
78{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000081 /* skip to next iteration group */
82 for (;;) {
83 if (gbo->currkey == NULL)
84 /* pass */;
85 else if (gbo->tgtkey == NULL)
86 break;
87 else {
88 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000089
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000090 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
91 gbo->currkey, Py_EQ);
92 if (rcmp == -1)
93 return NULL;
94 else if (rcmp == 0)
95 break;
96 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000097
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000098 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000101
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000102 if (gbo->keyfunc == Py_None) {
103 newkey = newvalue;
104 Py_INCREF(newvalue);
105 } else {
106 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
107 newvalue, NULL);
108 if (newkey == NULL) {
109 Py_DECREF(newvalue);
110 return NULL;
111 }
112 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000113
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000114 tmp = gbo->currkey;
115 gbo->currkey = newkey;
116 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000117
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000118 tmp = gbo->currvalue;
119 gbo->currvalue = newvalue;
120 Py_XDECREF(tmp);
121 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000123 Py_INCREF(gbo->currkey);
124 tmp = gbo->tgtkey;
125 gbo->tgtkey = gbo->currkey;
126 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000127
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000131
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000135}
136
137PyDoc_STRVAR(groupby_doc,
138"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
139(key, sub-iterator) grouped by each value of key(value).\n");
140
141static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000142 PyVarObject_HEAD_INIT(NULL, 0)
143 "itertools.groupby", /* tp_name */
144 sizeof(groupbyobject), /* tp_basicsize */
145 0, /* tp_itemsize */
146 /* methods */
147 (destructor)groupby_dealloc, /* tp_dealloc */
148 0, /* tp_print */
149 0, /* tp_getattr */
150 0, /* tp_setattr */
151 0, /* tp_reserved */
152 0, /* tp_repr */
153 0, /* tp_as_number */
154 0, /* tp_as_sequence */
155 0, /* tp_as_mapping */
156 0, /* tp_hash */
157 0, /* tp_call */
158 0, /* tp_str */
159 PyObject_GenericGetAttr, /* tp_getattro */
160 0, /* tp_setattro */
161 0, /* tp_as_buffer */
162 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
163 Py_TPFLAGS_BASETYPE, /* tp_flags */
164 groupby_doc, /* tp_doc */
165 (traverseproc)groupby_traverse, /* tp_traverse */
166 0, /* tp_clear */
167 0, /* tp_richcompare */
168 0, /* tp_weaklistoffset */
169 PyObject_SelfIter, /* tp_iter */
170 (iternextfunc)groupby_next, /* tp_iternext */
171 0, /* tp_methods */
172 0, /* tp_members */
173 0, /* tp_getset */
174 0, /* tp_base */
175 0, /* tp_dict */
176 0, /* tp_descr_get */
177 0, /* tp_descr_set */
178 0, /* tp_dictoffset */
179 0, /* tp_init */
180 0, /* tp_alloc */
181 groupby_new, /* tp_new */
182 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000183};
184
185
186/* _grouper object (internal) ************************************************/
187
188typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000189 PyObject_HEAD
190 PyObject *parent;
191 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000192} _grouperobject;
193
194static PyTypeObject _grouper_type;
195
196static PyObject *
197_grouper_create(groupbyobject *parent, PyObject *tgtkey)
198{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000199 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000200
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000201 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
202 if (igo == NULL)
203 return NULL;
204 igo->parent = (PyObject *)parent;
205 Py_INCREF(parent);
206 igo->tgtkey = tgtkey;
207 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000208
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000209 PyObject_GC_Track(igo);
210 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000211}
212
213static void
214_grouper_dealloc(_grouperobject *igo)
215{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000216 PyObject_GC_UnTrack(igo);
217 Py_DECREF(igo->parent);
218 Py_DECREF(igo->tgtkey);
219 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000220}
221
222static int
223_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
224{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000225 Py_VISIT(igo->parent);
226 Py_VISIT(igo->tgtkey);
227 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000228}
229
230static PyObject *
231_grouper_next(_grouperobject *igo)
232{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000233 groupbyobject *gbo = (groupbyobject *)igo->parent;
234 PyObject *newvalue, *newkey, *r;
235 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000236
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000237 if (gbo->currvalue == NULL) {
238 newvalue = PyIter_Next(gbo->it);
239 if (newvalue == NULL)
240 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000241
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000242 if (gbo->keyfunc == Py_None) {
243 newkey = newvalue;
244 Py_INCREF(newvalue);
245 } else {
246 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
247 newvalue, NULL);
248 if (newkey == NULL) {
249 Py_DECREF(newvalue);
250 return NULL;
251 }
252 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000253
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000254 assert(gbo->currkey == NULL);
255 gbo->currkey = newkey;
256 gbo->currvalue = newvalue;
257 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000259 assert(gbo->currkey != NULL);
260 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
261 if (rcmp <= 0)
262 /* got any error or current group is end */
263 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000264
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000265 r = gbo->currvalue;
266 gbo->currvalue = NULL;
267 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000268
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000269 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000270}
271
272static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000273 PyVarObject_HEAD_INIT(NULL, 0)
274 "itertools._grouper", /* tp_name */
275 sizeof(_grouperobject), /* tp_basicsize */
276 0, /* tp_itemsize */
277 /* methods */
278 (destructor)_grouper_dealloc, /* tp_dealloc */
279 0, /* tp_print */
280 0, /* tp_getattr */
281 0, /* tp_setattr */
282 0, /* tp_reserved */
283 0, /* tp_repr */
284 0, /* tp_as_number */
285 0, /* tp_as_sequence */
286 0, /* tp_as_mapping */
287 0, /* tp_hash */
288 0, /* tp_call */
289 0, /* tp_str */
290 PyObject_GenericGetAttr, /* tp_getattro */
291 0, /* tp_setattro */
292 0, /* tp_as_buffer */
293 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
294 0, /* tp_doc */
295 (traverseproc)_grouper_traverse,/* tp_traverse */
296 0, /* tp_clear */
297 0, /* tp_richcompare */
298 0, /* tp_weaklistoffset */
299 PyObject_SelfIter, /* tp_iter */
300 (iternextfunc)_grouper_next, /* tp_iternext */
301 0, /* tp_methods */
302 0, /* tp_members */
303 0, /* tp_getset */
304 0, /* tp_base */
305 0, /* tp_dict */
306 0, /* tp_descr_get */
307 0, /* tp_descr_set */
308 0, /* tp_dictoffset */
309 0, /* tp_init */
310 0, /* tp_alloc */
311 0, /* tp_new */
312 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000313};
314
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000315
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000316
Raymond Hettingerad983e72003-11-12 14:32:26 +0000317/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000318
Raymond Hettingerad983e72003-11-12 14:32:26 +0000319/* The teedataobject pre-allocates space for LINKCELLS number of objects.
320 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000321 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000322 the other structure members including PyHEAD overhead. The larger the
323 value, the less memory overhead per object and the less time spent
324 allocating/deallocating new links. The smaller the number, the less
325 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000326*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000327#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000328
329typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000330 PyObject_HEAD
331 PyObject *it;
332 int numread;
333 PyObject *nextlink;
334 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000335} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000336
337typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000338 PyObject_HEAD
339 teedataobject *dataobj;
340 int index;
341 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000342} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000343
Raymond Hettingerad983e72003-11-12 14:32:26 +0000344static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000345
346static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000347teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000348{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000349 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000350
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000351 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
352 if (tdo == NULL)
353 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000354
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000355 tdo->numread = 0;
356 tdo->nextlink = NULL;
357 Py_INCREF(it);
358 tdo->it = it;
359 PyObject_GC_Track(tdo);
360 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000361}
362
363static PyObject *
364teedataobject_jumplink(teedataobject *tdo)
365{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000366 if (tdo->nextlink == NULL)
367 tdo->nextlink = teedataobject_new(tdo->it);
368 Py_XINCREF(tdo->nextlink);
369 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000370}
371
372static PyObject *
373teedataobject_getitem(teedataobject *tdo, int i)
374{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000375 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000376
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000377 assert(i < LINKCELLS);
378 if (i < tdo->numread)
379 value = tdo->values[i];
380 else {
381 /* this is the lead iterator, so fetch more data */
382 assert(i == tdo->numread);
383 value = PyIter_Next(tdo->it);
384 if (value == NULL)
385 return NULL;
386 tdo->numread++;
387 tdo->values[i] = value;
388 }
389 Py_INCREF(value);
390 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000391}
392
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000393static int
394teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
395{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000396 int i;
397 Py_VISIT(tdo->it);
398 for (i = 0; i < tdo->numread; i++)
399 Py_VISIT(tdo->values[i]);
400 Py_VISIT(tdo->nextlink);
401 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000402}
403
404static int
405teedataobject_clear(teedataobject *tdo)
406{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000407 int i;
408 Py_CLEAR(tdo->it);
409 for (i=0 ; i<tdo->numread ; i++)
410 Py_CLEAR(tdo->values[i]);
411 Py_CLEAR(tdo->nextlink);
412 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000413}
414
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000416teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000418 PyObject_GC_UnTrack(tdo);
419 teedataobject_clear(tdo);
420 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000421}
422
Raymond Hettingerad983e72003-11-12 14:32:26 +0000423PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000424
Raymond Hettingerad983e72003-11-12 14:32:26 +0000425static PyTypeObject teedataobject_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000426 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
427 "itertools.tee_dataobject", /* tp_name */
428 sizeof(teedataobject), /* tp_basicsize */
429 0, /* tp_itemsize */
430 /* methods */
431 (destructor)teedataobject_dealloc, /* tp_dealloc */
432 0, /* tp_print */
433 0, /* tp_getattr */
434 0, /* tp_setattr */
435 0, /* tp_reserved */
436 0, /* tp_repr */
437 0, /* tp_as_number */
438 0, /* tp_as_sequence */
439 0, /* tp_as_mapping */
440 0, /* tp_hash */
441 0, /* tp_call */
442 0, /* tp_str */
443 PyObject_GenericGetAttr, /* tp_getattro */
444 0, /* tp_setattro */
445 0, /* tp_as_buffer */
446 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
447 teedataobject_doc, /* tp_doc */
448 (traverseproc)teedataobject_traverse, /* tp_traverse */
449 (inquiry)teedataobject_clear, /* tp_clear */
450 0, /* tp_richcompare */
451 0, /* tp_weaklistoffset */
452 0, /* tp_iter */
453 0, /* tp_iternext */
454 0, /* tp_methods */
455 0, /* tp_members */
456 0, /* tp_getset */
457 0, /* tp_base */
458 0, /* tp_dict */
459 0, /* tp_descr_get */
460 0, /* tp_descr_set */
461 0, /* tp_dictoffset */
462 0, /* tp_init */
463 0, /* tp_alloc */
464 0, /* tp_new */
465 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000466};
467
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000468
469static PyTypeObject tee_type;
470
471static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000472tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000473{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000474 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000475
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000476 if (to->index >= LINKCELLS) {
477 link = teedataobject_jumplink(to->dataobj);
478 Py_DECREF(to->dataobj);
479 to->dataobj = (teedataobject *)link;
480 to->index = 0;
481 }
482 value = teedataobject_getitem(to->dataobj, to->index);
483 if (value == NULL)
484 return NULL;
485 to->index++;
486 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000487}
488
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000489static int
490tee_traverse(teeobject *to, visitproc visit, void *arg)
491{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000492 Py_VISIT((PyObject *)to->dataobj);
493 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000494}
495
Raymond Hettingerad983e72003-11-12 14:32:26 +0000496static PyObject *
497tee_copy(teeobject *to)
498{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000499 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000500
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000501 newto = PyObject_GC_New(teeobject, &tee_type);
502 if (newto == NULL)
503 return NULL;
504 Py_INCREF(to->dataobj);
505 newto->dataobj = to->dataobj;
506 newto->index = to->index;
507 newto->weakreflist = NULL;
508 PyObject_GC_Track(newto);
509 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000510}
511
512PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
513
514static PyObject *
515tee_fromiterable(PyObject *iterable)
516{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000517 teeobject *to;
518 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000519
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000520 it = PyObject_GetIter(iterable);
521 if (it == NULL)
522 return NULL;
523 if (PyObject_TypeCheck(it, &tee_type)) {
524 to = (teeobject *)tee_copy((teeobject *)it);
525 goto done;
526 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000527
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000528 to = PyObject_GC_New(teeobject, &tee_type);
529 if (to == NULL)
530 goto done;
531 to->dataobj = (teedataobject *)teedataobject_new(it);
532 if (!to->dataobj) {
533 PyObject_GC_Del(to);
534 to = NULL;
535 goto done;
536 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000537
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000538 to->index = 0;
539 to->weakreflist = NULL;
540 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000541done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000542 Py_XDECREF(it);
543 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000544}
545
546static PyObject *
547tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
548{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000549 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000550
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000551 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
552 return NULL;
553 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000554}
555
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000556static int
557tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000558{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000559 if (to->weakreflist != NULL)
560 PyObject_ClearWeakRefs((PyObject *) to);
561 Py_CLEAR(to->dataobj);
562 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000563}
564
565static void
566tee_dealloc(teeobject *to)
567{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000568 PyObject_GC_UnTrack(to);
569 tee_clear(to);
570 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000571}
572
Raymond Hettingerad983e72003-11-12 14:32:26 +0000573PyDoc_STRVAR(teeobject_doc,
574"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000575
Raymond Hettingerad983e72003-11-12 14:32:26 +0000576static PyMethodDef tee_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000577 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
578 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000579};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000580
581static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000582 PyVarObject_HEAD_INIT(NULL, 0)
583 "itertools.tee", /* tp_name */
584 sizeof(teeobject), /* tp_basicsize */
585 0, /* tp_itemsize */
586 /* methods */
587 (destructor)tee_dealloc, /* tp_dealloc */
588 0, /* tp_print */
589 0, /* tp_getattr */
590 0, /* tp_setattr */
591 0, /* tp_reserved */
592 0, /* tp_repr */
593 0, /* tp_as_number */
594 0, /* tp_as_sequence */
595 0, /* tp_as_mapping */
596 0, /* tp_hash */
597 0, /* tp_call */
598 0, /* tp_str */
599 0, /* tp_getattro */
600 0, /* tp_setattro */
601 0, /* tp_as_buffer */
602 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
603 teeobject_doc, /* tp_doc */
604 (traverseproc)tee_traverse, /* tp_traverse */
605 (inquiry)tee_clear, /* tp_clear */
606 0, /* tp_richcompare */
607 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
608 PyObject_SelfIter, /* tp_iter */
609 (iternextfunc)tee_next, /* tp_iternext */
610 tee_methods, /* tp_methods */
611 0, /* tp_members */
612 0, /* tp_getset */
613 0, /* tp_base */
614 0, /* tp_dict */
615 0, /* tp_descr_get */
616 0, /* tp_descr_set */
617 0, /* tp_dictoffset */
618 0, /* tp_init */
619 0, /* tp_alloc */
620 tee_new, /* tp_new */
621 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000622};
623
Raymond Hettingerad983e72003-11-12 14:32:26 +0000624static PyObject *
625tee(PyObject *self, PyObject *args)
626{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000627 Py_ssize_t i, n=2;
628 PyObject *it, *iterable, *copyable, *result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000629
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000630 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
631 return NULL;
632 if (n < 0) {
633 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
634 return NULL;
635 }
636 result = PyTuple_New(n);
637 if (result == NULL)
638 return NULL;
639 if (n == 0)
640 return result;
641 it = PyObject_GetIter(iterable);
642 if (it == NULL) {
643 Py_DECREF(result);
644 return NULL;
645 }
646 if (!PyObject_HasAttrString(it, "__copy__")) {
647 copyable = tee_fromiterable(it);
648 Py_DECREF(it);
649 if (copyable == NULL) {
650 Py_DECREF(result);
651 return NULL;
652 }
653 } else
654 copyable = it;
655 PyTuple_SET_ITEM(result, 0, copyable);
656 for (i=1 ; i<n ; i++) {
657 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
658 if (copyable == NULL) {
659 Py_DECREF(result);
660 return NULL;
661 }
662 PyTuple_SET_ITEM(result, i, copyable);
663 }
664 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000665}
666
667PyDoc_STRVAR(tee_doc,
668"tee(iterable, n=2) --> tuple of n independent iterators.");
669
670
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000671/* cycle object **********************************************************/
672
673typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000674 PyObject_HEAD
675 PyObject *it;
676 PyObject *saved;
677 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000678} cycleobject;
679
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000680static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000681
682static PyObject *
683cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
684{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000685 PyObject *it;
686 PyObject *iterable;
687 PyObject *saved;
688 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000689
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000690 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
691 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000692
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000693 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
694 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000695
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000696 /* Get iterator. */
697 it = PyObject_GetIter(iterable);
698 if (it == NULL)
699 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000700
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000701 saved = PyList_New(0);
702 if (saved == NULL) {
703 Py_DECREF(it);
704 return NULL;
705 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000706
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000707 /* create cycleobject structure */
708 lz = (cycleobject *)type->tp_alloc(type, 0);
709 if (lz == NULL) {
710 Py_DECREF(it);
711 Py_DECREF(saved);
712 return NULL;
713 }
714 lz->it = it;
715 lz->saved = saved;
716 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000717
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000718 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000719}
720
721static void
722cycle_dealloc(cycleobject *lz)
723{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000724 PyObject_GC_UnTrack(lz);
725 Py_XDECREF(lz->saved);
726 Py_XDECREF(lz->it);
727 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000728}
729
730static int
731cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
732{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000733 Py_VISIT(lz->it);
734 Py_VISIT(lz->saved);
735 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000736}
737
738static PyObject *
739cycle_next(cycleobject *lz)
740{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000741 PyObject *item;
742 PyObject *it;
743 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000744
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000745 while (1) {
746 item = PyIter_Next(lz->it);
747 if (item != NULL) {
748 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
749 Py_DECREF(item);
750 return NULL;
751 }
752 return item;
753 }
754 if (PyErr_Occurred()) {
755 if (PyErr_ExceptionMatches(PyExc_StopIteration))
756 PyErr_Clear();
757 else
758 return NULL;
759 }
760 if (PyList_Size(lz->saved) == 0)
761 return NULL;
762 it = PyObject_GetIter(lz->saved);
763 if (it == NULL)
764 return NULL;
765 tmp = lz->it;
766 lz->it = it;
767 lz->firstpass = 1;
768 Py_DECREF(tmp);
769 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000770}
771
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000772PyDoc_STRVAR(cycle_doc,
773"cycle(iterable) --> cycle object\n\
774\n\
775Return elements from the iterable until it is exhausted.\n\
776Then repeat the sequence indefinitely.");
777
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000778static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000779 PyVarObject_HEAD_INIT(NULL, 0)
780 "itertools.cycle", /* tp_name */
781 sizeof(cycleobject), /* tp_basicsize */
782 0, /* tp_itemsize */
783 /* methods */
784 (destructor)cycle_dealloc, /* tp_dealloc */
785 0, /* tp_print */
786 0, /* tp_getattr */
787 0, /* tp_setattr */
788 0, /* tp_reserved */
789 0, /* tp_repr */
790 0, /* tp_as_number */
791 0, /* tp_as_sequence */
792 0, /* tp_as_mapping */
793 0, /* tp_hash */
794 0, /* tp_call */
795 0, /* tp_str */
796 PyObject_GenericGetAttr, /* tp_getattro */
797 0, /* tp_setattro */
798 0, /* tp_as_buffer */
799 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
800 Py_TPFLAGS_BASETYPE, /* tp_flags */
801 cycle_doc, /* tp_doc */
802 (traverseproc)cycle_traverse, /* tp_traverse */
803 0, /* tp_clear */
804 0, /* tp_richcompare */
805 0, /* tp_weaklistoffset */
806 PyObject_SelfIter, /* tp_iter */
807 (iternextfunc)cycle_next, /* tp_iternext */
808 0, /* tp_methods */
809 0, /* tp_members */
810 0, /* tp_getset */
811 0, /* tp_base */
812 0, /* tp_dict */
813 0, /* tp_descr_get */
814 0, /* tp_descr_set */
815 0, /* tp_dictoffset */
816 0, /* tp_init */
817 0, /* tp_alloc */
818 cycle_new, /* tp_new */
819 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000820};
821
822
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000823/* dropwhile object **********************************************************/
824
825typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000826 PyObject_HEAD
827 PyObject *func;
828 PyObject *it;
829 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000830} dropwhileobject;
831
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000832static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000833
834static PyObject *
835dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
836{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000837 PyObject *func, *seq;
838 PyObject *it;
839 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000840
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000841 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
842 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000843
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000844 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
845 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000846
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000847 /* Get iterator. */
848 it = PyObject_GetIter(seq);
849 if (it == NULL)
850 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000851
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000852 /* create dropwhileobject structure */
853 lz = (dropwhileobject *)type->tp_alloc(type, 0);
854 if (lz == NULL) {
855 Py_DECREF(it);
856 return NULL;
857 }
858 Py_INCREF(func);
859 lz->func = func;
860 lz->it = it;
861 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000862
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000863 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000864}
865
866static void
867dropwhile_dealloc(dropwhileobject *lz)
868{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000869 PyObject_GC_UnTrack(lz);
870 Py_XDECREF(lz->func);
871 Py_XDECREF(lz->it);
872 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000873}
874
875static int
876dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
877{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000878 Py_VISIT(lz->it);
879 Py_VISIT(lz->func);
880 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000881}
882
883static PyObject *
884dropwhile_next(dropwhileobject *lz)
885{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000886 PyObject *item, *good;
887 PyObject *it = lz->it;
888 long ok;
889 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000890
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000891 iternext = *Py_TYPE(it)->tp_iternext;
892 for (;;) {
893 item = iternext(it);
894 if (item == NULL)
895 return NULL;
896 if (lz->start == 1)
897 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000898
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000899 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
900 if (good == NULL) {
901 Py_DECREF(item);
902 return NULL;
903 }
904 ok = PyObject_IsTrue(good);
905 Py_DECREF(good);
906 if (!ok) {
907 lz->start = 1;
908 return item;
909 }
910 Py_DECREF(item);
911 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000912}
913
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000914PyDoc_STRVAR(dropwhile_doc,
915"dropwhile(predicate, iterable) --> dropwhile object\n\
916\n\
917Drop items from the iterable while predicate(item) is true.\n\
918Afterwards, return every element until the iterable is exhausted.");
919
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000920static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000921 PyVarObject_HEAD_INIT(NULL, 0)
922 "itertools.dropwhile", /* tp_name */
923 sizeof(dropwhileobject), /* tp_basicsize */
924 0, /* tp_itemsize */
925 /* methods */
926 (destructor)dropwhile_dealloc, /* tp_dealloc */
927 0, /* tp_print */
928 0, /* tp_getattr */
929 0, /* tp_setattr */
930 0, /* tp_reserved */
931 0, /* tp_repr */
932 0, /* tp_as_number */
933 0, /* tp_as_sequence */
934 0, /* tp_as_mapping */
935 0, /* tp_hash */
936 0, /* tp_call */
937 0, /* tp_str */
938 PyObject_GenericGetAttr, /* tp_getattro */
939 0, /* tp_setattro */
940 0, /* tp_as_buffer */
941 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
942 Py_TPFLAGS_BASETYPE, /* tp_flags */
943 dropwhile_doc, /* tp_doc */
944 (traverseproc)dropwhile_traverse, /* tp_traverse */
945 0, /* tp_clear */
946 0, /* tp_richcompare */
947 0, /* tp_weaklistoffset */
948 PyObject_SelfIter, /* tp_iter */
949 (iternextfunc)dropwhile_next, /* tp_iternext */
950 0, /* tp_methods */
951 0, /* tp_members */
952 0, /* tp_getset */
953 0, /* tp_base */
954 0, /* tp_dict */
955 0, /* tp_descr_get */
956 0, /* tp_descr_set */
957 0, /* tp_dictoffset */
958 0, /* tp_init */
959 0, /* tp_alloc */
960 dropwhile_new, /* tp_new */
961 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000962};
963
964
965/* takewhile object **********************************************************/
966
967typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000968 PyObject_HEAD
969 PyObject *func;
970 PyObject *it;
971 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000972} takewhileobject;
973
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000974static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000975
976static PyObject *
977takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
978{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000979 PyObject *func, *seq;
980 PyObject *it;
981 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000982
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000983 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
984 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000985
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000986 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
987 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000988
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000989 /* Get iterator. */
990 it = PyObject_GetIter(seq);
991 if (it == NULL)
992 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000993
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000994 /* create takewhileobject structure */
995 lz = (takewhileobject *)type->tp_alloc(type, 0);
996 if (lz == NULL) {
997 Py_DECREF(it);
998 return NULL;
999 }
1000 Py_INCREF(func);
1001 lz->func = func;
1002 lz->it = it;
1003 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001005 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001006}
1007
1008static void
1009takewhile_dealloc(takewhileobject *lz)
1010{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001011 PyObject_GC_UnTrack(lz);
1012 Py_XDECREF(lz->func);
1013 Py_XDECREF(lz->it);
1014 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001015}
1016
1017static int
1018takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1019{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001020 Py_VISIT(lz->it);
1021 Py_VISIT(lz->func);
1022 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001023}
1024
1025static PyObject *
1026takewhile_next(takewhileobject *lz)
1027{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001028 PyObject *item, *good;
1029 PyObject *it = lz->it;
1030 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001031
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001032 if (lz->stop == 1)
1033 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001034
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001035 item = (*Py_TYPE(it)->tp_iternext)(it);
1036 if (item == NULL)
1037 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001038
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001039 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1040 if (good == NULL) {
1041 Py_DECREF(item);
1042 return NULL;
1043 }
1044 ok = PyObject_IsTrue(good);
1045 Py_DECREF(good);
1046 if (ok)
1047 return item;
1048 Py_DECREF(item);
1049 lz->stop = 1;
1050 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001051}
1052
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053PyDoc_STRVAR(takewhile_doc,
1054"takewhile(predicate, iterable) --> takewhile object\n\
1055\n\
1056Return successive entries from an iterable as long as the \n\
1057predicate evaluates to true for each entry.");
1058
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001059static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001060 PyVarObject_HEAD_INIT(NULL, 0)
1061 "itertools.takewhile", /* tp_name */
1062 sizeof(takewhileobject), /* tp_basicsize */
1063 0, /* tp_itemsize */
1064 /* methods */
1065 (destructor)takewhile_dealloc, /* tp_dealloc */
1066 0, /* tp_print */
1067 0, /* tp_getattr */
1068 0, /* tp_setattr */
1069 0, /* tp_reserved */
1070 0, /* tp_repr */
1071 0, /* tp_as_number */
1072 0, /* tp_as_sequence */
1073 0, /* tp_as_mapping */
1074 0, /* tp_hash */
1075 0, /* tp_call */
1076 0, /* tp_str */
1077 PyObject_GenericGetAttr, /* tp_getattro */
1078 0, /* tp_setattro */
1079 0, /* tp_as_buffer */
1080 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1081 Py_TPFLAGS_BASETYPE, /* tp_flags */
1082 takewhile_doc, /* tp_doc */
1083 (traverseproc)takewhile_traverse, /* tp_traverse */
1084 0, /* tp_clear */
1085 0, /* tp_richcompare */
1086 0, /* tp_weaklistoffset */
1087 PyObject_SelfIter, /* tp_iter */
1088 (iternextfunc)takewhile_next, /* tp_iternext */
1089 0, /* tp_methods */
1090 0, /* tp_members */
1091 0, /* tp_getset */
1092 0, /* tp_base */
1093 0, /* tp_dict */
1094 0, /* tp_descr_get */
1095 0, /* tp_descr_set */
1096 0, /* tp_dictoffset */
1097 0, /* tp_init */
1098 0, /* tp_alloc */
1099 takewhile_new, /* tp_new */
1100 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001101};
1102
1103
1104/* islice object ************************************************************/
1105
1106typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001107 PyObject_HEAD
1108 PyObject *it;
1109 Py_ssize_t next;
1110 Py_ssize_t stop;
1111 Py_ssize_t step;
1112 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001113} isliceobject;
1114
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001115static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116
1117static PyObject *
1118islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1119{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001120 PyObject *seq;
1121 Py_ssize_t start=0, stop=-1, step=1;
1122 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1123 Py_ssize_t numargs;
1124 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001125
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001126 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1127 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001128
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001129 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1130 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001131
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001132 numargs = PyTuple_Size(args);
1133 if (numargs == 2) {
1134 if (a1 != Py_None) {
1135 stop = PyLong_AsSsize_t(a1);
1136 if (stop == -1) {
1137 if (PyErr_Occurred())
1138 PyErr_Clear();
1139 PyErr_SetString(PyExc_ValueError,
1140 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1141 return NULL;
1142 }
1143 }
1144 } else {
1145 if (a1 != Py_None)
1146 start = PyLong_AsSsize_t(a1);
1147 if (start == -1 && PyErr_Occurred())
1148 PyErr_Clear();
1149 if (a2 != Py_None) {
1150 stop = PyLong_AsSsize_t(a2);
1151 if (stop == -1) {
1152 if (PyErr_Occurred())
1153 PyErr_Clear();
1154 PyErr_SetString(PyExc_ValueError,
1155 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1156 return NULL;
1157 }
1158 }
1159 }
1160 if (start<0 || stop<-1) {
1161 PyErr_SetString(PyExc_ValueError,
1162 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1163 return NULL;
1164 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001165
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001166 if (a3 != NULL) {
1167 if (a3 != Py_None)
1168 step = PyLong_AsSsize_t(a3);
1169 if (step == -1 && PyErr_Occurred())
1170 PyErr_Clear();
1171 }
1172 if (step<1) {
1173 PyErr_SetString(PyExc_ValueError,
1174 "Step for islice() must be a positive integer or None.");
1175 return NULL;
1176 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001177
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001178 /* Get iterator. */
1179 it = PyObject_GetIter(seq);
1180 if (it == NULL)
1181 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001182
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001183 /* create isliceobject structure */
1184 lz = (isliceobject *)type->tp_alloc(type, 0);
1185 if (lz == NULL) {
1186 Py_DECREF(it);
1187 return NULL;
1188 }
1189 lz->it = it;
1190 lz->next = start;
1191 lz->stop = stop;
1192 lz->step = step;
1193 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001194
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001195 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196}
1197
1198static void
1199islice_dealloc(isliceobject *lz)
1200{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001201 PyObject_GC_UnTrack(lz);
1202 Py_XDECREF(lz->it);
1203 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204}
1205
1206static int
1207islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1208{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001209 Py_VISIT(lz->it);
1210 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001211}
1212
1213static PyObject *
1214islice_next(isliceobject *lz)
1215{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001216 PyObject *item;
1217 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001218 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001219 Py_ssize_t oldnext;
1220 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001221
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001222 iternext = *Py_TYPE(it)->tp_iternext;
1223 while (lz->cnt < lz->next) {
1224 item = iternext(it);
1225 if (item == NULL)
1226 return NULL;
1227 Py_DECREF(item);
1228 lz->cnt++;
1229 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001230 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001231 return NULL;
1232 item = iternext(it);
1233 if (item == NULL)
1234 return NULL;
1235 lz->cnt++;
1236 oldnext = lz->next;
1237 lz->next += lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001238 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1239 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001240 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001241}
1242
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001243PyDoc_STRVAR(islice_doc,
1244"islice(iterable, [start,] stop [, step]) --> islice object\n\
1245\n\
1246Return an iterator whose next() method returns selected values from an\n\
1247iterable. If start is specified, will skip all preceding elements;\n\
1248otherwise, start defaults to zero. Step defaults to one. If\n\
1249specified as another value, step determines how many values are \n\
1250skipped between successive calls. Works like a slice() on a list\n\
1251but returns an iterator.");
1252
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001253static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001254 PyVarObject_HEAD_INIT(NULL, 0)
1255 "itertools.islice", /* tp_name */
1256 sizeof(isliceobject), /* tp_basicsize */
1257 0, /* tp_itemsize */
1258 /* methods */
1259 (destructor)islice_dealloc, /* tp_dealloc */
1260 0, /* tp_print */
1261 0, /* tp_getattr */
1262 0, /* tp_setattr */
1263 0, /* tp_reserved */
1264 0, /* tp_repr */
1265 0, /* tp_as_number */
1266 0, /* tp_as_sequence */
1267 0, /* tp_as_mapping */
1268 0, /* tp_hash */
1269 0, /* tp_call */
1270 0, /* tp_str */
1271 PyObject_GenericGetAttr, /* tp_getattro */
1272 0, /* tp_setattro */
1273 0, /* tp_as_buffer */
1274 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1275 Py_TPFLAGS_BASETYPE, /* tp_flags */
1276 islice_doc, /* tp_doc */
1277 (traverseproc)islice_traverse, /* tp_traverse */
1278 0, /* tp_clear */
1279 0, /* tp_richcompare */
1280 0, /* tp_weaklistoffset */
1281 PyObject_SelfIter, /* tp_iter */
1282 (iternextfunc)islice_next, /* tp_iternext */
1283 0, /* tp_methods */
1284 0, /* tp_members */
1285 0, /* tp_getset */
1286 0, /* tp_base */
1287 0, /* tp_dict */
1288 0, /* tp_descr_get */
1289 0, /* tp_descr_set */
1290 0, /* tp_dictoffset */
1291 0, /* tp_init */
1292 0, /* tp_alloc */
1293 islice_new, /* tp_new */
1294 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001295};
1296
1297
1298/* starmap object ************************************************************/
1299
1300typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001301 PyObject_HEAD
1302 PyObject *func;
1303 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001304} starmapobject;
1305
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001306static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001307
1308static PyObject *
1309starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1310{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001311 PyObject *func, *seq;
1312 PyObject *it;
1313 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001314
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001315 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1316 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001317
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001318 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1319 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001320
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001321 /* Get iterator. */
1322 it = PyObject_GetIter(seq);
1323 if (it == NULL)
1324 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001325
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001326 /* create starmapobject structure */
1327 lz = (starmapobject *)type->tp_alloc(type, 0);
1328 if (lz == NULL) {
1329 Py_DECREF(it);
1330 return NULL;
1331 }
1332 Py_INCREF(func);
1333 lz->func = func;
1334 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001335
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001336 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001337}
1338
1339static void
1340starmap_dealloc(starmapobject *lz)
1341{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001342 PyObject_GC_UnTrack(lz);
1343 Py_XDECREF(lz->func);
1344 Py_XDECREF(lz->it);
1345 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001346}
1347
1348static int
1349starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1350{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001351 Py_VISIT(lz->it);
1352 Py_VISIT(lz->func);
1353 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001354}
1355
1356static PyObject *
1357starmap_next(starmapobject *lz)
1358{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001359 PyObject *args;
1360 PyObject *result;
1361 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001362
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001363 args = (*Py_TYPE(it)->tp_iternext)(it);
1364 if (args == NULL)
1365 return NULL;
1366 if (!PyTuple_CheckExact(args)) {
1367 PyObject *newargs = PySequence_Tuple(args);
1368 Py_DECREF(args);
1369 if (newargs == NULL)
1370 return NULL;
1371 args = newargs;
1372 }
1373 result = PyObject_Call(lz->func, args, NULL);
1374 Py_DECREF(args);
1375 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001376}
1377
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001378PyDoc_STRVAR(starmap_doc,
1379"starmap(function, sequence) --> starmap object\n\
1380\n\
1381Return an iterator whose values are returned from the function evaluated\n\
1382with a argument tuple taken from the given sequence.");
1383
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001384static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001385 PyVarObject_HEAD_INIT(NULL, 0)
1386 "itertools.starmap", /* tp_name */
1387 sizeof(starmapobject), /* tp_basicsize */
1388 0, /* tp_itemsize */
1389 /* methods */
1390 (destructor)starmap_dealloc, /* tp_dealloc */
1391 0, /* tp_print */
1392 0, /* tp_getattr */
1393 0, /* tp_setattr */
1394 0, /* tp_reserved */
1395 0, /* tp_repr */
1396 0, /* tp_as_number */
1397 0, /* tp_as_sequence */
1398 0, /* tp_as_mapping */
1399 0, /* tp_hash */
1400 0, /* tp_call */
1401 0, /* tp_str */
1402 PyObject_GenericGetAttr, /* tp_getattro */
1403 0, /* tp_setattro */
1404 0, /* tp_as_buffer */
1405 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1406 Py_TPFLAGS_BASETYPE, /* tp_flags */
1407 starmap_doc, /* tp_doc */
1408 (traverseproc)starmap_traverse, /* tp_traverse */
1409 0, /* tp_clear */
1410 0, /* tp_richcompare */
1411 0, /* tp_weaklistoffset */
1412 PyObject_SelfIter, /* tp_iter */
1413 (iternextfunc)starmap_next, /* tp_iternext */
1414 0, /* tp_methods */
1415 0, /* tp_members */
1416 0, /* tp_getset */
1417 0, /* tp_base */
1418 0, /* tp_dict */
1419 0, /* tp_descr_get */
1420 0, /* tp_descr_set */
1421 0, /* tp_dictoffset */
1422 0, /* tp_init */
1423 0, /* tp_alloc */
1424 starmap_new, /* tp_new */
1425 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001426};
1427
1428
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001429/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001430
1431typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001432 PyObject_HEAD
1433 PyObject *source; /* Iterator over input iterables */
1434 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001435} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001436
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001437static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001438
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001439static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001440chain_new_internal(PyTypeObject *type, PyObject *source)
1441{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001442 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001443
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001444 lz = (chainobject *)type->tp_alloc(type, 0);
1445 if (lz == NULL) {
1446 Py_DECREF(source);
1447 return NULL;
1448 }
1449
1450 lz->source = source;
1451 lz->active = NULL;
1452 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001453}
1454
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001455static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001456chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001458 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001460 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1461 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001462
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001463 source = PyObject_GetIter(args);
1464 if (source == NULL)
1465 return NULL;
1466
1467 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001468}
1469
1470static PyObject *
1471chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1472{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001473 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001474
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001475 source = PyObject_GetIter(arg);
1476 if (source == NULL)
1477 return NULL;
1478
1479 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001480}
1481
1482static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001483chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001484{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001485 PyObject_GC_UnTrack(lz);
1486 Py_XDECREF(lz->active);
1487 Py_XDECREF(lz->source);
1488 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001489}
1490
Raymond Hettinger2012f172003-02-07 05:32:58 +00001491static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001492chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001493{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001494 Py_VISIT(lz->source);
1495 Py_VISIT(lz->active);
1496 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001497}
1498
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001499static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001500chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001501{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001502 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001503
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001504 if (lz->source == NULL)
1505 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001506
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001507 if (lz->active == NULL) {
1508 PyObject *iterable = PyIter_Next(lz->source);
1509 if (iterable == NULL) {
1510 Py_CLEAR(lz->source);
1511 return NULL; /* no more input sources */
1512 }
1513 lz->active = PyObject_GetIter(iterable);
1514 Py_DECREF(iterable);
1515 if (lz->active == NULL) {
1516 Py_CLEAR(lz->source);
1517 return NULL; /* input not iterable */
1518 }
1519 }
1520 item = PyIter_Next(lz->active);
1521 if (item != NULL)
1522 return item;
1523 if (PyErr_Occurred()) {
1524 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1525 PyErr_Clear();
1526 else
1527 return NULL; /* input raised an exception */
1528 }
1529 Py_CLEAR(lz->active);
1530 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001531}
1532
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001533PyDoc_STRVAR(chain_doc,
1534"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001535\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001536Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001537first iterable until it is exhausted, then elements from the next\n\
1538iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001539
Christian Heimesf16baeb2008-02-29 14:57:44 +00001540PyDoc_STRVAR(chain_from_iterable_doc,
1541"chain.from_iterable(iterable) --> chain object\n\
1542\n\
1543Alternate chain() contructor taking a single iterable argument\n\
1544that evaluates lazily.");
1545
1546static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001547 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1548 chain_from_iterable_doc},
1549 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001550};
1551
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001552static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001553 PyVarObject_HEAD_INIT(NULL, 0)
1554 "itertools.chain", /* tp_name */
1555 sizeof(chainobject), /* tp_basicsize */
1556 0, /* tp_itemsize */
1557 /* methods */
1558 (destructor)chain_dealloc, /* tp_dealloc */
1559 0, /* tp_print */
1560 0, /* tp_getattr */
1561 0, /* tp_setattr */
1562 0, /* tp_reserved */
1563 0, /* tp_repr */
1564 0, /* tp_as_number */
1565 0, /* tp_as_sequence */
1566 0, /* tp_as_mapping */
1567 0, /* tp_hash */
1568 0, /* tp_call */
1569 0, /* tp_str */
1570 PyObject_GenericGetAttr, /* tp_getattro */
1571 0, /* tp_setattro */
1572 0, /* tp_as_buffer */
1573 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1574 Py_TPFLAGS_BASETYPE, /* tp_flags */
1575 chain_doc, /* tp_doc */
1576 (traverseproc)chain_traverse, /* tp_traverse */
1577 0, /* tp_clear */
1578 0, /* tp_richcompare */
1579 0, /* tp_weaklistoffset */
1580 PyObject_SelfIter, /* tp_iter */
1581 (iternextfunc)chain_next, /* tp_iternext */
1582 chain_methods, /* tp_methods */
1583 0, /* tp_members */
1584 0, /* tp_getset */
1585 0, /* tp_base */
1586 0, /* tp_dict */
1587 0, /* tp_descr_get */
1588 0, /* tp_descr_set */
1589 0, /* tp_dictoffset */
1590 0, /* tp_init */
1591 0, /* tp_alloc */
1592 chain_new, /* tp_new */
1593 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001594};
1595
1596
Christian Heimesc3f30c42008-02-22 16:37:40 +00001597/* product object ************************************************************/
1598
1599typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001600 PyObject_HEAD
1601 PyObject *pools; /* tuple of pool tuples */
1602 Py_ssize_t *indices; /* one index per pool */
1603 PyObject *result; /* most recently returned result tuple */
1604 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001605} productobject;
1606
1607static PyTypeObject product_type;
1608
1609static PyObject *
1610product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1611{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001612 productobject *lz;
1613 Py_ssize_t nargs, npools, repeat=1;
1614 PyObject *pools = NULL;
1615 Py_ssize_t *indices = NULL;
1616 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001617
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001618 if (kwds != NULL) {
1619 char *kwlist[] = {"repeat", 0};
1620 PyObject *tmpargs = PyTuple_New(0);
1621 if (tmpargs == NULL)
1622 return NULL;
1623 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1624 Py_DECREF(tmpargs);
1625 return NULL;
1626 }
1627 Py_DECREF(tmpargs);
1628 if (repeat < 0) {
1629 PyErr_SetString(PyExc_ValueError,
1630 "repeat argument cannot be negative");
1631 return NULL;
1632 }
1633 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001634
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001635 assert(PyTuple_Check(args));
1636 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1637 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001638
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001639 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1640 if (indices == NULL) {
1641 PyErr_NoMemory();
1642 goto error;
1643 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001644
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001645 pools = PyTuple_New(npools);
1646 if (pools == NULL)
1647 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001648
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001649 for (i=0; i < nargs ; ++i) {
1650 PyObject *item = PyTuple_GET_ITEM(args, i);
1651 PyObject *pool = PySequence_Tuple(item);
1652 if (pool == NULL)
1653 goto error;
1654 PyTuple_SET_ITEM(pools, i, pool);
1655 indices[i] = 0;
1656 }
1657 for ( ; i < npools; ++i) {
1658 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1659 Py_INCREF(pool);
1660 PyTuple_SET_ITEM(pools, i, pool);
1661 indices[i] = 0;
1662 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001663
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001664 /* create productobject structure */
1665 lz = (productobject *)type->tp_alloc(type, 0);
1666 if (lz == NULL)
1667 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001668
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001669 lz->pools = pools;
1670 lz->indices = indices;
1671 lz->result = NULL;
1672 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001673
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001674 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001675
1676error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001677 if (indices != NULL)
1678 PyMem_Free(indices);
1679 Py_XDECREF(pools);
1680 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001681}
1682
1683static void
1684product_dealloc(productobject *lz)
1685{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001686 PyObject_GC_UnTrack(lz);
1687 Py_XDECREF(lz->pools);
1688 Py_XDECREF(lz->result);
1689 if (lz->indices != NULL)
1690 PyMem_Free(lz->indices);
1691 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001692}
1693
1694static int
1695product_traverse(productobject *lz, visitproc visit, void *arg)
1696{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001697 Py_VISIT(lz->pools);
1698 Py_VISIT(lz->result);
1699 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001700}
1701
1702static PyObject *
1703product_next(productobject *lz)
1704{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001705 PyObject *pool;
1706 PyObject *elem;
1707 PyObject *oldelem;
1708 PyObject *pools = lz->pools;
1709 PyObject *result = lz->result;
1710 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1711 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001712
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001713 if (lz->stopped)
1714 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001715
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001716 if (result == NULL) {
1717 /* On the first pass, return an initial tuple filled with the
1718 first element from each pool. */
1719 result = PyTuple_New(npools);
1720 if (result == NULL)
1721 goto empty;
1722 lz->result = result;
1723 for (i=0; i < npools; i++) {
1724 pool = PyTuple_GET_ITEM(pools, i);
1725 if (PyTuple_GET_SIZE(pool) == 0)
1726 goto empty;
1727 elem = PyTuple_GET_ITEM(pool, 0);
1728 Py_INCREF(elem);
1729 PyTuple_SET_ITEM(result, i, elem);
1730 }
1731 } else {
1732 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001733
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001734 /* Copy the previous result tuple or re-use it if available */
1735 if (Py_REFCNT(result) > 1) {
1736 PyObject *old_result = result;
1737 result = PyTuple_New(npools);
1738 if (result == NULL)
1739 goto empty;
1740 lz->result = result;
1741 for (i=0; i < npools; i++) {
1742 elem = PyTuple_GET_ITEM(old_result, i);
1743 Py_INCREF(elem);
1744 PyTuple_SET_ITEM(result, i, elem);
1745 }
1746 Py_DECREF(old_result);
1747 }
1748 /* Now, we've got the only copy so we can update it in-place */
1749 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001750
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001751 /* Update the pool indices right-to-left. Only advance to the
1752 next pool when the previous one rolls-over */
1753 for (i=npools-1 ; i >= 0 ; i--) {
1754 pool = PyTuple_GET_ITEM(pools, i);
1755 indices[i]++;
1756 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1757 /* Roll-over and advance to next pool */
1758 indices[i] = 0;
1759 elem = PyTuple_GET_ITEM(pool, 0);
1760 Py_INCREF(elem);
1761 oldelem = PyTuple_GET_ITEM(result, i);
1762 PyTuple_SET_ITEM(result, i, elem);
1763 Py_DECREF(oldelem);
1764 } else {
1765 /* No rollover. Just increment and stop here. */
1766 elem = PyTuple_GET_ITEM(pool, indices[i]);
1767 Py_INCREF(elem);
1768 oldelem = PyTuple_GET_ITEM(result, i);
1769 PyTuple_SET_ITEM(result, i, elem);
1770 Py_DECREF(oldelem);
1771 break;
1772 }
1773 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001774
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001775 /* If i is negative, then the indices have all rolled-over
1776 and we're done. */
1777 if (i < 0)
1778 goto empty;
1779 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001780
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001781 Py_INCREF(result);
1782 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001783
1784empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001785 lz->stopped = 1;
1786 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001787}
1788
1789PyDoc_STRVAR(product_doc,
1790"product(*iterables) --> product object\n\
1791\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001792Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001793For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1794The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1795cycle in a manner similar to an odometer (with the rightmost element changing\n\
1796on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00001797To compute the product of an iterable with itself, specify the number\n\
1798of repetitions with the optional repeat keyword argument. For example,\n\
1799product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001800product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1801product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1802
1803static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001804 PyVarObject_HEAD_INIT(NULL, 0)
1805 "itertools.product", /* tp_name */
1806 sizeof(productobject), /* tp_basicsize */
1807 0, /* tp_itemsize */
1808 /* methods */
1809 (destructor)product_dealloc, /* tp_dealloc */
1810 0, /* tp_print */
1811 0, /* tp_getattr */
1812 0, /* tp_setattr */
1813 0, /* tp_reserved */
1814 0, /* tp_repr */
1815 0, /* tp_as_number */
1816 0, /* tp_as_sequence */
1817 0, /* tp_as_mapping */
1818 0, /* tp_hash */
1819 0, /* tp_call */
1820 0, /* tp_str */
1821 PyObject_GenericGetAttr, /* tp_getattro */
1822 0, /* tp_setattro */
1823 0, /* tp_as_buffer */
1824 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1825 Py_TPFLAGS_BASETYPE, /* tp_flags */
1826 product_doc, /* tp_doc */
1827 (traverseproc)product_traverse, /* tp_traverse */
1828 0, /* tp_clear */
1829 0, /* tp_richcompare */
1830 0, /* tp_weaklistoffset */
1831 PyObject_SelfIter, /* tp_iter */
1832 (iternextfunc)product_next, /* tp_iternext */
1833 0, /* tp_methods */
1834 0, /* tp_members */
1835 0, /* tp_getset */
1836 0, /* tp_base */
1837 0, /* tp_dict */
1838 0, /* tp_descr_get */
1839 0, /* tp_descr_set */
1840 0, /* tp_dictoffset */
1841 0, /* tp_init */
1842 0, /* tp_alloc */
1843 product_new, /* tp_new */
1844 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001845};
1846
1847
Christian Heimes380f7f22008-02-28 11:19:05 +00001848/* combinations object ************************************************************/
1849
1850typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001851 PyObject_HEAD
1852 PyObject *pool; /* input converted to a tuple */
1853 Py_ssize_t *indices; /* one index per result element */
1854 PyObject *result; /* most recently returned result tuple */
1855 Py_ssize_t r; /* size of result tuple */
1856 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00001857} combinationsobject;
1858
1859static PyTypeObject combinations_type;
1860
1861static PyObject *
1862combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1863{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001864 combinationsobject *co;
1865 Py_ssize_t n;
1866 Py_ssize_t r;
1867 PyObject *pool = NULL;
1868 PyObject *iterable = NULL;
1869 Py_ssize_t *indices = NULL;
1870 Py_ssize_t i;
1871 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00001872
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001873 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
1874 &iterable, &r))
1875 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001876
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001877 pool = PySequence_Tuple(iterable);
1878 if (pool == NULL)
1879 goto error;
1880 n = PyTuple_GET_SIZE(pool);
1881 if (r < 0) {
1882 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
1883 goto error;
1884 }
Christian Heimes380f7f22008-02-28 11:19:05 +00001885
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001886 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
1887 if (indices == NULL) {
1888 PyErr_NoMemory();
1889 goto error;
1890 }
Christian Heimes380f7f22008-02-28 11:19:05 +00001891
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001892 for (i=0 ; i<r ; i++)
1893 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00001894
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001895 /* create combinationsobject structure */
1896 co = (combinationsobject *)type->tp_alloc(type, 0);
1897 if (co == NULL)
1898 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00001899
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001900 co->pool = pool;
1901 co->indices = indices;
1902 co->result = NULL;
1903 co->r = r;
1904 co->stopped = r > n ? 1 : 0;
1905
1906 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00001907
1908error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001909 if (indices != NULL)
1910 PyMem_Free(indices);
1911 Py_XDECREF(pool);
1912 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001913}
1914
1915static void
1916combinations_dealloc(combinationsobject *co)
1917{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001918 PyObject_GC_UnTrack(co);
1919 Py_XDECREF(co->pool);
1920 Py_XDECREF(co->result);
1921 if (co->indices != NULL)
1922 PyMem_Free(co->indices);
1923 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00001924}
1925
1926static int
1927combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
1928{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001929 Py_VISIT(co->pool);
1930 Py_VISIT(co->result);
1931 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00001932}
1933
1934static PyObject *
1935combinations_next(combinationsobject *co)
1936{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001937 PyObject *elem;
1938 PyObject *oldelem;
1939 PyObject *pool = co->pool;
1940 Py_ssize_t *indices = co->indices;
1941 PyObject *result = co->result;
1942 Py_ssize_t n = PyTuple_GET_SIZE(pool);
1943 Py_ssize_t r = co->r;
1944 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00001945
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001946 if (co->stopped)
1947 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001948
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001949 if (result == NULL) {
1950 /* On the first pass, initialize result tuple using the indices */
1951 result = PyTuple_New(r);
1952 if (result == NULL)
1953 goto empty;
1954 co->result = result;
1955 for (i=0; i<r ; i++) {
1956 index = indices[i];
1957 elem = PyTuple_GET_ITEM(pool, index);
1958 Py_INCREF(elem);
1959 PyTuple_SET_ITEM(result, i, elem);
1960 }
1961 } else {
1962 /* Copy the previous result tuple or re-use it if available */
1963 if (Py_REFCNT(result) > 1) {
1964 PyObject *old_result = result;
1965 result = PyTuple_New(r);
1966 if (result == NULL)
1967 goto empty;
1968 co->result = result;
1969 for (i=0; i<r ; i++) {
1970 elem = PyTuple_GET_ITEM(old_result, i);
1971 Py_INCREF(elem);
1972 PyTuple_SET_ITEM(result, i, elem);
1973 }
1974 Py_DECREF(old_result);
1975 }
1976 /* Now, we've got the only copy so we can update it in-place
1977 * CPython's empty tuple is a singleton and cached in
1978 * PyTuple's freelist.
1979 */
1980 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00001981
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001982 /* Scan indices right-to-left until finding one that is not
1983 at its maximum (i + n - r). */
1984 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
1985 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00001986
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001987 /* If i is negative, then the indices are all at
1988 their maximum value and we're done. */
1989 if (i < 0)
1990 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00001991
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001992 /* Increment the current index which we know is not at its
1993 maximum. Then move back to the right setting each index
1994 to its lowest possible value (one higher than the index
1995 to its left -- this maintains the sort order invariant). */
1996 indices[i]++;
1997 for (j=i+1 ; j<r ; j++)
1998 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00001999
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002000 /* Update the result tuple for the new indices
2001 starting with i, the leftmost index that changed */
2002 for ( ; i<r ; i++) {
2003 index = indices[i];
2004 elem = PyTuple_GET_ITEM(pool, index);
2005 Py_INCREF(elem);
2006 oldelem = PyTuple_GET_ITEM(result, i);
2007 PyTuple_SET_ITEM(result, i, elem);
2008 Py_DECREF(oldelem);
2009 }
2010 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002011
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002012 Py_INCREF(result);
2013 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002014
2015empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002016 co->stopped = 1;
2017 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002018}
2019
2020PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002021"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002022\n\
2023Return successive r-length combinations of elements in the iterable.\n\n\
2024combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2025
2026static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002027 PyVarObject_HEAD_INIT(NULL, 0)
2028 "itertools.combinations", /* tp_name */
2029 sizeof(combinationsobject), /* tp_basicsize */
2030 0, /* tp_itemsize */
2031 /* methods */
2032 (destructor)combinations_dealloc, /* tp_dealloc */
2033 0, /* tp_print */
2034 0, /* tp_getattr */
2035 0, /* tp_setattr */
2036 0, /* tp_reserved */
2037 0, /* tp_repr */
2038 0, /* tp_as_number */
2039 0, /* tp_as_sequence */
2040 0, /* tp_as_mapping */
2041 0, /* tp_hash */
2042 0, /* tp_call */
2043 0, /* tp_str */
2044 PyObject_GenericGetAttr, /* tp_getattro */
2045 0, /* tp_setattro */
2046 0, /* tp_as_buffer */
2047 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2048 Py_TPFLAGS_BASETYPE, /* tp_flags */
2049 combinations_doc, /* tp_doc */
2050 (traverseproc)combinations_traverse, /* tp_traverse */
2051 0, /* tp_clear */
2052 0, /* tp_richcompare */
2053 0, /* tp_weaklistoffset */
2054 PyObject_SelfIter, /* tp_iter */
2055 (iternextfunc)combinations_next, /* tp_iternext */
2056 0, /* tp_methods */
2057 0, /* tp_members */
2058 0, /* tp_getset */
2059 0, /* tp_base */
2060 0, /* tp_dict */
2061 0, /* tp_descr_get */
2062 0, /* tp_descr_set */
2063 0, /* tp_dictoffset */
2064 0, /* tp_init */
2065 0, /* tp_alloc */
2066 combinations_new, /* tp_new */
2067 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002068};
2069
2070
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002071/* combinations with replacement object *******************************************/
2072
2073/* Equivalent to:
2074
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002075 def combinations_with_replacement(iterable, r):
2076 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2077 # number items returned: (n+r-1)! / r! / (n-1)!
2078 pool = tuple(iterable)
2079 n = len(pool)
2080 indices = [0] * r
2081 yield tuple(pool[i] for i in indices)
2082 while 1:
2083 for i in reversed(range(r)):
2084 if indices[i] != n - 1:
2085 break
2086 else:
2087 return
2088 indices[i:] = [indices[i] + 1] * (r - i)
2089 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002090
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002091 def combinations_with_replacement2(iterable, r):
2092 'Alternate version that filters from product()'
2093 pool = tuple(iterable)
2094 n = len(pool)
2095 for indices in product(range(n), repeat=r):
2096 if sorted(indices) == list(indices):
2097 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002098*/
2099typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002100 PyObject_HEAD
2101 PyObject *pool; /* input converted to a tuple */
2102 Py_ssize_t *indices; /* one index per result element */
2103 PyObject *result; /* most recently returned result tuple */
2104 Py_ssize_t r; /* size of result tuple */
2105 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002106} cwrobject;
2107
2108static PyTypeObject cwr_type;
2109
2110static PyObject *
2111cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2112{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002113 cwrobject *co;
2114 Py_ssize_t n;
2115 Py_ssize_t r;
2116 PyObject *pool = NULL;
2117 PyObject *iterable = NULL;
2118 Py_ssize_t *indices = NULL;
2119 Py_ssize_t i;
2120 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002121
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002122 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2123 &iterable, &r))
2124 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002125
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002126 pool = PySequence_Tuple(iterable);
2127 if (pool == NULL)
2128 goto error;
2129 n = PyTuple_GET_SIZE(pool);
2130 if (r < 0) {
2131 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2132 goto error;
2133 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002134
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002135 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2136 if (indices == NULL) {
2137 PyErr_NoMemory();
2138 goto error;
2139 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002140
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002141 for (i=0 ; i<r ; i++)
2142 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002143
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002144 /* create cwrobject structure */
2145 co = (cwrobject *)type->tp_alloc(type, 0);
2146 if (co == NULL)
2147 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002148
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002149 co->pool = pool;
2150 co->indices = indices;
2151 co->result = NULL;
2152 co->r = r;
2153 co->stopped = !n && r;
2154
2155 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002156
2157error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002158 if (indices != NULL)
2159 PyMem_Free(indices);
2160 Py_XDECREF(pool);
2161 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002162}
2163
2164static void
2165cwr_dealloc(cwrobject *co)
2166{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002167 PyObject_GC_UnTrack(co);
2168 Py_XDECREF(co->pool);
2169 Py_XDECREF(co->result);
2170 if (co->indices != NULL)
2171 PyMem_Free(co->indices);
2172 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002173}
2174
2175static int
2176cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2177{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002178 Py_VISIT(co->pool);
2179 Py_VISIT(co->result);
2180 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002181}
2182
2183static PyObject *
2184cwr_next(cwrobject *co)
2185{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002186 PyObject *elem;
2187 PyObject *oldelem;
2188 PyObject *pool = co->pool;
2189 Py_ssize_t *indices = co->indices;
2190 PyObject *result = co->result;
2191 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2192 Py_ssize_t r = co->r;
2193 Py_ssize_t i, j, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002194
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002195 if (co->stopped)
2196 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002197
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002198 if (result == NULL) {
2199 /* On the first pass, initialize result tuple using the indices */
2200 result = PyTuple_New(r);
2201 if (result == NULL)
2202 goto empty;
2203 co->result = result;
2204 for (i=0; i<r ; i++) {
2205 index = indices[i];
2206 elem = PyTuple_GET_ITEM(pool, index);
2207 Py_INCREF(elem);
2208 PyTuple_SET_ITEM(result, i, elem);
2209 }
2210 } else {
2211 /* Copy the previous result tuple or re-use it if available */
2212 if (Py_REFCNT(result) > 1) {
2213 PyObject *old_result = result;
2214 result = PyTuple_New(r);
2215 if (result == NULL)
2216 goto empty;
2217 co->result = result;
2218 for (i=0; i<r ; i++) {
2219 elem = PyTuple_GET_ITEM(old_result, i);
2220 Py_INCREF(elem);
2221 PyTuple_SET_ITEM(result, i, elem);
2222 }
2223 Py_DECREF(old_result);
2224 }
2225 /* Now, we've got the only copy so we can update it in-place CPython's
2226 empty tuple is a singleton and cached in PyTuple's freelist. */
2227 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002228
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002229 /* Scan indices right-to-left until finding one that is not
2230 * at its maximum (n-1). */
2231 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2232 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002233
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002234 /* If i is negative, then the indices are all at
2235 their maximum value and we're done. */
2236 if (i < 0)
2237 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002238
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002239 /* Increment the current index which we know is not at its
2240 maximum. Then set all to the right to the same value. */
2241 indices[i]++;
2242 for (j=i+1 ; j<r ; j++)
2243 indices[j] = indices[j-1];
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002244
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002245 /* Update the result tuple for the new indices
2246 starting with i, the leftmost index that changed */
2247 for ( ; i<r ; i++) {
2248 index = indices[i];
2249 elem = PyTuple_GET_ITEM(pool, index);
2250 Py_INCREF(elem);
2251 oldelem = PyTuple_GET_ITEM(result, i);
2252 PyTuple_SET_ITEM(result, i, elem);
2253 Py_DECREF(oldelem);
2254 }
2255 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002256
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002257 Py_INCREF(result);
2258 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002259
2260empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002261 co->stopped = 1;
2262 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002263}
2264
2265PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002266"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002267\n\
2268Return successive r-length combinations of elements in the iterable\n\
2269allowing individual elements to have successive repeats.\n\
2270combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2271
2272static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002273 PyVarObject_HEAD_INIT(NULL, 0)
2274 "itertools.combinations_with_replacement", /* tp_name */
2275 sizeof(cwrobject), /* tp_basicsize */
2276 0, /* tp_itemsize */
2277 /* methods */
2278 (destructor)cwr_dealloc, /* tp_dealloc */
2279 0, /* tp_print */
2280 0, /* tp_getattr */
2281 0, /* tp_setattr */
2282 0, /* tp_reserved */
2283 0, /* tp_repr */
2284 0, /* tp_as_number */
2285 0, /* tp_as_sequence */
2286 0, /* tp_as_mapping */
2287 0, /* tp_hash */
2288 0, /* tp_call */
2289 0, /* tp_str */
2290 PyObject_GenericGetAttr, /* tp_getattro */
2291 0, /* tp_setattro */
2292 0, /* tp_as_buffer */
2293 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2294 Py_TPFLAGS_BASETYPE, /* tp_flags */
2295 cwr_doc, /* tp_doc */
2296 (traverseproc)cwr_traverse, /* tp_traverse */
2297 0, /* tp_clear */
2298 0, /* tp_richcompare */
2299 0, /* tp_weaklistoffset */
2300 PyObject_SelfIter, /* tp_iter */
2301 (iternextfunc)cwr_next, /* tp_iternext */
2302 0, /* tp_methods */
2303 0, /* tp_members */
2304 0, /* tp_getset */
2305 0, /* tp_base */
2306 0, /* tp_dict */
2307 0, /* tp_descr_get */
2308 0, /* tp_descr_set */
2309 0, /* tp_dictoffset */
2310 0, /* tp_init */
2311 0, /* tp_alloc */
2312 cwr_new, /* tp_new */
2313 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002314};
2315
2316
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002317/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002318
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002319def permutations(iterable, r=None):
2320 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2321 pool = tuple(iterable)
2322 n = len(pool)
2323 r = n if r is None else r
2324 indices = range(n)
2325 cycles = range(n-r+1, n+1)[::-1]
2326 yield tuple(pool[i] for i in indices[:r])
2327 while n:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002328 for i in reversed(range(r)):
2329 cycles[i] -= 1
2330 if cycles[i] == 0:
2331 indices[i:] = indices[i+1:] + indices[i:i+1]
2332 cycles[i] = n - i
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002333 else:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002334 j = cycles[i]
2335 indices[i], indices[-j] = indices[-j], indices[i]
2336 yield tuple(pool[i] for i in indices[:r])
2337 break
2338 else:
2339 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002340*/
2341
2342typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002343 PyObject_HEAD
2344 PyObject *pool; /* input converted to a tuple */
2345 Py_ssize_t *indices; /* one index per element in the pool */
2346 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2347 PyObject *result; /* most recently returned result tuple */
2348 Py_ssize_t r; /* size of result tuple */
2349 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002350} permutationsobject;
2351
2352static PyTypeObject permutations_type;
2353
2354static PyObject *
2355permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2356{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002357 permutationsobject *po;
2358 Py_ssize_t n;
2359 Py_ssize_t r;
2360 PyObject *robj = Py_None;
2361 PyObject *pool = NULL;
2362 PyObject *iterable = NULL;
2363 Py_ssize_t *indices = NULL;
2364 Py_ssize_t *cycles = NULL;
2365 Py_ssize_t i;
2366 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002367
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002368 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2369 &iterable, &robj))
2370 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002371
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002372 pool = PySequence_Tuple(iterable);
2373 if (pool == NULL)
2374 goto error;
2375 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002376
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002377 r = n;
2378 if (robj != Py_None) {
2379 if (!PyLong_Check(robj)) {
2380 PyErr_SetString(PyExc_TypeError, "Expected int as r");
2381 goto error;
2382 }
2383 r = PyLong_AsSsize_t(robj);
2384 if (r == -1 && PyErr_Occurred())
2385 goto error;
2386 }
2387 if (r < 0) {
2388 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2389 goto error;
2390 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002391
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002392 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2393 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2394 if (indices == NULL || cycles == NULL) {
2395 PyErr_NoMemory();
2396 goto error;
2397 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002398
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002399 for (i=0 ; i<n ; i++)
2400 indices[i] = i;
2401 for (i=0 ; i<r ; i++)
2402 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002403
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002404 /* create permutationsobject structure */
2405 po = (permutationsobject *)type->tp_alloc(type, 0);
2406 if (po == NULL)
2407 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002408
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002409 po->pool = pool;
2410 po->indices = indices;
2411 po->cycles = cycles;
2412 po->result = NULL;
2413 po->r = r;
2414 po->stopped = r > n ? 1 : 0;
2415
2416 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002417
2418error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002419 if (indices != NULL)
2420 PyMem_Free(indices);
2421 if (cycles != NULL)
2422 PyMem_Free(cycles);
2423 Py_XDECREF(pool);
2424 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002425}
2426
2427static void
2428permutations_dealloc(permutationsobject *po)
2429{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002430 PyObject_GC_UnTrack(po);
2431 Py_XDECREF(po->pool);
2432 Py_XDECREF(po->result);
2433 PyMem_Free(po->indices);
2434 PyMem_Free(po->cycles);
2435 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002436}
2437
2438static int
2439permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2440{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002441 Py_VISIT(po->pool);
2442 Py_VISIT(po->result);
2443 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002444}
2445
2446static PyObject *
2447permutations_next(permutationsobject *po)
2448{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002449 PyObject *elem;
2450 PyObject *oldelem;
2451 PyObject *pool = po->pool;
2452 Py_ssize_t *indices = po->indices;
2453 Py_ssize_t *cycles = po->cycles;
2454 PyObject *result = po->result;
2455 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2456 Py_ssize_t r = po->r;
2457 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002458
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002459 if (po->stopped)
2460 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002461
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002462 if (result == NULL) {
2463 /* On the first pass, initialize result tuple using the indices */
2464 result = PyTuple_New(r);
2465 if (result == NULL)
2466 goto empty;
2467 po->result = result;
2468 for (i=0; i<r ; i++) {
2469 index = indices[i];
2470 elem = PyTuple_GET_ITEM(pool, index);
2471 Py_INCREF(elem);
2472 PyTuple_SET_ITEM(result, i, elem);
2473 }
2474 } else {
2475 if (n == 0)
2476 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002477
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002478 /* Copy the previous result tuple or re-use it if available */
2479 if (Py_REFCNT(result) > 1) {
2480 PyObject *old_result = result;
2481 result = PyTuple_New(r);
2482 if (result == NULL)
2483 goto empty;
2484 po->result = result;
2485 for (i=0; i<r ; i++) {
2486 elem = PyTuple_GET_ITEM(old_result, i);
2487 Py_INCREF(elem);
2488 PyTuple_SET_ITEM(result, i, elem);
2489 }
2490 Py_DECREF(old_result);
2491 }
2492 /* Now, we've got the only copy so we can update it in-place */
2493 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002494
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002495 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2496 for (i=r-1 ; i>=0 ; i--) {
2497 cycles[i] -= 1;
2498 if (cycles[i] == 0) {
2499 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2500 index = indices[i];
2501 for (j=i ; j<n-1 ; j++)
2502 indices[j] = indices[j+1];
2503 indices[n-1] = index;
2504 cycles[i] = n - i;
2505 } else {
2506 j = cycles[i];
2507 index = indices[i];
2508 indices[i] = indices[n-j];
2509 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002510
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002511 for (k=i; k<r ; k++) {
2512 /* start with i, the leftmost element that changed */
2513 /* yield tuple(pool[k] for k in indices[:r]) */
2514 index = indices[k];
2515 elem = PyTuple_GET_ITEM(pool, index);
2516 Py_INCREF(elem);
2517 oldelem = PyTuple_GET_ITEM(result, k);
2518 PyTuple_SET_ITEM(result, k, elem);
2519 Py_DECREF(oldelem);
2520 }
2521 break;
2522 }
2523 }
2524 /* If i is negative, then the cycles have all
2525 rolled-over and we're done. */
2526 if (i < 0)
2527 goto empty;
2528 }
2529 Py_INCREF(result);
2530 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002531
2532empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002533 po->stopped = 1;
2534 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002535}
2536
2537PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00002538"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002539\n\
2540Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00002541permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002542
2543static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002544 PyVarObject_HEAD_INIT(NULL, 0)
2545 "itertools.permutations", /* tp_name */
2546 sizeof(permutationsobject), /* tp_basicsize */
2547 0, /* tp_itemsize */
2548 /* methods */
2549 (destructor)permutations_dealloc, /* tp_dealloc */
2550 0, /* tp_print */
2551 0, /* tp_getattr */
2552 0, /* tp_setattr */
2553 0, /* tp_reserved */
2554 0, /* tp_repr */
2555 0, /* tp_as_number */
2556 0, /* tp_as_sequence */
2557 0, /* tp_as_mapping */
2558 0, /* tp_hash */
2559 0, /* tp_call */
2560 0, /* tp_str */
2561 PyObject_GenericGetAttr, /* tp_getattro */
2562 0, /* tp_setattro */
2563 0, /* tp_as_buffer */
2564 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2565 Py_TPFLAGS_BASETYPE, /* tp_flags */
2566 permutations_doc, /* tp_doc */
2567 (traverseproc)permutations_traverse, /* tp_traverse */
2568 0, /* tp_clear */
2569 0, /* tp_richcompare */
2570 0, /* tp_weaklistoffset */
2571 PyObject_SelfIter, /* tp_iter */
2572 (iternextfunc)permutations_next, /* tp_iternext */
2573 0, /* tp_methods */
2574 0, /* tp_members */
2575 0, /* tp_getset */
2576 0, /* tp_base */
2577 0, /* tp_dict */
2578 0, /* tp_descr_get */
2579 0, /* tp_descr_set */
2580 0, /* tp_dictoffset */
2581 0, /* tp_init */
2582 0, /* tp_alloc */
2583 permutations_new, /* tp_new */
2584 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002585};
2586
Raymond Hettinger482ba772010-12-01 22:48:00 +00002587/* accumulate object ************************************************************/
2588
2589typedef struct {
2590 PyObject_HEAD
2591 PyObject *total;
2592 PyObject *it;
2593} accumulateobject;
2594
2595static PyTypeObject accumulate_type;
2596
2597static PyObject *
2598accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2599{
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002600 static char *kwargs[] = {"iterable", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00002601 PyObject *iterable;
2602 PyObject *it;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002603 accumulateobject *lz;
2604
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002605 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O:accumulate", kwargs, &iterable))
2606 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002607
2608 /* Get iterator. */
2609 it = PyObject_GetIter(iterable);
2610 if (it == NULL)
2611 return NULL;
2612
Raymond Hettinger482ba772010-12-01 22:48:00 +00002613 /* create accumulateobject structure */
2614 lz = (accumulateobject *)type->tp_alloc(type, 0);
2615 if (lz == NULL) {
2616 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002617 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002618 }
2619
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002620 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002621 lz->it = it;
2622 return (PyObject *)lz;
2623}
2624
2625static void
2626accumulate_dealloc(accumulateobject *lz)
2627{
2628 PyObject_GC_UnTrack(lz);
2629 Py_XDECREF(lz->total);
2630 Py_XDECREF(lz->it);
2631 Py_TYPE(lz)->tp_free(lz);
2632}
2633
2634static int
2635accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
2636{
2637 Py_VISIT(lz->it);
2638 Py_VISIT(lz->total);
2639 return 0;
2640}
2641
2642static PyObject *
2643accumulate_next(accumulateobject *lz)
2644{
2645 PyObject *val, *oldtotal, *newtotal;
2646
2647 val = PyIter_Next(lz->it);
2648 if (val == NULL)
2649 return NULL;
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002650
2651 if (lz->total == NULL) {
2652 Py_INCREF(val);
2653 lz->total = val;
2654 return lz->total;
2655 }
2656
Raymond Hettinger482ba772010-12-01 22:48:00 +00002657 newtotal = PyNumber_Add(lz->total, val);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002658 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00002659 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002660 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002661
2662 oldtotal = lz->total;
2663 lz->total = newtotal;
2664 Py_DECREF(oldtotal);
2665
2666 Py_INCREF(newtotal);
2667 return newtotal;
2668}
2669
2670PyDoc_STRVAR(accumulate_doc,
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002671"accumulate(iterable) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00002672\n\
2673Return series of accumulated sums.");
2674
2675static PyTypeObject accumulate_type = {
2676 PyVarObject_HEAD_INIT(NULL, 0)
2677 "itertools.accumulate", /* tp_name */
2678 sizeof(accumulateobject), /* tp_basicsize */
2679 0, /* tp_itemsize */
2680 /* methods */
2681 (destructor)accumulate_dealloc, /* tp_dealloc */
2682 0, /* tp_print */
2683 0, /* tp_getattr */
2684 0, /* tp_setattr */
2685 0, /* tp_reserved */
2686 0, /* tp_repr */
2687 0, /* tp_as_number */
2688 0, /* tp_as_sequence */
2689 0, /* tp_as_mapping */
2690 0, /* tp_hash */
2691 0, /* tp_call */
2692 0, /* tp_str */
2693 PyObject_GenericGetAttr, /* tp_getattro */
2694 0, /* tp_setattro */
2695 0, /* tp_as_buffer */
2696 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2697 Py_TPFLAGS_BASETYPE, /* tp_flags */
2698 accumulate_doc, /* tp_doc */
2699 (traverseproc)accumulate_traverse, /* tp_traverse */
2700 0, /* tp_clear */
2701 0, /* tp_richcompare */
2702 0, /* tp_weaklistoffset */
2703 PyObject_SelfIter, /* tp_iter */
2704 (iternextfunc)accumulate_next, /* tp_iternext */
2705 0, /* tp_methods */
2706 0, /* tp_members */
2707 0, /* tp_getset */
2708 0, /* tp_base */
2709 0, /* tp_dict */
2710 0, /* tp_descr_get */
2711 0, /* tp_descr_set */
2712 0, /* tp_dictoffset */
2713 0, /* tp_init */
2714 0, /* tp_alloc */
2715 accumulate_new, /* tp_new */
2716 PyObject_GC_Del, /* tp_free */
2717};
2718
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002719
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002720/* compress object ************************************************************/
2721
2722/* Equivalent to:
2723
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002724 def compress(data, selectors):
2725 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2726 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002727*/
2728
2729typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002730 PyObject_HEAD
2731 PyObject *data;
2732 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002733} compressobject;
2734
2735static PyTypeObject compress_type;
2736
2737static PyObject *
2738compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2739{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002740 PyObject *seq1, *seq2;
2741 PyObject *data=NULL, *selectors=NULL;
2742 compressobject *lz;
2743 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002744
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002745 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2746 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002747
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002748 data = PyObject_GetIter(seq1);
2749 if (data == NULL)
2750 goto fail;
2751 selectors = PyObject_GetIter(seq2);
2752 if (selectors == NULL)
2753 goto fail;
2754
2755 /* create compressobject structure */
2756 lz = (compressobject *)type->tp_alloc(type, 0);
2757 if (lz == NULL)
2758 goto fail;
2759 lz->data = data;
2760 lz->selectors = selectors;
2761 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002762
2763fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002764 Py_XDECREF(data);
2765 Py_XDECREF(selectors);
2766 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002767}
2768
2769static void
2770compress_dealloc(compressobject *lz)
2771{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002772 PyObject_GC_UnTrack(lz);
2773 Py_XDECREF(lz->data);
2774 Py_XDECREF(lz->selectors);
2775 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002776}
2777
2778static int
2779compress_traverse(compressobject *lz, visitproc visit, void *arg)
2780{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002781 Py_VISIT(lz->data);
2782 Py_VISIT(lz->selectors);
2783 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002784}
2785
2786static PyObject *
2787compress_next(compressobject *lz)
2788{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002789 PyObject *data = lz->data, *selectors = lz->selectors;
2790 PyObject *datum, *selector;
2791 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2792 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2793 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002794
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002795 while (1) {
2796 /* Steps: get datum, get selector, evaluate selector.
2797 Order is important (to match the pure python version
2798 in terms of which input gets a chance to raise an
2799 exception first).
2800 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002801
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002802 datum = datanext(data);
2803 if (datum == NULL)
2804 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002805
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002806 selector = selectornext(selectors);
2807 if (selector == NULL) {
2808 Py_DECREF(datum);
2809 return NULL;
2810 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002811
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002812 ok = PyObject_IsTrue(selector);
2813 Py_DECREF(selector);
2814 if (ok == 1)
2815 return datum;
2816 Py_DECREF(datum);
2817 if (ok == -1)
2818 return NULL;
2819 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002820}
2821
2822PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00002823"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002824\n\
2825Return data elements corresponding to true selector elements.\n\
2826Forms a shorter iterator from selected data elements using the\n\
2827selectors to choose the data elements.");
2828
2829static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002830 PyVarObject_HEAD_INIT(NULL, 0)
2831 "itertools.compress", /* tp_name */
2832 sizeof(compressobject), /* tp_basicsize */
2833 0, /* tp_itemsize */
2834 /* methods */
2835 (destructor)compress_dealloc, /* tp_dealloc */
2836 0, /* tp_print */
2837 0, /* tp_getattr */
2838 0, /* tp_setattr */
2839 0, /* tp_reserved */
2840 0, /* tp_repr */
2841 0, /* tp_as_number */
2842 0, /* tp_as_sequence */
2843 0, /* tp_as_mapping */
2844 0, /* tp_hash */
2845 0, /* tp_call */
2846 0, /* tp_str */
2847 PyObject_GenericGetAttr, /* tp_getattro */
2848 0, /* tp_setattro */
2849 0, /* tp_as_buffer */
2850 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2851 Py_TPFLAGS_BASETYPE, /* tp_flags */
2852 compress_doc, /* tp_doc */
2853 (traverseproc)compress_traverse, /* tp_traverse */
2854 0, /* tp_clear */
2855 0, /* tp_richcompare */
2856 0, /* tp_weaklistoffset */
2857 PyObject_SelfIter, /* tp_iter */
2858 (iternextfunc)compress_next, /* tp_iternext */
2859 0, /* tp_methods */
2860 0, /* tp_members */
2861 0, /* tp_getset */
2862 0, /* tp_base */
2863 0, /* tp_dict */
2864 0, /* tp_descr_get */
2865 0, /* tp_descr_set */
2866 0, /* tp_dictoffset */
2867 0, /* tp_init */
2868 0, /* tp_alloc */
2869 compress_new, /* tp_new */
2870 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002871};
2872
2873
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002874/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00002875
2876typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002877 PyObject_HEAD
2878 PyObject *func;
2879 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002880} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002881
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002882static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002883
2884static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002885filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002886{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002887 PyObject *func, *seq;
2888 PyObject *it;
2889 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002890
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002891 if (type == &filterfalse_type &&
2892 !_PyArg_NoKeywords("filterfalse()", kwds))
2893 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002894
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002895 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
2896 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002897
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002898 /* Get iterator. */
2899 it = PyObject_GetIter(seq);
2900 if (it == NULL)
2901 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002902
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002903 /* create filterfalseobject structure */
2904 lz = (filterfalseobject *)type->tp_alloc(type, 0);
2905 if (lz == NULL) {
2906 Py_DECREF(it);
2907 return NULL;
2908 }
2909 Py_INCREF(func);
2910 lz->func = func;
2911 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002912
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002913 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002914}
2915
2916static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002917filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002918{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002919 PyObject_GC_UnTrack(lz);
2920 Py_XDECREF(lz->func);
2921 Py_XDECREF(lz->it);
2922 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002923}
2924
2925static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002926filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002927{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002928 Py_VISIT(lz->it);
2929 Py_VISIT(lz->func);
2930 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002931}
2932
2933static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002934filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002935{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002936 PyObject *item;
2937 PyObject *it = lz->it;
2938 long ok;
2939 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002940
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002941 iternext = *Py_TYPE(it)->tp_iternext;
2942 for (;;) {
2943 item = iternext(it);
2944 if (item == NULL)
2945 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002946
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002947 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2948 ok = PyObject_IsTrue(item);
2949 } else {
2950 PyObject *good;
2951 good = PyObject_CallFunctionObjArgs(lz->func,
2952 item, NULL);
2953 if (good == NULL) {
2954 Py_DECREF(item);
2955 return NULL;
2956 }
2957 ok = PyObject_IsTrue(good);
2958 Py_DECREF(good);
2959 }
2960 if (!ok)
2961 return item;
2962 Py_DECREF(item);
2963 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00002964}
2965
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002966PyDoc_STRVAR(filterfalse_doc,
2967"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002968\n\
2969Return those items of sequence for which function(item) is false.\n\
2970If function is None, return the items that are false.");
2971
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002972static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002973 PyVarObject_HEAD_INIT(NULL, 0)
2974 "itertools.filterfalse", /* tp_name */
2975 sizeof(filterfalseobject), /* tp_basicsize */
2976 0, /* tp_itemsize */
2977 /* methods */
2978 (destructor)filterfalse_dealloc, /* tp_dealloc */
2979 0, /* tp_print */
2980 0, /* tp_getattr */
2981 0, /* tp_setattr */
2982 0, /* tp_reserved */
2983 0, /* tp_repr */
2984 0, /* tp_as_number */
2985 0, /* tp_as_sequence */
2986 0, /* tp_as_mapping */
2987 0, /* tp_hash */
2988 0, /* tp_call */
2989 0, /* tp_str */
2990 PyObject_GenericGetAttr, /* tp_getattro */
2991 0, /* tp_setattro */
2992 0, /* tp_as_buffer */
2993 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2994 Py_TPFLAGS_BASETYPE, /* tp_flags */
2995 filterfalse_doc, /* tp_doc */
2996 (traverseproc)filterfalse_traverse, /* tp_traverse */
2997 0, /* tp_clear */
2998 0, /* tp_richcompare */
2999 0, /* tp_weaklistoffset */
3000 PyObject_SelfIter, /* tp_iter */
3001 (iternextfunc)filterfalse_next, /* tp_iternext */
3002 0, /* tp_methods */
3003 0, /* tp_members */
3004 0, /* tp_getset */
3005 0, /* tp_base */
3006 0, /* tp_dict */
3007 0, /* tp_descr_get */
3008 0, /* tp_descr_set */
3009 0, /* tp_dictoffset */
3010 0, /* tp_init */
3011 0, /* tp_alloc */
3012 filterfalse_new, /* tp_new */
3013 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003014};
3015
3016
3017/* count object ************************************************************/
3018
3019typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003020 PyObject_HEAD
3021 Py_ssize_t cnt;
3022 PyObject *long_cnt;
3023 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003024} countobject;
3025
Raymond Hettinger30729212009-02-12 06:28:27 +00003026/* Counting logic and invariants:
3027
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003028fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003029
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003030 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3031 Advances with: cnt += 1
3032 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003033
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003034slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003035
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003036 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3037 All counting is done with python objects (no overflows or underflows).
3038 Advances with: long_cnt += long_step
3039 Step may be zero -- effectively a slow version of repeat(cnt).
3040 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003041*/
3042
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003043static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003044
3045static PyObject *
3046count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3047{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003048 countobject *lz;
3049 int slow_mode = 0;
3050 Py_ssize_t cnt = 0;
3051 PyObject *long_cnt = NULL;
3052 PyObject *long_step = NULL;
3053 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003054
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003055 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3056 kwlist, &long_cnt, &long_step))
3057 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003058
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003059 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3060 (long_step != NULL && !PyNumber_Check(long_step))) {
3061 PyErr_SetString(PyExc_TypeError, "a number is required");
3062 return NULL;
3063 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003064
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003065 if (long_cnt != NULL) {
3066 cnt = PyLong_AsSsize_t(long_cnt);
3067 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3068 PyErr_Clear();
3069 slow_mode = 1;
3070 }
3071 Py_INCREF(long_cnt);
3072 } else {
3073 cnt = 0;
3074 long_cnt = PyLong_FromLong(0);
3075 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003076
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003077 /* If not specified, step defaults to 1 */
3078 if (long_step == NULL) {
3079 long_step = PyLong_FromLong(1);
3080 if (long_step == NULL) {
3081 Py_DECREF(long_cnt);
3082 return NULL;
3083 }
3084 } else
3085 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003086
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003087 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003088
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003089 /* Fast mode only works when the step is 1 */
3090 if (!PyLong_Check(long_step) ||
3091 PyLong_AS_LONG(long_step) != 1) {
3092 slow_mode = 1;
3093 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003094
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003095 if (slow_mode)
3096 cnt = PY_SSIZE_T_MAX;
3097 else
3098 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003099
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003100 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3101 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3102 assert(slow_mode ||
3103 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003104
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003105 /* create countobject structure */
3106 lz = (countobject *)type->tp_alloc(type, 0);
3107 if (lz == NULL) {
3108 Py_XDECREF(long_cnt);
3109 return NULL;
3110 }
3111 lz->cnt = cnt;
3112 lz->long_cnt = long_cnt;
3113 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003114
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003115 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003116}
3117
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003118static void
3119count_dealloc(countobject *lz)
3120{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003121 PyObject_GC_UnTrack(lz);
3122 Py_XDECREF(lz->long_cnt);
3123 Py_XDECREF(lz->long_step);
3124 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003125}
3126
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003127static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003128count_traverse(countobject *lz, visitproc visit, void *arg)
3129{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003130 Py_VISIT(lz->long_cnt);
3131 Py_VISIT(lz->long_step);
3132 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003133}
3134
3135static PyObject *
3136count_nextlong(countobject *lz)
3137{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003138 PyObject *long_cnt;
3139 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003140
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003141 long_cnt = lz->long_cnt;
3142 if (long_cnt == NULL) {
3143 /* Switch to slow_mode */
3144 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3145 if (long_cnt == NULL)
3146 return NULL;
3147 }
3148 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00003149
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003150 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3151 if (stepped_up == NULL)
3152 return NULL;
3153 lz->long_cnt = stepped_up;
3154 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003155}
3156
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003157static PyObject *
3158count_next(countobject *lz)
3159{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003160 if (lz->cnt == PY_SSIZE_T_MAX)
3161 return count_nextlong(lz);
3162 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003163}
3164
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003165static PyObject *
3166count_repr(countobject *lz)
3167{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003168 if (lz->cnt != PY_SSIZE_T_MAX)
3169 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003170
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003171 if (PyLong_Check(lz->long_step)) {
3172 long step = PyLong_AsLong(lz->long_step);
3173 if (step == -1 && PyErr_Occurred()) {
3174 PyErr_Clear();
3175 }
3176 if (step == 1) {
3177 /* Don't display step when it is an integer equal to 1 */
3178 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
3179 }
3180 }
3181 return PyUnicode_FromFormat("count(%R, %R)",
3182 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003183}
3184
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003185static PyObject *
3186count_reduce(countobject *lz)
3187{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003188 if (lz->cnt == PY_SSIZE_T_MAX)
3189 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3190 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003191}
3192
3193PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3194
3195static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003196 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3197 count_reduce_doc},
3198 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003199};
3200
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003201PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00003202 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003203\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003204Return a count object whose .__next__() method returns consecutive values.\n\
3205Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00003206 def count(firstval=0, step=1):\n\
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003207 x = firstval\n\
3208 while 1:\n\
3209 yield x\n\
3210 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003211
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003212static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003213 PyVarObject_HEAD_INIT(NULL, 0)
3214 "itertools.count", /* tp_name */
3215 sizeof(countobject), /* tp_basicsize */
3216 0, /* tp_itemsize */
3217 /* methods */
3218 (destructor)count_dealloc, /* tp_dealloc */
3219 0, /* tp_print */
3220 0, /* tp_getattr */
3221 0, /* tp_setattr */
3222 0, /* tp_reserved */
3223 (reprfunc)count_repr, /* tp_repr */
3224 0, /* tp_as_number */
3225 0, /* tp_as_sequence */
3226 0, /* tp_as_mapping */
3227 0, /* tp_hash */
3228 0, /* tp_call */
3229 0, /* tp_str */
3230 PyObject_GenericGetAttr, /* tp_getattro */
3231 0, /* tp_setattro */
3232 0, /* tp_as_buffer */
3233 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3234 Py_TPFLAGS_BASETYPE, /* tp_flags */
3235 count_doc, /* tp_doc */
3236 (traverseproc)count_traverse, /* tp_traverse */
3237 0, /* tp_clear */
3238 0, /* tp_richcompare */
3239 0, /* tp_weaklistoffset */
3240 PyObject_SelfIter, /* tp_iter */
3241 (iternextfunc)count_next, /* tp_iternext */
3242 count_methods, /* tp_methods */
3243 0, /* tp_members */
3244 0, /* tp_getset */
3245 0, /* tp_base */
3246 0, /* tp_dict */
3247 0, /* tp_descr_get */
3248 0, /* tp_descr_set */
3249 0, /* tp_dictoffset */
3250 0, /* tp_init */
3251 0, /* tp_alloc */
3252 count_new, /* tp_new */
3253 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003254};
3255
3256
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003257/* repeat object ************************************************************/
3258
3259typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003260 PyObject_HEAD
3261 PyObject *element;
3262 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003263} repeatobject;
3264
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003265static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003266
3267static PyObject *
3268repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3269{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003270 repeatobject *ro;
3271 PyObject *element;
3272 Py_ssize_t cnt = -1;
3273 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003274
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003275 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3276 &element, &cnt))
3277 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003278
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003279 if (PyTuple_Size(args) == 2 && cnt < 0)
3280 cnt = 0;
3281
3282 ro = (repeatobject *)type->tp_alloc(type, 0);
3283 if (ro == NULL)
3284 return NULL;
3285 Py_INCREF(element);
3286 ro->element = element;
3287 ro->cnt = cnt;
3288 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003289}
3290
3291static void
3292repeat_dealloc(repeatobject *ro)
3293{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003294 PyObject_GC_UnTrack(ro);
3295 Py_XDECREF(ro->element);
3296 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003297}
3298
3299static int
3300repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3301{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003302 Py_VISIT(ro->element);
3303 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003304}
3305
3306static PyObject *
3307repeat_next(repeatobject *ro)
3308{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003309 if (ro->cnt == 0)
3310 return NULL;
3311 if (ro->cnt > 0)
3312 ro->cnt--;
3313 Py_INCREF(ro->element);
3314 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003315}
3316
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003317static PyObject *
3318repeat_repr(repeatobject *ro)
3319{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003320 if (ro->cnt == -1)
3321 return PyUnicode_FromFormat("repeat(%R)", ro->element);
3322 else
3323 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
3324}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003325
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003326static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003327repeat_len(repeatobject *ro)
3328{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003329 if (ro->cnt == -1) {
3330 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3331 return NULL;
3332 }
3333 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003334}
3335
Armin Rigof5b3e362006-02-11 21:32:43 +00003336PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003337
3338static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003339 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3340 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003341};
3342
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003343PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00003344"repeat(object [,times]) -> create an iterator which returns the object\n\
3345for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003346endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003347
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003348static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003349 PyVarObject_HEAD_INIT(NULL, 0)
3350 "itertools.repeat", /* tp_name */
3351 sizeof(repeatobject), /* tp_basicsize */
3352 0, /* tp_itemsize */
3353 /* methods */
3354 (destructor)repeat_dealloc, /* tp_dealloc */
3355 0, /* tp_print */
3356 0, /* tp_getattr */
3357 0, /* tp_setattr */
3358 0, /* tp_reserved */
3359 (reprfunc)repeat_repr, /* tp_repr */
3360 0, /* tp_as_number */
3361 0, /* tp_as_sequence */
3362 0, /* tp_as_mapping */
3363 0, /* tp_hash */
3364 0, /* tp_call */
3365 0, /* tp_str */
3366 PyObject_GenericGetAttr, /* tp_getattro */
3367 0, /* tp_setattro */
3368 0, /* tp_as_buffer */
3369 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3370 Py_TPFLAGS_BASETYPE, /* tp_flags */
3371 repeat_doc, /* tp_doc */
3372 (traverseproc)repeat_traverse, /* tp_traverse */
3373 0, /* tp_clear */
3374 0, /* tp_richcompare */
3375 0, /* tp_weaklistoffset */
3376 PyObject_SelfIter, /* tp_iter */
3377 (iternextfunc)repeat_next, /* tp_iternext */
3378 repeat_methods, /* tp_methods */
3379 0, /* tp_members */
3380 0, /* tp_getset */
3381 0, /* tp_base */
3382 0, /* tp_dict */
3383 0, /* tp_descr_get */
3384 0, /* tp_descr_set */
3385 0, /* tp_dictoffset */
3386 0, /* tp_init */
3387 0, /* tp_alloc */
3388 repeat_new, /* tp_new */
3389 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003390};
3391
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003392/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00003393
3394#include "Python.h"
3395
3396typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003397 PyObject_HEAD
3398 Py_ssize_t tuplesize;
3399 Py_ssize_t numactive;
3400 PyObject *ittuple; /* tuple of iterators */
3401 PyObject *result;
3402 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003403} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003404
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003405static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003406
3407static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003408zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003409{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003410 ziplongestobject *lz;
3411 Py_ssize_t i;
3412 PyObject *ittuple; /* tuple of iterators */
3413 PyObject *result;
3414 PyObject *fillvalue = Py_None;
3415 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003416
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003417 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3418 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3419 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3420 PyErr_SetString(PyExc_TypeError,
3421 "zip_longest() got an unexpected keyword argument");
3422 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003423 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003424 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003425
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003426 /* args must be a tuple */
3427 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00003428
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003429 /* obtain iterators */
3430 ittuple = PyTuple_New(tuplesize);
3431 if (ittuple == NULL)
3432 return NULL;
3433 for (i=0; i < tuplesize; ++i) {
3434 PyObject *item = PyTuple_GET_ITEM(args, i);
3435 PyObject *it = PyObject_GetIter(item);
3436 if (it == NULL) {
3437 if (PyErr_ExceptionMatches(PyExc_TypeError))
3438 PyErr_Format(PyExc_TypeError,
3439 "zip_longest argument #%zd must support iteration",
3440 i+1);
3441 Py_DECREF(ittuple);
3442 return NULL;
3443 }
3444 PyTuple_SET_ITEM(ittuple, i, it);
3445 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003446
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003447 /* create a result holder */
3448 result = PyTuple_New(tuplesize);
3449 if (result == NULL) {
3450 Py_DECREF(ittuple);
3451 return NULL;
3452 }
3453 for (i=0 ; i < tuplesize ; i++) {
3454 Py_INCREF(Py_None);
3455 PyTuple_SET_ITEM(result, i, Py_None);
3456 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003457
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003458 /* create ziplongestobject structure */
3459 lz = (ziplongestobject *)type->tp_alloc(type, 0);
3460 if (lz == NULL) {
3461 Py_DECREF(ittuple);
3462 Py_DECREF(result);
3463 return NULL;
3464 }
3465 lz->ittuple = ittuple;
3466 lz->tuplesize = tuplesize;
3467 lz->numactive = tuplesize;
3468 lz->result = result;
3469 Py_INCREF(fillvalue);
3470 lz->fillvalue = fillvalue;
3471 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003472}
3473
3474static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003475zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003476{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003477 PyObject_GC_UnTrack(lz);
3478 Py_XDECREF(lz->ittuple);
3479 Py_XDECREF(lz->result);
3480 Py_XDECREF(lz->fillvalue);
3481 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003482}
3483
3484static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003485zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003486{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003487 Py_VISIT(lz->ittuple);
3488 Py_VISIT(lz->result);
3489 Py_VISIT(lz->fillvalue);
3490 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003491}
3492
3493static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003494zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003495{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003496 Py_ssize_t i;
3497 Py_ssize_t tuplesize = lz->tuplesize;
3498 PyObject *result = lz->result;
3499 PyObject *it;
3500 PyObject *item;
3501 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003502
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003503 if (tuplesize == 0)
3504 return NULL;
3505 if (lz->numactive == 0)
3506 return NULL;
3507 if (Py_REFCNT(result) == 1) {
3508 Py_INCREF(result);
3509 for (i=0 ; i < tuplesize ; i++) {
3510 it = PyTuple_GET_ITEM(lz->ittuple, i);
3511 if (it == NULL) {
3512 Py_INCREF(lz->fillvalue);
3513 item = lz->fillvalue;
3514 } else {
3515 item = PyIter_Next(it);
3516 if (item == NULL) {
3517 lz->numactive -= 1;
3518 if (lz->numactive == 0 || PyErr_Occurred()) {
3519 lz->numactive = 0;
3520 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00003521 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003522 } else {
3523 Py_INCREF(lz->fillvalue);
3524 item = lz->fillvalue;
3525 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3526 Py_DECREF(it);
3527 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00003528 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003529 }
3530 olditem = PyTuple_GET_ITEM(result, i);
3531 PyTuple_SET_ITEM(result, i, item);
3532 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00003533 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003534 } else {
3535 result = PyTuple_New(tuplesize);
3536 if (result == NULL)
3537 return NULL;
3538 for (i=0 ; i < tuplesize ; i++) {
3539 it = PyTuple_GET_ITEM(lz->ittuple, i);
3540 if (it == NULL) {
3541 Py_INCREF(lz->fillvalue);
3542 item = lz->fillvalue;
3543 } else {
3544 item = PyIter_Next(it);
3545 if (item == NULL) {
3546 lz->numactive -= 1;
3547 if (lz->numactive == 0 || PyErr_Occurred()) {
3548 lz->numactive = 0;
3549 Py_DECREF(result);
3550 return NULL;
3551 } else {
3552 Py_INCREF(lz->fillvalue);
3553 item = lz->fillvalue;
3554 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3555 Py_DECREF(it);
3556 }
3557 }
3558 }
3559 PyTuple_SET_ITEM(result, i, item);
3560 }
3561 }
3562 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003563}
3564
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003565PyDoc_STRVAR(zip_longest_doc,
3566"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003567\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003568Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00003569the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003570method continues until the longest iterable in the argument sequence\n\
3571is exhausted and then it raises StopIteration. When the shorter iterables\n\
3572are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3573defaults to None or can be specified by a keyword argument.\n\
3574");
3575
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003576static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003577 PyVarObject_HEAD_INIT(NULL, 0)
3578 "itertools.zip_longest", /* tp_name */
3579 sizeof(ziplongestobject), /* tp_basicsize */
3580 0, /* tp_itemsize */
3581 /* methods */
3582 (destructor)zip_longest_dealloc, /* tp_dealloc */
3583 0, /* tp_print */
3584 0, /* tp_getattr */
3585 0, /* tp_setattr */
3586 0, /* tp_reserved */
3587 0, /* tp_repr */
3588 0, /* tp_as_number */
3589 0, /* tp_as_sequence */
3590 0, /* tp_as_mapping */
3591 0, /* tp_hash */
3592 0, /* tp_call */
3593 0, /* tp_str */
3594 PyObject_GenericGetAttr, /* tp_getattro */
3595 0, /* tp_setattro */
3596 0, /* tp_as_buffer */
3597 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3598 Py_TPFLAGS_BASETYPE, /* tp_flags */
3599 zip_longest_doc, /* tp_doc */
3600 (traverseproc)zip_longest_traverse, /* tp_traverse */
3601 0, /* tp_clear */
3602 0, /* tp_richcompare */
3603 0, /* tp_weaklistoffset */
3604 PyObject_SelfIter, /* tp_iter */
3605 (iternextfunc)zip_longest_next, /* tp_iternext */
3606 0, /* tp_methods */
3607 0, /* tp_members */
3608 0, /* tp_getset */
3609 0, /* tp_base */
3610 0, /* tp_dict */
3611 0, /* tp_descr_get */
3612 0, /* tp_descr_set */
3613 0, /* tp_dictoffset */
3614 0, /* tp_init */
3615 0, /* tp_alloc */
3616 zip_longest_new, /* tp_new */
3617 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003618};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003619
3620/* module level code ********************************************************/
3621
3622PyDoc_STRVAR(module_doc,
3623"Functional tools for creating and using iterators.\n\
3624\n\
3625Infinite iterators:\n\
3626count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003627cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003628repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003629\n\
3630Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003631accumulate(p, start=0) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003632chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
3633compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
3634dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
3635groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003636filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003637islice(seq, [start,] stop [, step]) --> elements from\n\
3638 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003639starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003640tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003641takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003642zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00003643\n\
3644Combinatoric generators:\n\
3645product(p, q, ... [repeat=1]) --> cartesian product\n\
3646permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00003647combinations(p, r)\n\
3648combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003649");
3650
3651
Raymond Hettingerad983e72003-11-12 14:32:26 +00003652static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003653 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3654 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00003655};
3656
Martin v. Löwis1a214512008-06-11 05:26:20 +00003657
3658static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003659 PyModuleDef_HEAD_INIT,
3660 "itertools",
3661 module_doc,
3662 -1,
3663 module_methods,
3664 NULL,
3665 NULL,
3666 NULL,
3667 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00003668};
3669
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003670PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00003671PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003672{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003673 int i;
3674 PyObject *m;
3675 char *name;
3676 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00003677 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003678 &combinations_type,
3679 &cwr_type,
3680 &cycle_type,
3681 &dropwhile_type,
3682 &takewhile_type,
3683 &islice_type,
3684 &starmap_type,
3685 &chain_type,
3686 &compress_type,
3687 &filterfalse_type,
3688 &count_type,
3689 &ziplongest_type,
3690 &permutations_type,
3691 &product_type,
3692 &repeat_type,
3693 &groupby_type,
3694 NULL
3695 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00003696
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003697 Py_TYPE(&teedataobject_type) = &PyType_Type;
3698 m = PyModule_Create(&itertoolsmodule);
3699 if (m == NULL)
3700 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003701
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003702 for (i=0 ; typelist[i] != NULL ; i++) {
3703 if (PyType_Ready(typelist[i]) < 0)
3704 return NULL;
3705 name = strchr(typelist[i]->tp_name, '.');
3706 assert (name != NULL);
3707 Py_INCREF(typelist[i]);
3708 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
3709 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003710
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003711 if (PyType_Ready(&teedataobject_type) < 0)
3712 return NULL;
3713 if (PyType_Ready(&tee_type) < 0)
3714 return NULL;
3715 if (PyType_Ready(&_grouper_type) < 0)
3716 return NULL;
3717 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003718}