blob: 15b0c17973d25312669c548582ed02a933a889ab [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;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200629 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000630
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000631 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
632 return NULL;
633 if (n < 0) {
634 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
635 return NULL;
636 }
637 result = PyTuple_New(n);
638 if (result == NULL)
639 return NULL;
640 if (n == 0)
641 return result;
642 it = PyObject_GetIter(iterable);
643 if (it == NULL) {
644 Py_DECREF(result);
645 return NULL;
646 }
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200647 if (!_PyObject_HasAttrId(it, &PyId___copy__)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000648 copyable = tee_fromiterable(it);
649 Py_DECREF(it);
650 if (copyable == NULL) {
651 Py_DECREF(result);
652 return NULL;
653 }
654 } else
655 copyable = it;
656 PyTuple_SET_ITEM(result, 0, copyable);
657 for (i=1 ; i<n ; i++) {
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200658
659 copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000660 if (copyable == NULL) {
661 Py_DECREF(result);
662 return NULL;
663 }
664 PyTuple_SET_ITEM(result, i, copyable);
665 }
666 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000667}
668
669PyDoc_STRVAR(tee_doc,
670"tee(iterable, n=2) --> tuple of n independent iterators.");
671
672
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000673/* cycle object **********************************************************/
674
675typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000676 PyObject_HEAD
677 PyObject *it;
678 PyObject *saved;
679 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000680} cycleobject;
681
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000682static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000683
684static PyObject *
685cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
686{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000687 PyObject *it;
688 PyObject *iterable;
689 PyObject *saved;
690 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000691
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000692 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
693 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000694
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000695 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
696 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000697
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000698 /* Get iterator. */
699 it = PyObject_GetIter(iterable);
700 if (it == NULL)
701 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000702
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000703 saved = PyList_New(0);
704 if (saved == NULL) {
705 Py_DECREF(it);
706 return NULL;
707 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000708
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000709 /* create cycleobject structure */
710 lz = (cycleobject *)type->tp_alloc(type, 0);
711 if (lz == NULL) {
712 Py_DECREF(it);
713 Py_DECREF(saved);
714 return NULL;
715 }
716 lz->it = it;
717 lz->saved = saved;
718 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000719
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000720 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000721}
722
723static void
724cycle_dealloc(cycleobject *lz)
725{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000726 PyObject_GC_UnTrack(lz);
727 Py_XDECREF(lz->saved);
728 Py_XDECREF(lz->it);
729 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000730}
731
732static int
733cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
734{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000735 Py_VISIT(lz->it);
736 Py_VISIT(lz->saved);
737 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000738}
739
740static PyObject *
741cycle_next(cycleobject *lz)
742{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000743 PyObject *item;
744 PyObject *it;
745 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000746
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000747 while (1) {
748 item = PyIter_Next(lz->it);
749 if (item != NULL) {
750 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
751 Py_DECREF(item);
752 return NULL;
753 }
754 return item;
755 }
756 if (PyErr_Occurred()) {
757 if (PyErr_ExceptionMatches(PyExc_StopIteration))
758 PyErr_Clear();
759 else
760 return NULL;
761 }
762 if (PyList_Size(lz->saved) == 0)
763 return NULL;
764 it = PyObject_GetIter(lz->saved);
765 if (it == NULL)
766 return NULL;
767 tmp = lz->it;
768 lz->it = it;
769 lz->firstpass = 1;
770 Py_DECREF(tmp);
771 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000772}
773
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000774PyDoc_STRVAR(cycle_doc,
775"cycle(iterable) --> cycle object\n\
776\n\
777Return elements from the iterable until it is exhausted.\n\
778Then repeat the sequence indefinitely.");
779
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000780static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000781 PyVarObject_HEAD_INIT(NULL, 0)
782 "itertools.cycle", /* tp_name */
783 sizeof(cycleobject), /* tp_basicsize */
784 0, /* tp_itemsize */
785 /* methods */
786 (destructor)cycle_dealloc, /* tp_dealloc */
787 0, /* tp_print */
788 0, /* tp_getattr */
789 0, /* tp_setattr */
790 0, /* tp_reserved */
791 0, /* tp_repr */
792 0, /* tp_as_number */
793 0, /* tp_as_sequence */
794 0, /* tp_as_mapping */
795 0, /* tp_hash */
796 0, /* tp_call */
797 0, /* tp_str */
798 PyObject_GenericGetAttr, /* tp_getattro */
799 0, /* tp_setattro */
800 0, /* tp_as_buffer */
801 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
802 Py_TPFLAGS_BASETYPE, /* tp_flags */
803 cycle_doc, /* tp_doc */
804 (traverseproc)cycle_traverse, /* tp_traverse */
805 0, /* tp_clear */
806 0, /* tp_richcompare */
807 0, /* tp_weaklistoffset */
808 PyObject_SelfIter, /* tp_iter */
809 (iternextfunc)cycle_next, /* tp_iternext */
810 0, /* tp_methods */
811 0, /* tp_members */
812 0, /* tp_getset */
813 0, /* tp_base */
814 0, /* tp_dict */
815 0, /* tp_descr_get */
816 0, /* tp_descr_set */
817 0, /* tp_dictoffset */
818 0, /* tp_init */
819 0, /* tp_alloc */
820 cycle_new, /* tp_new */
821 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000822};
823
824
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000825/* dropwhile object **********************************************************/
826
827typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000828 PyObject_HEAD
829 PyObject *func;
830 PyObject *it;
831 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000832} dropwhileobject;
833
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000834static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000835
836static PyObject *
837dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
838{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000839 PyObject *func, *seq;
840 PyObject *it;
841 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000842
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000843 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
844 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000845
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000846 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
847 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000848
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000849 /* Get iterator. */
850 it = PyObject_GetIter(seq);
851 if (it == NULL)
852 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000853
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000854 /* create dropwhileobject structure */
855 lz = (dropwhileobject *)type->tp_alloc(type, 0);
856 if (lz == NULL) {
857 Py_DECREF(it);
858 return NULL;
859 }
860 Py_INCREF(func);
861 lz->func = func;
862 lz->it = it;
863 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000864
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000865 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000866}
867
868static void
869dropwhile_dealloc(dropwhileobject *lz)
870{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000871 PyObject_GC_UnTrack(lz);
872 Py_XDECREF(lz->func);
873 Py_XDECREF(lz->it);
874 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000875}
876
877static int
878dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
879{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000880 Py_VISIT(lz->it);
881 Py_VISIT(lz->func);
882 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000883}
884
885static PyObject *
886dropwhile_next(dropwhileobject *lz)
887{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000888 PyObject *item, *good;
889 PyObject *it = lz->it;
890 long ok;
891 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000892
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000893 iternext = *Py_TYPE(it)->tp_iternext;
894 for (;;) {
895 item = iternext(it);
896 if (item == NULL)
897 return NULL;
898 if (lz->start == 1)
899 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000900
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000901 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
902 if (good == NULL) {
903 Py_DECREF(item);
904 return NULL;
905 }
906 ok = PyObject_IsTrue(good);
907 Py_DECREF(good);
908 if (!ok) {
909 lz->start = 1;
910 return item;
911 }
912 Py_DECREF(item);
913 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000914}
915
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000916PyDoc_STRVAR(dropwhile_doc,
917"dropwhile(predicate, iterable) --> dropwhile object\n\
918\n\
919Drop items from the iterable while predicate(item) is true.\n\
920Afterwards, return every element until the iterable is exhausted.");
921
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000922static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000923 PyVarObject_HEAD_INIT(NULL, 0)
924 "itertools.dropwhile", /* tp_name */
925 sizeof(dropwhileobject), /* tp_basicsize */
926 0, /* tp_itemsize */
927 /* methods */
928 (destructor)dropwhile_dealloc, /* tp_dealloc */
929 0, /* tp_print */
930 0, /* tp_getattr */
931 0, /* tp_setattr */
932 0, /* tp_reserved */
933 0, /* tp_repr */
934 0, /* tp_as_number */
935 0, /* tp_as_sequence */
936 0, /* tp_as_mapping */
937 0, /* tp_hash */
938 0, /* tp_call */
939 0, /* tp_str */
940 PyObject_GenericGetAttr, /* tp_getattro */
941 0, /* tp_setattro */
942 0, /* tp_as_buffer */
943 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
944 Py_TPFLAGS_BASETYPE, /* tp_flags */
945 dropwhile_doc, /* tp_doc */
946 (traverseproc)dropwhile_traverse, /* tp_traverse */
947 0, /* tp_clear */
948 0, /* tp_richcompare */
949 0, /* tp_weaklistoffset */
950 PyObject_SelfIter, /* tp_iter */
951 (iternextfunc)dropwhile_next, /* tp_iternext */
952 0, /* tp_methods */
953 0, /* tp_members */
954 0, /* tp_getset */
955 0, /* tp_base */
956 0, /* tp_dict */
957 0, /* tp_descr_get */
958 0, /* tp_descr_set */
959 0, /* tp_dictoffset */
960 0, /* tp_init */
961 0, /* tp_alloc */
962 dropwhile_new, /* tp_new */
963 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000964};
965
966
967/* takewhile object **********************************************************/
968
969typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000970 PyObject_HEAD
971 PyObject *func;
972 PyObject *it;
973 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000974} takewhileobject;
975
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000976static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000977
978static PyObject *
979takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
980{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000981 PyObject *func, *seq;
982 PyObject *it;
983 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000984
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000985 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
986 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000987
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000988 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
989 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000990
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000991 /* Get iterator. */
992 it = PyObject_GetIter(seq);
993 if (it == NULL)
994 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000995
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000996 /* create takewhileobject structure */
997 lz = (takewhileobject *)type->tp_alloc(type, 0);
998 if (lz == NULL) {
999 Py_DECREF(it);
1000 return NULL;
1001 }
1002 Py_INCREF(func);
1003 lz->func = func;
1004 lz->it = it;
1005 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001006
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001007 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001008}
1009
1010static void
1011takewhile_dealloc(takewhileobject *lz)
1012{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001013 PyObject_GC_UnTrack(lz);
1014 Py_XDECREF(lz->func);
1015 Py_XDECREF(lz->it);
1016 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001017}
1018
1019static int
1020takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1021{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001022 Py_VISIT(lz->it);
1023 Py_VISIT(lz->func);
1024 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001025}
1026
1027static PyObject *
1028takewhile_next(takewhileobject *lz)
1029{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001030 PyObject *item, *good;
1031 PyObject *it = lz->it;
1032 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001033
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001034 if (lz->stop == 1)
1035 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001036
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001037 item = (*Py_TYPE(it)->tp_iternext)(it);
1038 if (item == NULL)
1039 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001040
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001041 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1042 if (good == NULL) {
1043 Py_DECREF(item);
1044 return NULL;
1045 }
1046 ok = PyObject_IsTrue(good);
1047 Py_DECREF(good);
1048 if (ok)
1049 return item;
1050 Py_DECREF(item);
1051 lz->stop = 1;
1052 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053}
1054
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001055PyDoc_STRVAR(takewhile_doc,
1056"takewhile(predicate, iterable) --> takewhile object\n\
1057\n\
1058Return successive entries from an iterable as long as the \n\
1059predicate evaluates to true for each entry.");
1060
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001061static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001062 PyVarObject_HEAD_INIT(NULL, 0)
1063 "itertools.takewhile", /* tp_name */
1064 sizeof(takewhileobject), /* tp_basicsize */
1065 0, /* tp_itemsize */
1066 /* methods */
1067 (destructor)takewhile_dealloc, /* tp_dealloc */
1068 0, /* tp_print */
1069 0, /* tp_getattr */
1070 0, /* tp_setattr */
1071 0, /* tp_reserved */
1072 0, /* tp_repr */
1073 0, /* tp_as_number */
1074 0, /* tp_as_sequence */
1075 0, /* tp_as_mapping */
1076 0, /* tp_hash */
1077 0, /* tp_call */
1078 0, /* tp_str */
1079 PyObject_GenericGetAttr, /* tp_getattro */
1080 0, /* tp_setattro */
1081 0, /* tp_as_buffer */
1082 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1083 Py_TPFLAGS_BASETYPE, /* tp_flags */
1084 takewhile_doc, /* tp_doc */
1085 (traverseproc)takewhile_traverse, /* tp_traverse */
1086 0, /* tp_clear */
1087 0, /* tp_richcompare */
1088 0, /* tp_weaklistoffset */
1089 PyObject_SelfIter, /* tp_iter */
1090 (iternextfunc)takewhile_next, /* tp_iternext */
1091 0, /* tp_methods */
1092 0, /* tp_members */
1093 0, /* tp_getset */
1094 0, /* tp_base */
1095 0, /* tp_dict */
1096 0, /* tp_descr_get */
1097 0, /* tp_descr_set */
1098 0, /* tp_dictoffset */
1099 0, /* tp_init */
1100 0, /* tp_alloc */
1101 takewhile_new, /* tp_new */
1102 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001103};
1104
1105
1106/* islice object ************************************************************/
1107
1108typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001109 PyObject_HEAD
1110 PyObject *it;
1111 Py_ssize_t next;
1112 Py_ssize_t stop;
1113 Py_ssize_t step;
1114 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001115} isliceobject;
1116
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001117static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001118
1119static PyObject *
1120islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1121{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001122 PyObject *seq;
1123 Py_ssize_t start=0, stop=-1, step=1;
1124 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1125 Py_ssize_t numargs;
1126 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001127
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001128 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1129 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001130
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001131 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1132 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001133
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001134 numargs = PyTuple_Size(args);
1135 if (numargs == 2) {
1136 if (a1 != Py_None) {
1137 stop = PyLong_AsSsize_t(a1);
1138 if (stop == -1) {
1139 if (PyErr_Occurred())
1140 PyErr_Clear();
1141 PyErr_SetString(PyExc_ValueError,
1142 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1143 return NULL;
1144 }
1145 }
1146 } else {
1147 if (a1 != Py_None)
1148 start = PyLong_AsSsize_t(a1);
1149 if (start == -1 && PyErr_Occurred())
1150 PyErr_Clear();
1151 if (a2 != Py_None) {
1152 stop = PyLong_AsSsize_t(a2);
1153 if (stop == -1) {
1154 if (PyErr_Occurred())
1155 PyErr_Clear();
1156 PyErr_SetString(PyExc_ValueError,
1157 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1158 return NULL;
1159 }
1160 }
1161 }
1162 if (start<0 || stop<-1) {
1163 PyErr_SetString(PyExc_ValueError,
1164 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1165 return NULL;
1166 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001167
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001168 if (a3 != NULL) {
1169 if (a3 != Py_None)
1170 step = PyLong_AsSsize_t(a3);
1171 if (step == -1 && PyErr_Occurred())
1172 PyErr_Clear();
1173 }
1174 if (step<1) {
1175 PyErr_SetString(PyExc_ValueError,
1176 "Step for islice() must be a positive integer or None.");
1177 return NULL;
1178 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001179
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001180 /* Get iterator. */
1181 it = PyObject_GetIter(seq);
1182 if (it == NULL)
1183 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001184
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001185 /* create isliceobject structure */
1186 lz = (isliceobject *)type->tp_alloc(type, 0);
1187 if (lz == NULL) {
1188 Py_DECREF(it);
1189 return NULL;
1190 }
1191 lz->it = it;
1192 lz->next = start;
1193 lz->stop = stop;
1194 lz->step = step;
1195 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001197 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001198}
1199
1200static void
1201islice_dealloc(isliceobject *lz)
1202{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001203 PyObject_GC_UnTrack(lz);
1204 Py_XDECREF(lz->it);
1205 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001206}
1207
1208static int
1209islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1210{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001211 Py_VISIT(lz->it);
1212 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001213}
1214
1215static PyObject *
1216islice_next(isliceobject *lz)
1217{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001218 PyObject *item;
1219 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001220 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001221 Py_ssize_t oldnext;
1222 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001223
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001224 iternext = *Py_TYPE(it)->tp_iternext;
1225 while (lz->cnt < lz->next) {
1226 item = iternext(it);
1227 if (item == NULL)
1228 return NULL;
1229 Py_DECREF(item);
1230 lz->cnt++;
1231 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001232 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001233 return NULL;
1234 item = iternext(it);
1235 if (item == NULL)
1236 return NULL;
1237 lz->cnt++;
1238 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001239 /* The (size_t) cast below avoids the danger of undefined
1240 behaviour from signed integer overflow. */
1241 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001242 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1243 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001244 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001245}
1246
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001247PyDoc_STRVAR(islice_doc,
1248"islice(iterable, [start,] stop [, step]) --> islice object\n\
1249\n\
1250Return an iterator whose next() method returns selected values from an\n\
1251iterable. If start is specified, will skip all preceding elements;\n\
1252otherwise, start defaults to zero. Step defaults to one. If\n\
1253specified as another value, step determines how many values are \n\
1254skipped between successive calls. Works like a slice() on a list\n\
1255but returns an iterator.");
1256
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001257static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001258 PyVarObject_HEAD_INIT(NULL, 0)
1259 "itertools.islice", /* tp_name */
1260 sizeof(isliceobject), /* tp_basicsize */
1261 0, /* tp_itemsize */
1262 /* methods */
1263 (destructor)islice_dealloc, /* tp_dealloc */
1264 0, /* tp_print */
1265 0, /* tp_getattr */
1266 0, /* tp_setattr */
1267 0, /* tp_reserved */
1268 0, /* tp_repr */
1269 0, /* tp_as_number */
1270 0, /* tp_as_sequence */
1271 0, /* tp_as_mapping */
1272 0, /* tp_hash */
1273 0, /* tp_call */
1274 0, /* tp_str */
1275 PyObject_GenericGetAttr, /* tp_getattro */
1276 0, /* tp_setattro */
1277 0, /* tp_as_buffer */
1278 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1279 Py_TPFLAGS_BASETYPE, /* tp_flags */
1280 islice_doc, /* tp_doc */
1281 (traverseproc)islice_traverse, /* tp_traverse */
1282 0, /* tp_clear */
1283 0, /* tp_richcompare */
1284 0, /* tp_weaklistoffset */
1285 PyObject_SelfIter, /* tp_iter */
1286 (iternextfunc)islice_next, /* tp_iternext */
1287 0, /* tp_methods */
1288 0, /* tp_members */
1289 0, /* tp_getset */
1290 0, /* tp_base */
1291 0, /* tp_dict */
1292 0, /* tp_descr_get */
1293 0, /* tp_descr_set */
1294 0, /* tp_dictoffset */
1295 0, /* tp_init */
1296 0, /* tp_alloc */
1297 islice_new, /* tp_new */
1298 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001299};
1300
1301
1302/* starmap object ************************************************************/
1303
1304typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001305 PyObject_HEAD
1306 PyObject *func;
1307 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001308} starmapobject;
1309
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001310static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001311
1312static PyObject *
1313starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1314{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001315 PyObject *func, *seq;
1316 PyObject *it;
1317 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001318
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001319 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1320 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001321
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001322 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1323 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001324
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001325 /* Get iterator. */
1326 it = PyObject_GetIter(seq);
1327 if (it == NULL)
1328 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001329
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001330 /* create starmapobject structure */
1331 lz = (starmapobject *)type->tp_alloc(type, 0);
1332 if (lz == NULL) {
1333 Py_DECREF(it);
1334 return NULL;
1335 }
1336 Py_INCREF(func);
1337 lz->func = func;
1338 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001339
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001340 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001341}
1342
1343static void
1344starmap_dealloc(starmapobject *lz)
1345{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001346 PyObject_GC_UnTrack(lz);
1347 Py_XDECREF(lz->func);
1348 Py_XDECREF(lz->it);
1349 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001350}
1351
1352static int
1353starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1354{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001355 Py_VISIT(lz->it);
1356 Py_VISIT(lz->func);
1357 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001358}
1359
1360static PyObject *
1361starmap_next(starmapobject *lz)
1362{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001363 PyObject *args;
1364 PyObject *result;
1365 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001366
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001367 args = (*Py_TYPE(it)->tp_iternext)(it);
1368 if (args == NULL)
1369 return NULL;
1370 if (!PyTuple_CheckExact(args)) {
1371 PyObject *newargs = PySequence_Tuple(args);
1372 Py_DECREF(args);
1373 if (newargs == NULL)
1374 return NULL;
1375 args = newargs;
1376 }
1377 result = PyObject_Call(lz->func, args, NULL);
1378 Py_DECREF(args);
1379 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001380}
1381
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001382PyDoc_STRVAR(starmap_doc,
1383"starmap(function, sequence) --> starmap object\n\
1384\n\
1385Return an iterator whose values are returned from the function evaluated\n\
1386with a argument tuple taken from the given sequence.");
1387
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001388static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001389 PyVarObject_HEAD_INIT(NULL, 0)
1390 "itertools.starmap", /* tp_name */
1391 sizeof(starmapobject), /* tp_basicsize */
1392 0, /* tp_itemsize */
1393 /* methods */
1394 (destructor)starmap_dealloc, /* tp_dealloc */
1395 0, /* tp_print */
1396 0, /* tp_getattr */
1397 0, /* tp_setattr */
1398 0, /* tp_reserved */
1399 0, /* tp_repr */
1400 0, /* tp_as_number */
1401 0, /* tp_as_sequence */
1402 0, /* tp_as_mapping */
1403 0, /* tp_hash */
1404 0, /* tp_call */
1405 0, /* tp_str */
1406 PyObject_GenericGetAttr, /* tp_getattro */
1407 0, /* tp_setattro */
1408 0, /* tp_as_buffer */
1409 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1410 Py_TPFLAGS_BASETYPE, /* tp_flags */
1411 starmap_doc, /* tp_doc */
1412 (traverseproc)starmap_traverse, /* tp_traverse */
1413 0, /* tp_clear */
1414 0, /* tp_richcompare */
1415 0, /* tp_weaklistoffset */
1416 PyObject_SelfIter, /* tp_iter */
1417 (iternextfunc)starmap_next, /* tp_iternext */
1418 0, /* tp_methods */
1419 0, /* tp_members */
1420 0, /* tp_getset */
1421 0, /* tp_base */
1422 0, /* tp_dict */
1423 0, /* tp_descr_get */
1424 0, /* tp_descr_set */
1425 0, /* tp_dictoffset */
1426 0, /* tp_init */
1427 0, /* tp_alloc */
1428 starmap_new, /* tp_new */
1429 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001430};
1431
1432
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001433/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001434
1435typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001436 PyObject_HEAD
1437 PyObject *source; /* Iterator over input iterables */
1438 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001439} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001440
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001441static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001442
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001443static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001444chain_new_internal(PyTypeObject *type, PyObject *source)
1445{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001446 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001447
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001448 lz = (chainobject *)type->tp_alloc(type, 0);
1449 if (lz == NULL) {
1450 Py_DECREF(source);
1451 return NULL;
1452 }
1453
1454 lz->source = source;
1455 lz->active = NULL;
1456 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001457}
1458
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001460chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001461{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001462 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001463
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001464 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1465 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001466
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001467 source = PyObject_GetIter(args);
1468 if (source == NULL)
1469 return NULL;
1470
1471 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001472}
1473
1474static PyObject *
1475chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1476{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001477 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001478
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001479 source = PyObject_GetIter(arg);
1480 if (source == NULL)
1481 return NULL;
1482
1483 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001484}
1485
1486static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001487chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001488{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001489 PyObject_GC_UnTrack(lz);
1490 Py_XDECREF(lz->active);
1491 Py_XDECREF(lz->source);
1492 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001493}
1494
Raymond Hettinger2012f172003-02-07 05:32:58 +00001495static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001496chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001497{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001498 Py_VISIT(lz->source);
1499 Py_VISIT(lz->active);
1500 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001501}
1502
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001503static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001504chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001505{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001506 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001507
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001508 if (lz->source == NULL)
1509 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001510
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001511 if (lz->active == NULL) {
1512 PyObject *iterable = PyIter_Next(lz->source);
1513 if (iterable == NULL) {
1514 Py_CLEAR(lz->source);
1515 return NULL; /* no more input sources */
1516 }
1517 lz->active = PyObject_GetIter(iterable);
1518 Py_DECREF(iterable);
1519 if (lz->active == NULL) {
1520 Py_CLEAR(lz->source);
1521 return NULL; /* input not iterable */
1522 }
1523 }
1524 item = PyIter_Next(lz->active);
1525 if (item != NULL)
1526 return item;
1527 if (PyErr_Occurred()) {
1528 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1529 PyErr_Clear();
1530 else
1531 return NULL; /* input raised an exception */
1532 }
1533 Py_CLEAR(lz->active);
1534 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001535}
1536
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001537PyDoc_STRVAR(chain_doc,
1538"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001539\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001540Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001541first iterable until it is exhausted, then elements from the next\n\
1542iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001543
Christian Heimesf16baeb2008-02-29 14:57:44 +00001544PyDoc_STRVAR(chain_from_iterable_doc,
1545"chain.from_iterable(iterable) --> chain object\n\
1546\n\
1547Alternate chain() contructor taking a single iterable argument\n\
1548that evaluates lazily.");
1549
1550static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001551 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1552 chain_from_iterable_doc},
1553 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001554};
1555
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001556static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001557 PyVarObject_HEAD_INIT(NULL, 0)
1558 "itertools.chain", /* tp_name */
1559 sizeof(chainobject), /* tp_basicsize */
1560 0, /* tp_itemsize */
1561 /* methods */
1562 (destructor)chain_dealloc, /* tp_dealloc */
1563 0, /* tp_print */
1564 0, /* tp_getattr */
1565 0, /* tp_setattr */
1566 0, /* tp_reserved */
1567 0, /* tp_repr */
1568 0, /* tp_as_number */
1569 0, /* tp_as_sequence */
1570 0, /* tp_as_mapping */
1571 0, /* tp_hash */
1572 0, /* tp_call */
1573 0, /* tp_str */
1574 PyObject_GenericGetAttr, /* tp_getattro */
1575 0, /* tp_setattro */
1576 0, /* tp_as_buffer */
1577 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1578 Py_TPFLAGS_BASETYPE, /* tp_flags */
1579 chain_doc, /* tp_doc */
1580 (traverseproc)chain_traverse, /* tp_traverse */
1581 0, /* tp_clear */
1582 0, /* tp_richcompare */
1583 0, /* tp_weaklistoffset */
1584 PyObject_SelfIter, /* tp_iter */
1585 (iternextfunc)chain_next, /* tp_iternext */
1586 chain_methods, /* tp_methods */
1587 0, /* tp_members */
1588 0, /* tp_getset */
1589 0, /* tp_base */
1590 0, /* tp_dict */
1591 0, /* tp_descr_get */
1592 0, /* tp_descr_set */
1593 0, /* tp_dictoffset */
1594 0, /* tp_init */
1595 0, /* tp_alloc */
1596 chain_new, /* tp_new */
1597 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001598};
1599
1600
Christian Heimesc3f30c42008-02-22 16:37:40 +00001601/* product object ************************************************************/
1602
1603typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001604 PyObject_HEAD
1605 PyObject *pools; /* tuple of pool tuples */
1606 Py_ssize_t *indices; /* one index per pool */
1607 PyObject *result; /* most recently returned result tuple */
1608 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001609} productobject;
1610
1611static PyTypeObject product_type;
1612
1613static PyObject *
1614product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1615{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001616 productobject *lz;
1617 Py_ssize_t nargs, npools, repeat=1;
1618 PyObject *pools = NULL;
1619 Py_ssize_t *indices = NULL;
1620 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001621
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001622 if (kwds != NULL) {
1623 char *kwlist[] = {"repeat", 0};
1624 PyObject *tmpargs = PyTuple_New(0);
1625 if (tmpargs == NULL)
1626 return NULL;
1627 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1628 Py_DECREF(tmpargs);
1629 return NULL;
1630 }
1631 Py_DECREF(tmpargs);
1632 if (repeat < 0) {
1633 PyErr_SetString(PyExc_ValueError,
1634 "repeat argument cannot be negative");
1635 return NULL;
1636 }
1637 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001638
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001639 assert(PyTuple_Check(args));
1640 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1641 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001642
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001643 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1644 if (indices == NULL) {
1645 PyErr_NoMemory();
1646 goto error;
1647 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001648
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001649 pools = PyTuple_New(npools);
1650 if (pools == NULL)
1651 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001652
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001653 for (i=0; i < nargs ; ++i) {
1654 PyObject *item = PyTuple_GET_ITEM(args, i);
1655 PyObject *pool = PySequence_Tuple(item);
1656 if (pool == NULL)
1657 goto error;
1658 PyTuple_SET_ITEM(pools, i, pool);
1659 indices[i] = 0;
1660 }
1661 for ( ; i < npools; ++i) {
1662 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1663 Py_INCREF(pool);
1664 PyTuple_SET_ITEM(pools, i, pool);
1665 indices[i] = 0;
1666 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001667
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001668 /* create productobject structure */
1669 lz = (productobject *)type->tp_alloc(type, 0);
1670 if (lz == NULL)
1671 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001672
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001673 lz->pools = pools;
1674 lz->indices = indices;
1675 lz->result = NULL;
1676 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001677
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001678 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001679
1680error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001681 if (indices != NULL)
1682 PyMem_Free(indices);
1683 Py_XDECREF(pools);
1684 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001685}
1686
1687static void
1688product_dealloc(productobject *lz)
1689{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001690 PyObject_GC_UnTrack(lz);
1691 Py_XDECREF(lz->pools);
1692 Py_XDECREF(lz->result);
1693 if (lz->indices != NULL)
1694 PyMem_Free(lz->indices);
1695 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001696}
1697
1698static int
1699product_traverse(productobject *lz, visitproc visit, void *arg)
1700{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001701 Py_VISIT(lz->pools);
1702 Py_VISIT(lz->result);
1703 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001704}
1705
1706static PyObject *
1707product_next(productobject *lz)
1708{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001709 PyObject *pool;
1710 PyObject *elem;
1711 PyObject *oldelem;
1712 PyObject *pools = lz->pools;
1713 PyObject *result = lz->result;
1714 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1715 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001716
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001717 if (lz->stopped)
1718 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001719
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001720 if (result == NULL) {
1721 /* On the first pass, return an initial tuple filled with the
1722 first element from each pool. */
1723 result = PyTuple_New(npools);
1724 if (result == NULL)
1725 goto empty;
1726 lz->result = result;
1727 for (i=0; i < npools; i++) {
1728 pool = PyTuple_GET_ITEM(pools, i);
1729 if (PyTuple_GET_SIZE(pool) == 0)
1730 goto empty;
1731 elem = PyTuple_GET_ITEM(pool, 0);
1732 Py_INCREF(elem);
1733 PyTuple_SET_ITEM(result, i, elem);
1734 }
1735 } else {
1736 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001737
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001738 /* Copy the previous result tuple or re-use it if available */
1739 if (Py_REFCNT(result) > 1) {
1740 PyObject *old_result = result;
1741 result = PyTuple_New(npools);
1742 if (result == NULL)
1743 goto empty;
1744 lz->result = result;
1745 for (i=0; i < npools; i++) {
1746 elem = PyTuple_GET_ITEM(old_result, i);
1747 Py_INCREF(elem);
1748 PyTuple_SET_ITEM(result, i, elem);
1749 }
1750 Py_DECREF(old_result);
1751 }
1752 /* Now, we've got the only copy so we can update it in-place */
1753 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001754
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001755 /* Update the pool indices right-to-left. Only advance to the
1756 next pool when the previous one rolls-over */
1757 for (i=npools-1 ; i >= 0 ; i--) {
1758 pool = PyTuple_GET_ITEM(pools, i);
1759 indices[i]++;
1760 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1761 /* Roll-over and advance to next pool */
1762 indices[i] = 0;
1763 elem = PyTuple_GET_ITEM(pool, 0);
1764 Py_INCREF(elem);
1765 oldelem = PyTuple_GET_ITEM(result, i);
1766 PyTuple_SET_ITEM(result, i, elem);
1767 Py_DECREF(oldelem);
1768 } else {
1769 /* No rollover. Just increment and stop here. */
1770 elem = PyTuple_GET_ITEM(pool, indices[i]);
1771 Py_INCREF(elem);
1772 oldelem = PyTuple_GET_ITEM(result, i);
1773 PyTuple_SET_ITEM(result, i, elem);
1774 Py_DECREF(oldelem);
1775 break;
1776 }
1777 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001778
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001779 /* If i is negative, then the indices have all rolled-over
1780 and we're done. */
1781 if (i < 0)
1782 goto empty;
1783 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001784
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001785 Py_INCREF(result);
1786 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001787
1788empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001789 lz->stopped = 1;
1790 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001791}
1792
1793PyDoc_STRVAR(product_doc,
1794"product(*iterables) --> product object\n\
1795\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001796Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001797For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1798The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1799cycle in a manner similar to an odometer (with the rightmost element changing\n\
1800on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00001801To compute the product of an iterable with itself, specify the number\n\
1802of repetitions with the optional repeat keyword argument. For example,\n\
1803product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001804product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1805product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1806
1807static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001808 PyVarObject_HEAD_INIT(NULL, 0)
1809 "itertools.product", /* tp_name */
1810 sizeof(productobject), /* tp_basicsize */
1811 0, /* tp_itemsize */
1812 /* methods */
1813 (destructor)product_dealloc, /* tp_dealloc */
1814 0, /* tp_print */
1815 0, /* tp_getattr */
1816 0, /* tp_setattr */
1817 0, /* tp_reserved */
1818 0, /* tp_repr */
1819 0, /* tp_as_number */
1820 0, /* tp_as_sequence */
1821 0, /* tp_as_mapping */
1822 0, /* tp_hash */
1823 0, /* tp_call */
1824 0, /* tp_str */
1825 PyObject_GenericGetAttr, /* tp_getattro */
1826 0, /* tp_setattro */
1827 0, /* tp_as_buffer */
1828 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1829 Py_TPFLAGS_BASETYPE, /* tp_flags */
1830 product_doc, /* tp_doc */
1831 (traverseproc)product_traverse, /* tp_traverse */
1832 0, /* tp_clear */
1833 0, /* tp_richcompare */
1834 0, /* tp_weaklistoffset */
1835 PyObject_SelfIter, /* tp_iter */
1836 (iternextfunc)product_next, /* tp_iternext */
1837 0, /* tp_methods */
1838 0, /* tp_members */
1839 0, /* tp_getset */
1840 0, /* tp_base */
1841 0, /* tp_dict */
1842 0, /* tp_descr_get */
1843 0, /* tp_descr_set */
1844 0, /* tp_dictoffset */
1845 0, /* tp_init */
1846 0, /* tp_alloc */
1847 product_new, /* tp_new */
1848 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001849};
1850
1851
Christian Heimes380f7f22008-02-28 11:19:05 +00001852/* combinations object ************************************************************/
1853
1854typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001855 PyObject_HEAD
1856 PyObject *pool; /* input converted to a tuple */
1857 Py_ssize_t *indices; /* one index per result element */
1858 PyObject *result; /* most recently returned result tuple */
1859 Py_ssize_t r; /* size of result tuple */
1860 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00001861} combinationsobject;
1862
1863static PyTypeObject combinations_type;
1864
1865static PyObject *
1866combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1867{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001868 combinationsobject *co;
1869 Py_ssize_t n;
1870 Py_ssize_t r;
1871 PyObject *pool = NULL;
1872 PyObject *iterable = NULL;
1873 Py_ssize_t *indices = NULL;
1874 Py_ssize_t i;
1875 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00001876
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001877 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
1878 &iterable, &r))
1879 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001880
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001881 pool = PySequence_Tuple(iterable);
1882 if (pool == NULL)
1883 goto error;
1884 n = PyTuple_GET_SIZE(pool);
1885 if (r < 0) {
1886 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
1887 goto error;
1888 }
Christian Heimes380f7f22008-02-28 11:19:05 +00001889
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001890 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
1891 if (indices == NULL) {
1892 PyErr_NoMemory();
1893 goto error;
1894 }
Christian Heimes380f7f22008-02-28 11:19:05 +00001895
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001896 for (i=0 ; i<r ; i++)
1897 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00001898
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001899 /* create combinationsobject structure */
1900 co = (combinationsobject *)type->tp_alloc(type, 0);
1901 if (co == NULL)
1902 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00001903
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001904 co->pool = pool;
1905 co->indices = indices;
1906 co->result = NULL;
1907 co->r = r;
1908 co->stopped = r > n ? 1 : 0;
1909
1910 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00001911
1912error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001913 if (indices != NULL)
1914 PyMem_Free(indices);
1915 Py_XDECREF(pool);
1916 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001917}
1918
1919static void
1920combinations_dealloc(combinationsobject *co)
1921{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001922 PyObject_GC_UnTrack(co);
1923 Py_XDECREF(co->pool);
1924 Py_XDECREF(co->result);
1925 if (co->indices != NULL)
1926 PyMem_Free(co->indices);
1927 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00001928}
1929
1930static int
1931combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
1932{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001933 Py_VISIT(co->pool);
1934 Py_VISIT(co->result);
1935 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00001936}
1937
1938static PyObject *
1939combinations_next(combinationsobject *co)
1940{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001941 PyObject *elem;
1942 PyObject *oldelem;
1943 PyObject *pool = co->pool;
1944 Py_ssize_t *indices = co->indices;
1945 PyObject *result = co->result;
1946 Py_ssize_t n = PyTuple_GET_SIZE(pool);
1947 Py_ssize_t r = co->r;
1948 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00001949
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001950 if (co->stopped)
1951 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001952
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001953 if (result == NULL) {
1954 /* On the first pass, initialize result tuple using the indices */
1955 result = PyTuple_New(r);
1956 if (result == NULL)
1957 goto empty;
1958 co->result = result;
1959 for (i=0; i<r ; i++) {
1960 index = indices[i];
1961 elem = PyTuple_GET_ITEM(pool, index);
1962 Py_INCREF(elem);
1963 PyTuple_SET_ITEM(result, i, elem);
1964 }
1965 } else {
1966 /* Copy the previous result tuple or re-use it if available */
1967 if (Py_REFCNT(result) > 1) {
1968 PyObject *old_result = result;
1969 result = PyTuple_New(r);
1970 if (result == NULL)
1971 goto empty;
1972 co->result = result;
1973 for (i=0; i<r ; i++) {
1974 elem = PyTuple_GET_ITEM(old_result, i);
1975 Py_INCREF(elem);
1976 PyTuple_SET_ITEM(result, i, elem);
1977 }
1978 Py_DECREF(old_result);
1979 }
1980 /* Now, we've got the only copy so we can update it in-place
1981 * CPython's empty tuple is a singleton and cached in
1982 * PyTuple's freelist.
1983 */
1984 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00001985
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001986 /* Scan indices right-to-left until finding one that is not
1987 at its maximum (i + n - r). */
1988 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
1989 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00001990
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001991 /* If i is negative, then the indices are all at
1992 their maximum value and we're done. */
1993 if (i < 0)
1994 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00001995
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001996 /* Increment the current index which we know is not at its
1997 maximum. Then move back to the right setting each index
1998 to its lowest possible value (one higher than the index
1999 to its left -- this maintains the sort order invariant). */
2000 indices[i]++;
2001 for (j=i+1 ; j<r ; j++)
2002 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002003
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002004 /* Update the result tuple for the new indices
2005 starting with i, the leftmost index that changed */
2006 for ( ; i<r ; i++) {
2007 index = indices[i];
2008 elem = PyTuple_GET_ITEM(pool, index);
2009 Py_INCREF(elem);
2010 oldelem = PyTuple_GET_ITEM(result, i);
2011 PyTuple_SET_ITEM(result, i, elem);
2012 Py_DECREF(oldelem);
2013 }
2014 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002015
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002016 Py_INCREF(result);
2017 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002018
2019empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002020 co->stopped = 1;
2021 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002022}
2023
2024PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002025"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002026\n\
2027Return successive r-length combinations of elements in the iterable.\n\n\
2028combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2029
2030static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002031 PyVarObject_HEAD_INIT(NULL, 0)
2032 "itertools.combinations", /* tp_name */
2033 sizeof(combinationsobject), /* tp_basicsize */
2034 0, /* tp_itemsize */
2035 /* methods */
2036 (destructor)combinations_dealloc, /* tp_dealloc */
2037 0, /* tp_print */
2038 0, /* tp_getattr */
2039 0, /* tp_setattr */
2040 0, /* tp_reserved */
2041 0, /* tp_repr */
2042 0, /* tp_as_number */
2043 0, /* tp_as_sequence */
2044 0, /* tp_as_mapping */
2045 0, /* tp_hash */
2046 0, /* tp_call */
2047 0, /* tp_str */
2048 PyObject_GenericGetAttr, /* tp_getattro */
2049 0, /* tp_setattro */
2050 0, /* tp_as_buffer */
2051 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2052 Py_TPFLAGS_BASETYPE, /* tp_flags */
2053 combinations_doc, /* tp_doc */
2054 (traverseproc)combinations_traverse, /* tp_traverse */
2055 0, /* tp_clear */
2056 0, /* tp_richcompare */
2057 0, /* tp_weaklistoffset */
2058 PyObject_SelfIter, /* tp_iter */
2059 (iternextfunc)combinations_next, /* tp_iternext */
2060 0, /* tp_methods */
2061 0, /* tp_members */
2062 0, /* tp_getset */
2063 0, /* tp_base */
2064 0, /* tp_dict */
2065 0, /* tp_descr_get */
2066 0, /* tp_descr_set */
2067 0, /* tp_dictoffset */
2068 0, /* tp_init */
2069 0, /* tp_alloc */
2070 combinations_new, /* tp_new */
2071 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002072};
2073
2074
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002075/* combinations with replacement object *******************************************/
2076
2077/* Equivalent to:
2078
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002079 def combinations_with_replacement(iterable, r):
2080 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2081 # number items returned: (n+r-1)! / r! / (n-1)!
2082 pool = tuple(iterable)
2083 n = len(pool)
2084 indices = [0] * r
2085 yield tuple(pool[i] for i in indices)
2086 while 1:
2087 for i in reversed(range(r)):
2088 if indices[i] != n - 1:
2089 break
2090 else:
2091 return
2092 indices[i:] = [indices[i] + 1] * (r - i)
2093 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002094
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002095 def combinations_with_replacement2(iterable, r):
2096 'Alternate version that filters from product()'
2097 pool = tuple(iterable)
2098 n = len(pool)
2099 for indices in product(range(n), repeat=r):
2100 if sorted(indices) == list(indices):
2101 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002102*/
2103typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002104 PyObject_HEAD
2105 PyObject *pool; /* input converted to a tuple */
2106 Py_ssize_t *indices; /* one index per result element */
2107 PyObject *result; /* most recently returned result tuple */
2108 Py_ssize_t r; /* size of result tuple */
2109 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002110} cwrobject;
2111
2112static PyTypeObject cwr_type;
2113
2114static PyObject *
2115cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2116{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002117 cwrobject *co;
2118 Py_ssize_t n;
2119 Py_ssize_t r;
2120 PyObject *pool = NULL;
2121 PyObject *iterable = NULL;
2122 Py_ssize_t *indices = NULL;
2123 Py_ssize_t i;
2124 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002125
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002126 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2127 &iterable, &r))
2128 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002129
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002130 pool = PySequence_Tuple(iterable);
2131 if (pool == NULL)
2132 goto error;
2133 n = PyTuple_GET_SIZE(pool);
2134 if (r < 0) {
2135 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2136 goto error;
2137 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002138
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002139 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2140 if (indices == NULL) {
2141 PyErr_NoMemory();
2142 goto error;
2143 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002144
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002145 for (i=0 ; i<r ; i++)
2146 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002147
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002148 /* create cwrobject structure */
2149 co = (cwrobject *)type->tp_alloc(type, 0);
2150 if (co == NULL)
2151 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002152
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002153 co->pool = pool;
2154 co->indices = indices;
2155 co->result = NULL;
2156 co->r = r;
2157 co->stopped = !n && r;
2158
2159 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002160
2161error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002162 if (indices != NULL)
2163 PyMem_Free(indices);
2164 Py_XDECREF(pool);
2165 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002166}
2167
2168static void
2169cwr_dealloc(cwrobject *co)
2170{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002171 PyObject_GC_UnTrack(co);
2172 Py_XDECREF(co->pool);
2173 Py_XDECREF(co->result);
2174 if (co->indices != NULL)
2175 PyMem_Free(co->indices);
2176 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002177}
2178
2179static int
2180cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2181{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002182 Py_VISIT(co->pool);
2183 Py_VISIT(co->result);
2184 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002185}
2186
2187static PyObject *
2188cwr_next(cwrobject *co)
2189{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002190 PyObject *elem;
2191 PyObject *oldelem;
2192 PyObject *pool = co->pool;
2193 Py_ssize_t *indices = co->indices;
2194 PyObject *result = co->result;
2195 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2196 Py_ssize_t r = co->r;
2197 Py_ssize_t i, j, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002198
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002199 if (co->stopped)
2200 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002201
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002202 if (result == NULL) {
2203 /* On the first pass, initialize result tuple using the indices */
2204 result = PyTuple_New(r);
2205 if (result == NULL)
2206 goto empty;
2207 co->result = result;
2208 for (i=0; i<r ; i++) {
2209 index = indices[i];
2210 elem = PyTuple_GET_ITEM(pool, index);
2211 Py_INCREF(elem);
2212 PyTuple_SET_ITEM(result, i, elem);
2213 }
2214 } else {
2215 /* Copy the previous result tuple or re-use it if available */
2216 if (Py_REFCNT(result) > 1) {
2217 PyObject *old_result = result;
2218 result = PyTuple_New(r);
2219 if (result == NULL)
2220 goto empty;
2221 co->result = result;
2222 for (i=0; i<r ; i++) {
2223 elem = PyTuple_GET_ITEM(old_result, i);
2224 Py_INCREF(elem);
2225 PyTuple_SET_ITEM(result, i, elem);
2226 }
2227 Py_DECREF(old_result);
2228 }
2229 /* Now, we've got the only copy so we can update it in-place CPython's
2230 empty tuple is a singleton and cached in PyTuple's freelist. */
2231 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002232
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002233 /* Scan indices right-to-left until finding one that is not
2234 * at its maximum (n-1). */
2235 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2236 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002237
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002238 /* If i is negative, then the indices are all at
2239 their maximum value and we're done. */
2240 if (i < 0)
2241 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002242
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002243 /* Increment the current index which we know is not at its
2244 maximum. Then set all to the right to the same value. */
2245 indices[i]++;
2246 for (j=i+1 ; j<r ; j++)
2247 indices[j] = indices[j-1];
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002248
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002249 /* Update the result tuple for the new indices
2250 starting with i, the leftmost index that changed */
2251 for ( ; i<r ; i++) {
2252 index = indices[i];
2253 elem = PyTuple_GET_ITEM(pool, index);
2254 Py_INCREF(elem);
2255 oldelem = PyTuple_GET_ITEM(result, i);
2256 PyTuple_SET_ITEM(result, i, elem);
2257 Py_DECREF(oldelem);
2258 }
2259 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002260
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002261 Py_INCREF(result);
2262 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002263
2264empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002265 co->stopped = 1;
2266 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002267}
2268
2269PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002270"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002271\n\
2272Return successive r-length combinations of elements in the iterable\n\
2273allowing individual elements to have successive repeats.\n\
2274combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2275
2276static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002277 PyVarObject_HEAD_INIT(NULL, 0)
2278 "itertools.combinations_with_replacement", /* tp_name */
2279 sizeof(cwrobject), /* tp_basicsize */
2280 0, /* tp_itemsize */
2281 /* methods */
2282 (destructor)cwr_dealloc, /* tp_dealloc */
2283 0, /* tp_print */
2284 0, /* tp_getattr */
2285 0, /* tp_setattr */
2286 0, /* tp_reserved */
2287 0, /* tp_repr */
2288 0, /* tp_as_number */
2289 0, /* tp_as_sequence */
2290 0, /* tp_as_mapping */
2291 0, /* tp_hash */
2292 0, /* tp_call */
2293 0, /* tp_str */
2294 PyObject_GenericGetAttr, /* tp_getattro */
2295 0, /* tp_setattro */
2296 0, /* tp_as_buffer */
2297 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2298 Py_TPFLAGS_BASETYPE, /* tp_flags */
2299 cwr_doc, /* tp_doc */
2300 (traverseproc)cwr_traverse, /* tp_traverse */
2301 0, /* tp_clear */
2302 0, /* tp_richcompare */
2303 0, /* tp_weaklistoffset */
2304 PyObject_SelfIter, /* tp_iter */
2305 (iternextfunc)cwr_next, /* tp_iternext */
2306 0, /* tp_methods */
2307 0, /* tp_members */
2308 0, /* tp_getset */
2309 0, /* tp_base */
2310 0, /* tp_dict */
2311 0, /* tp_descr_get */
2312 0, /* tp_descr_set */
2313 0, /* tp_dictoffset */
2314 0, /* tp_init */
2315 0, /* tp_alloc */
2316 cwr_new, /* tp_new */
2317 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002318};
2319
2320
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002321/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002322
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002323def permutations(iterable, r=None):
2324 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2325 pool = tuple(iterable)
2326 n = len(pool)
2327 r = n if r is None else r
2328 indices = range(n)
2329 cycles = range(n-r+1, n+1)[::-1]
2330 yield tuple(pool[i] for i in indices[:r])
2331 while n:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002332 for i in reversed(range(r)):
2333 cycles[i] -= 1
2334 if cycles[i] == 0:
2335 indices[i:] = indices[i+1:] + indices[i:i+1]
2336 cycles[i] = n - i
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002337 else:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002338 j = cycles[i]
2339 indices[i], indices[-j] = indices[-j], indices[i]
2340 yield tuple(pool[i] for i in indices[:r])
2341 break
2342 else:
2343 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002344*/
2345
2346typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002347 PyObject_HEAD
2348 PyObject *pool; /* input converted to a tuple */
2349 Py_ssize_t *indices; /* one index per element in the pool */
2350 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2351 PyObject *result; /* most recently returned result tuple */
2352 Py_ssize_t r; /* size of result tuple */
2353 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002354} permutationsobject;
2355
2356static PyTypeObject permutations_type;
2357
2358static PyObject *
2359permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2360{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002361 permutationsobject *po;
2362 Py_ssize_t n;
2363 Py_ssize_t r;
2364 PyObject *robj = Py_None;
2365 PyObject *pool = NULL;
2366 PyObject *iterable = NULL;
2367 Py_ssize_t *indices = NULL;
2368 Py_ssize_t *cycles = NULL;
2369 Py_ssize_t i;
2370 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002371
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002372 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2373 &iterable, &robj))
2374 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002375
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002376 pool = PySequence_Tuple(iterable);
2377 if (pool == NULL)
2378 goto error;
2379 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002380
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002381 r = n;
2382 if (robj != Py_None) {
2383 if (!PyLong_Check(robj)) {
2384 PyErr_SetString(PyExc_TypeError, "Expected int as r");
2385 goto error;
2386 }
2387 r = PyLong_AsSsize_t(robj);
2388 if (r == -1 && PyErr_Occurred())
2389 goto error;
2390 }
2391 if (r < 0) {
2392 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2393 goto error;
2394 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002395
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002396 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2397 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2398 if (indices == NULL || cycles == NULL) {
2399 PyErr_NoMemory();
2400 goto error;
2401 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002402
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002403 for (i=0 ; i<n ; i++)
2404 indices[i] = i;
2405 for (i=0 ; i<r ; i++)
2406 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002407
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002408 /* create permutationsobject structure */
2409 po = (permutationsobject *)type->tp_alloc(type, 0);
2410 if (po == NULL)
2411 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002412
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002413 po->pool = pool;
2414 po->indices = indices;
2415 po->cycles = cycles;
2416 po->result = NULL;
2417 po->r = r;
2418 po->stopped = r > n ? 1 : 0;
2419
2420 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002421
2422error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002423 if (indices != NULL)
2424 PyMem_Free(indices);
2425 if (cycles != NULL)
2426 PyMem_Free(cycles);
2427 Py_XDECREF(pool);
2428 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002429}
2430
2431static void
2432permutations_dealloc(permutationsobject *po)
2433{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002434 PyObject_GC_UnTrack(po);
2435 Py_XDECREF(po->pool);
2436 Py_XDECREF(po->result);
2437 PyMem_Free(po->indices);
2438 PyMem_Free(po->cycles);
2439 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002440}
2441
2442static int
2443permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2444{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002445 Py_VISIT(po->pool);
2446 Py_VISIT(po->result);
2447 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002448}
2449
2450static PyObject *
2451permutations_next(permutationsobject *po)
2452{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002453 PyObject *elem;
2454 PyObject *oldelem;
2455 PyObject *pool = po->pool;
2456 Py_ssize_t *indices = po->indices;
2457 Py_ssize_t *cycles = po->cycles;
2458 PyObject *result = po->result;
2459 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2460 Py_ssize_t r = po->r;
2461 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002462
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002463 if (po->stopped)
2464 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002465
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002466 if (result == NULL) {
2467 /* On the first pass, initialize result tuple using the indices */
2468 result = PyTuple_New(r);
2469 if (result == NULL)
2470 goto empty;
2471 po->result = result;
2472 for (i=0; i<r ; i++) {
2473 index = indices[i];
2474 elem = PyTuple_GET_ITEM(pool, index);
2475 Py_INCREF(elem);
2476 PyTuple_SET_ITEM(result, i, elem);
2477 }
2478 } else {
2479 if (n == 0)
2480 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002481
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002482 /* Copy the previous result tuple or re-use it if available */
2483 if (Py_REFCNT(result) > 1) {
2484 PyObject *old_result = result;
2485 result = PyTuple_New(r);
2486 if (result == NULL)
2487 goto empty;
2488 po->result = result;
2489 for (i=0; i<r ; i++) {
2490 elem = PyTuple_GET_ITEM(old_result, i);
2491 Py_INCREF(elem);
2492 PyTuple_SET_ITEM(result, i, elem);
2493 }
2494 Py_DECREF(old_result);
2495 }
2496 /* Now, we've got the only copy so we can update it in-place */
2497 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002498
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002499 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2500 for (i=r-1 ; i>=0 ; i--) {
2501 cycles[i] -= 1;
2502 if (cycles[i] == 0) {
2503 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2504 index = indices[i];
2505 for (j=i ; j<n-1 ; j++)
2506 indices[j] = indices[j+1];
2507 indices[n-1] = index;
2508 cycles[i] = n - i;
2509 } else {
2510 j = cycles[i];
2511 index = indices[i];
2512 indices[i] = indices[n-j];
2513 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002514
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002515 for (k=i; k<r ; k++) {
2516 /* start with i, the leftmost element that changed */
2517 /* yield tuple(pool[k] for k in indices[:r]) */
2518 index = indices[k];
2519 elem = PyTuple_GET_ITEM(pool, index);
2520 Py_INCREF(elem);
2521 oldelem = PyTuple_GET_ITEM(result, k);
2522 PyTuple_SET_ITEM(result, k, elem);
2523 Py_DECREF(oldelem);
2524 }
2525 break;
2526 }
2527 }
2528 /* If i is negative, then the cycles have all
2529 rolled-over and we're done. */
2530 if (i < 0)
2531 goto empty;
2532 }
2533 Py_INCREF(result);
2534 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002535
2536empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002537 po->stopped = 1;
2538 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002539}
2540
2541PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00002542"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002543\n\
2544Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00002545permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002546
2547static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002548 PyVarObject_HEAD_INIT(NULL, 0)
2549 "itertools.permutations", /* tp_name */
2550 sizeof(permutationsobject), /* tp_basicsize */
2551 0, /* tp_itemsize */
2552 /* methods */
2553 (destructor)permutations_dealloc, /* tp_dealloc */
2554 0, /* tp_print */
2555 0, /* tp_getattr */
2556 0, /* tp_setattr */
2557 0, /* tp_reserved */
2558 0, /* tp_repr */
2559 0, /* tp_as_number */
2560 0, /* tp_as_sequence */
2561 0, /* tp_as_mapping */
2562 0, /* tp_hash */
2563 0, /* tp_call */
2564 0, /* tp_str */
2565 PyObject_GenericGetAttr, /* tp_getattro */
2566 0, /* tp_setattro */
2567 0, /* tp_as_buffer */
2568 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2569 Py_TPFLAGS_BASETYPE, /* tp_flags */
2570 permutations_doc, /* tp_doc */
2571 (traverseproc)permutations_traverse, /* tp_traverse */
2572 0, /* tp_clear */
2573 0, /* tp_richcompare */
2574 0, /* tp_weaklistoffset */
2575 PyObject_SelfIter, /* tp_iter */
2576 (iternextfunc)permutations_next, /* tp_iternext */
2577 0, /* tp_methods */
2578 0, /* tp_members */
2579 0, /* tp_getset */
2580 0, /* tp_base */
2581 0, /* tp_dict */
2582 0, /* tp_descr_get */
2583 0, /* tp_descr_set */
2584 0, /* tp_dictoffset */
2585 0, /* tp_init */
2586 0, /* tp_alloc */
2587 permutations_new, /* tp_new */
2588 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002589};
2590
Raymond Hettinger482ba772010-12-01 22:48:00 +00002591/* accumulate object ************************************************************/
2592
2593typedef struct {
2594 PyObject_HEAD
2595 PyObject *total;
2596 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07002597 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002598} accumulateobject;
2599
2600static PyTypeObject accumulate_type;
2601
2602static PyObject *
2603accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2604{
Raymond Hettinger5d446132011-03-27 18:52:10 -07002605 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00002606 PyObject *iterable;
2607 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07002608 PyObject *binop = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002609 accumulateobject *lz;
2610
Raymond Hettinger5d446132011-03-27 18:52:10 -07002611 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
2612 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002613 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002614
2615 /* Get iterator. */
2616 it = PyObject_GetIter(iterable);
2617 if (it == NULL)
2618 return NULL;
2619
Raymond Hettinger482ba772010-12-01 22:48:00 +00002620 /* create accumulateobject structure */
2621 lz = (accumulateobject *)type->tp_alloc(type, 0);
2622 if (lz == NULL) {
2623 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002624 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002625 }
2626
Raymond Hettinger5d446132011-03-27 18:52:10 -07002627 Py_XINCREF(binop);
2628 lz->binop = binop;
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002629 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002630 lz->it = it;
2631 return (PyObject *)lz;
2632}
2633
2634static void
2635accumulate_dealloc(accumulateobject *lz)
2636{
2637 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07002638 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00002639 Py_XDECREF(lz->total);
2640 Py_XDECREF(lz->it);
2641 Py_TYPE(lz)->tp_free(lz);
2642}
2643
2644static int
2645accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
2646{
Raymond Hettinger5d446132011-03-27 18:52:10 -07002647 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00002648 Py_VISIT(lz->it);
2649 Py_VISIT(lz->total);
2650 return 0;
2651}
2652
2653static PyObject *
2654accumulate_next(accumulateobject *lz)
2655{
2656 PyObject *val, *oldtotal, *newtotal;
2657
2658 val = PyIter_Next(lz->it);
2659 if (val == NULL)
2660 return NULL;
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002661
2662 if (lz->total == NULL) {
2663 Py_INCREF(val);
2664 lz->total = val;
2665 return lz->total;
2666 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07002667
2668 if (lz->binop == NULL)
2669 newtotal = PyNumber_Add(lz->total, val);
2670 else
2671 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002672 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00002673 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002674 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002675
2676 oldtotal = lz->total;
2677 lz->total = newtotal;
2678 Py_DECREF(oldtotal);
2679
2680 Py_INCREF(newtotal);
2681 return newtotal;
2682}
2683
2684PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07002685"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00002686\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07002687Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00002688
2689static PyTypeObject accumulate_type = {
2690 PyVarObject_HEAD_INIT(NULL, 0)
2691 "itertools.accumulate", /* tp_name */
2692 sizeof(accumulateobject), /* tp_basicsize */
2693 0, /* tp_itemsize */
2694 /* methods */
2695 (destructor)accumulate_dealloc, /* tp_dealloc */
2696 0, /* tp_print */
2697 0, /* tp_getattr */
2698 0, /* tp_setattr */
2699 0, /* tp_reserved */
2700 0, /* tp_repr */
2701 0, /* tp_as_number */
2702 0, /* tp_as_sequence */
2703 0, /* tp_as_mapping */
2704 0, /* tp_hash */
2705 0, /* tp_call */
2706 0, /* tp_str */
2707 PyObject_GenericGetAttr, /* tp_getattro */
2708 0, /* tp_setattro */
2709 0, /* tp_as_buffer */
2710 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2711 Py_TPFLAGS_BASETYPE, /* tp_flags */
2712 accumulate_doc, /* tp_doc */
2713 (traverseproc)accumulate_traverse, /* tp_traverse */
2714 0, /* tp_clear */
2715 0, /* tp_richcompare */
2716 0, /* tp_weaklistoffset */
2717 PyObject_SelfIter, /* tp_iter */
2718 (iternextfunc)accumulate_next, /* tp_iternext */
2719 0, /* tp_methods */
2720 0, /* tp_members */
2721 0, /* tp_getset */
2722 0, /* tp_base */
2723 0, /* tp_dict */
2724 0, /* tp_descr_get */
2725 0, /* tp_descr_set */
2726 0, /* tp_dictoffset */
2727 0, /* tp_init */
2728 0, /* tp_alloc */
2729 accumulate_new, /* tp_new */
2730 PyObject_GC_Del, /* tp_free */
2731};
2732
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002733
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002734/* compress object ************************************************************/
2735
2736/* Equivalent to:
2737
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002738 def compress(data, selectors):
2739 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2740 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002741*/
2742
2743typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002744 PyObject_HEAD
2745 PyObject *data;
2746 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002747} compressobject;
2748
2749static PyTypeObject compress_type;
2750
2751static PyObject *
2752compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2753{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002754 PyObject *seq1, *seq2;
2755 PyObject *data=NULL, *selectors=NULL;
2756 compressobject *lz;
2757 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002758
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002759 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2760 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002761
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002762 data = PyObject_GetIter(seq1);
2763 if (data == NULL)
2764 goto fail;
2765 selectors = PyObject_GetIter(seq2);
2766 if (selectors == NULL)
2767 goto fail;
2768
2769 /* create compressobject structure */
2770 lz = (compressobject *)type->tp_alloc(type, 0);
2771 if (lz == NULL)
2772 goto fail;
2773 lz->data = data;
2774 lz->selectors = selectors;
2775 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002776
2777fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002778 Py_XDECREF(data);
2779 Py_XDECREF(selectors);
2780 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002781}
2782
2783static void
2784compress_dealloc(compressobject *lz)
2785{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002786 PyObject_GC_UnTrack(lz);
2787 Py_XDECREF(lz->data);
2788 Py_XDECREF(lz->selectors);
2789 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002790}
2791
2792static int
2793compress_traverse(compressobject *lz, visitproc visit, void *arg)
2794{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002795 Py_VISIT(lz->data);
2796 Py_VISIT(lz->selectors);
2797 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002798}
2799
2800static PyObject *
2801compress_next(compressobject *lz)
2802{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002803 PyObject *data = lz->data, *selectors = lz->selectors;
2804 PyObject *datum, *selector;
2805 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2806 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2807 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002808
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002809 while (1) {
2810 /* Steps: get datum, get selector, evaluate selector.
2811 Order is important (to match the pure python version
2812 in terms of which input gets a chance to raise an
2813 exception first).
2814 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002815
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002816 datum = datanext(data);
2817 if (datum == NULL)
2818 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002819
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002820 selector = selectornext(selectors);
2821 if (selector == NULL) {
2822 Py_DECREF(datum);
2823 return NULL;
2824 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002825
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002826 ok = PyObject_IsTrue(selector);
2827 Py_DECREF(selector);
2828 if (ok == 1)
2829 return datum;
2830 Py_DECREF(datum);
2831 if (ok == -1)
2832 return NULL;
2833 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002834}
2835
2836PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00002837"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002838\n\
2839Return data elements corresponding to true selector elements.\n\
2840Forms a shorter iterator from selected data elements using the\n\
2841selectors to choose the data elements.");
2842
2843static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002844 PyVarObject_HEAD_INIT(NULL, 0)
2845 "itertools.compress", /* tp_name */
2846 sizeof(compressobject), /* tp_basicsize */
2847 0, /* tp_itemsize */
2848 /* methods */
2849 (destructor)compress_dealloc, /* tp_dealloc */
2850 0, /* tp_print */
2851 0, /* tp_getattr */
2852 0, /* tp_setattr */
2853 0, /* tp_reserved */
2854 0, /* tp_repr */
2855 0, /* tp_as_number */
2856 0, /* tp_as_sequence */
2857 0, /* tp_as_mapping */
2858 0, /* tp_hash */
2859 0, /* tp_call */
2860 0, /* tp_str */
2861 PyObject_GenericGetAttr, /* tp_getattro */
2862 0, /* tp_setattro */
2863 0, /* tp_as_buffer */
2864 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2865 Py_TPFLAGS_BASETYPE, /* tp_flags */
2866 compress_doc, /* tp_doc */
2867 (traverseproc)compress_traverse, /* tp_traverse */
2868 0, /* tp_clear */
2869 0, /* tp_richcompare */
2870 0, /* tp_weaklistoffset */
2871 PyObject_SelfIter, /* tp_iter */
2872 (iternextfunc)compress_next, /* tp_iternext */
2873 0, /* tp_methods */
2874 0, /* tp_members */
2875 0, /* tp_getset */
2876 0, /* tp_base */
2877 0, /* tp_dict */
2878 0, /* tp_descr_get */
2879 0, /* tp_descr_set */
2880 0, /* tp_dictoffset */
2881 0, /* tp_init */
2882 0, /* tp_alloc */
2883 compress_new, /* tp_new */
2884 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002885};
2886
2887
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002888/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00002889
2890typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002891 PyObject_HEAD
2892 PyObject *func;
2893 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002894} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002895
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002896static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002897
2898static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002899filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002900{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002901 PyObject *func, *seq;
2902 PyObject *it;
2903 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002904
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002905 if (type == &filterfalse_type &&
2906 !_PyArg_NoKeywords("filterfalse()", kwds))
2907 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002908
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002909 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
2910 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002911
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002912 /* Get iterator. */
2913 it = PyObject_GetIter(seq);
2914 if (it == NULL)
2915 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002916
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002917 /* create filterfalseobject structure */
2918 lz = (filterfalseobject *)type->tp_alloc(type, 0);
2919 if (lz == NULL) {
2920 Py_DECREF(it);
2921 return NULL;
2922 }
2923 Py_INCREF(func);
2924 lz->func = func;
2925 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002926
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002927 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002928}
2929
2930static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002931filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002932{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002933 PyObject_GC_UnTrack(lz);
2934 Py_XDECREF(lz->func);
2935 Py_XDECREF(lz->it);
2936 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002937}
2938
2939static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002940filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002941{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002942 Py_VISIT(lz->it);
2943 Py_VISIT(lz->func);
2944 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002945}
2946
2947static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002948filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002949{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002950 PyObject *item;
2951 PyObject *it = lz->it;
2952 long ok;
2953 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002954
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002955 iternext = *Py_TYPE(it)->tp_iternext;
2956 for (;;) {
2957 item = iternext(it);
2958 if (item == NULL)
2959 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002960
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002961 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2962 ok = PyObject_IsTrue(item);
2963 } else {
2964 PyObject *good;
2965 good = PyObject_CallFunctionObjArgs(lz->func,
2966 item, NULL);
2967 if (good == NULL) {
2968 Py_DECREF(item);
2969 return NULL;
2970 }
2971 ok = PyObject_IsTrue(good);
2972 Py_DECREF(good);
2973 }
2974 if (!ok)
2975 return item;
2976 Py_DECREF(item);
2977 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00002978}
2979
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002980PyDoc_STRVAR(filterfalse_doc,
2981"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002982\n\
2983Return those items of sequence for which function(item) is false.\n\
2984If function is None, return the items that are false.");
2985
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002986static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002987 PyVarObject_HEAD_INIT(NULL, 0)
2988 "itertools.filterfalse", /* tp_name */
2989 sizeof(filterfalseobject), /* tp_basicsize */
2990 0, /* tp_itemsize */
2991 /* methods */
2992 (destructor)filterfalse_dealloc, /* tp_dealloc */
2993 0, /* tp_print */
2994 0, /* tp_getattr */
2995 0, /* tp_setattr */
2996 0, /* tp_reserved */
2997 0, /* tp_repr */
2998 0, /* tp_as_number */
2999 0, /* tp_as_sequence */
3000 0, /* tp_as_mapping */
3001 0, /* tp_hash */
3002 0, /* tp_call */
3003 0, /* tp_str */
3004 PyObject_GenericGetAttr, /* tp_getattro */
3005 0, /* tp_setattro */
3006 0, /* tp_as_buffer */
3007 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3008 Py_TPFLAGS_BASETYPE, /* tp_flags */
3009 filterfalse_doc, /* tp_doc */
3010 (traverseproc)filterfalse_traverse, /* tp_traverse */
3011 0, /* tp_clear */
3012 0, /* tp_richcompare */
3013 0, /* tp_weaklistoffset */
3014 PyObject_SelfIter, /* tp_iter */
3015 (iternextfunc)filterfalse_next, /* tp_iternext */
3016 0, /* tp_methods */
3017 0, /* tp_members */
3018 0, /* tp_getset */
3019 0, /* tp_base */
3020 0, /* tp_dict */
3021 0, /* tp_descr_get */
3022 0, /* tp_descr_set */
3023 0, /* tp_dictoffset */
3024 0, /* tp_init */
3025 0, /* tp_alloc */
3026 filterfalse_new, /* tp_new */
3027 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003028};
3029
3030
3031/* count object ************************************************************/
3032
3033typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003034 PyObject_HEAD
3035 Py_ssize_t cnt;
3036 PyObject *long_cnt;
3037 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003038} countobject;
3039
Raymond Hettinger30729212009-02-12 06:28:27 +00003040/* Counting logic and invariants:
3041
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003042fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003043
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003044 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3045 Advances with: cnt += 1
3046 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003047
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003048slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003049
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003050 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3051 All counting is done with python objects (no overflows or underflows).
3052 Advances with: long_cnt += long_step
3053 Step may be zero -- effectively a slow version of repeat(cnt).
3054 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003055*/
3056
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003057static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003058
3059static PyObject *
3060count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3061{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003062 countobject *lz;
3063 int slow_mode = 0;
3064 Py_ssize_t cnt = 0;
3065 PyObject *long_cnt = NULL;
3066 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003067 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003068 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003069
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003070 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3071 kwlist, &long_cnt, &long_step))
3072 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003073
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003074 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3075 (long_step != NULL && !PyNumber_Check(long_step))) {
3076 PyErr_SetString(PyExc_TypeError, "a number is required");
3077 return NULL;
3078 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003079
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003080 if (long_cnt != NULL) {
3081 cnt = PyLong_AsSsize_t(long_cnt);
3082 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3083 PyErr_Clear();
3084 slow_mode = 1;
3085 }
3086 Py_INCREF(long_cnt);
3087 } else {
3088 cnt = 0;
3089 long_cnt = PyLong_FromLong(0);
3090 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003091
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003092 /* If not specified, step defaults to 1 */
3093 if (long_step == NULL) {
3094 long_step = PyLong_FromLong(1);
3095 if (long_step == NULL) {
3096 Py_DECREF(long_cnt);
3097 return NULL;
3098 }
3099 } else
3100 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003101
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003102 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003103
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003104 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003105 step = PyLong_AsLong(long_step);
3106 if (step != 1) {
3107 slow_mode = 1;
3108 if (step == -1 && PyErr_Occurred())
3109 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003110 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003111
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003112 if (slow_mode)
3113 cnt = PY_SSIZE_T_MAX;
3114 else
3115 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003116
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003117 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3118 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3119 assert(slow_mode ||
3120 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003121
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003122 /* create countobject structure */
3123 lz = (countobject *)type->tp_alloc(type, 0);
3124 if (lz == NULL) {
3125 Py_XDECREF(long_cnt);
3126 return NULL;
3127 }
3128 lz->cnt = cnt;
3129 lz->long_cnt = long_cnt;
3130 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003131
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003132 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003133}
3134
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003135static void
3136count_dealloc(countobject *lz)
3137{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003138 PyObject_GC_UnTrack(lz);
3139 Py_XDECREF(lz->long_cnt);
3140 Py_XDECREF(lz->long_step);
3141 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003142}
3143
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003144static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003145count_traverse(countobject *lz, visitproc visit, void *arg)
3146{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003147 Py_VISIT(lz->long_cnt);
3148 Py_VISIT(lz->long_step);
3149 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003150}
3151
3152static PyObject *
3153count_nextlong(countobject *lz)
3154{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003155 PyObject *long_cnt;
3156 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003157
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003158 long_cnt = lz->long_cnt;
3159 if (long_cnt == NULL) {
3160 /* Switch to slow_mode */
3161 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3162 if (long_cnt == NULL)
3163 return NULL;
3164 }
3165 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00003166
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003167 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3168 if (stepped_up == NULL)
3169 return NULL;
3170 lz->long_cnt = stepped_up;
3171 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003172}
3173
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003174static PyObject *
3175count_next(countobject *lz)
3176{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003177 if (lz->cnt == PY_SSIZE_T_MAX)
3178 return count_nextlong(lz);
3179 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003180}
3181
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003182static PyObject *
3183count_repr(countobject *lz)
3184{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003185 if (lz->cnt != PY_SSIZE_T_MAX)
3186 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003187
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003188 if (PyLong_Check(lz->long_step)) {
3189 long step = PyLong_AsLong(lz->long_step);
3190 if (step == -1 && PyErr_Occurred()) {
3191 PyErr_Clear();
3192 }
3193 if (step == 1) {
3194 /* Don't display step when it is an integer equal to 1 */
3195 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
3196 }
3197 }
3198 return PyUnicode_FromFormat("count(%R, %R)",
3199 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003200}
3201
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003202static PyObject *
3203count_reduce(countobject *lz)
3204{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003205 if (lz->cnt == PY_SSIZE_T_MAX)
3206 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3207 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003208}
3209
3210PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3211
3212static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003213 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3214 count_reduce_doc},
3215 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003216};
3217
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003218PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00003219 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003220\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003221Return a count object whose .__next__() method returns consecutive values.\n\
3222Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00003223 def count(firstval=0, step=1):\n\
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003224 x = firstval\n\
3225 while 1:\n\
3226 yield x\n\
3227 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003228
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003229static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003230 PyVarObject_HEAD_INIT(NULL, 0)
3231 "itertools.count", /* tp_name */
3232 sizeof(countobject), /* tp_basicsize */
3233 0, /* tp_itemsize */
3234 /* methods */
3235 (destructor)count_dealloc, /* tp_dealloc */
3236 0, /* tp_print */
3237 0, /* tp_getattr */
3238 0, /* tp_setattr */
3239 0, /* tp_reserved */
3240 (reprfunc)count_repr, /* tp_repr */
3241 0, /* tp_as_number */
3242 0, /* tp_as_sequence */
3243 0, /* tp_as_mapping */
3244 0, /* tp_hash */
3245 0, /* tp_call */
3246 0, /* tp_str */
3247 PyObject_GenericGetAttr, /* tp_getattro */
3248 0, /* tp_setattro */
3249 0, /* tp_as_buffer */
3250 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3251 Py_TPFLAGS_BASETYPE, /* tp_flags */
3252 count_doc, /* tp_doc */
3253 (traverseproc)count_traverse, /* tp_traverse */
3254 0, /* tp_clear */
3255 0, /* tp_richcompare */
3256 0, /* tp_weaklistoffset */
3257 PyObject_SelfIter, /* tp_iter */
3258 (iternextfunc)count_next, /* tp_iternext */
3259 count_methods, /* tp_methods */
3260 0, /* tp_members */
3261 0, /* tp_getset */
3262 0, /* tp_base */
3263 0, /* tp_dict */
3264 0, /* tp_descr_get */
3265 0, /* tp_descr_set */
3266 0, /* tp_dictoffset */
3267 0, /* tp_init */
3268 0, /* tp_alloc */
3269 count_new, /* tp_new */
3270 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003271};
3272
3273
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003274/* repeat object ************************************************************/
3275
3276typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003277 PyObject_HEAD
3278 PyObject *element;
3279 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003280} repeatobject;
3281
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003282static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003283
3284static PyObject *
3285repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3286{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003287 repeatobject *ro;
3288 PyObject *element;
3289 Py_ssize_t cnt = -1;
3290 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003291
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003292 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3293 &element, &cnt))
3294 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003295
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003296 if (PyTuple_Size(args) == 2 && cnt < 0)
3297 cnt = 0;
3298
3299 ro = (repeatobject *)type->tp_alloc(type, 0);
3300 if (ro == NULL)
3301 return NULL;
3302 Py_INCREF(element);
3303 ro->element = element;
3304 ro->cnt = cnt;
3305 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003306}
3307
3308static void
3309repeat_dealloc(repeatobject *ro)
3310{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003311 PyObject_GC_UnTrack(ro);
3312 Py_XDECREF(ro->element);
3313 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003314}
3315
3316static int
3317repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3318{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003319 Py_VISIT(ro->element);
3320 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003321}
3322
3323static PyObject *
3324repeat_next(repeatobject *ro)
3325{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003326 if (ro->cnt == 0)
3327 return NULL;
3328 if (ro->cnt > 0)
3329 ro->cnt--;
3330 Py_INCREF(ro->element);
3331 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003332}
3333
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003334static PyObject *
3335repeat_repr(repeatobject *ro)
3336{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003337 if (ro->cnt == -1)
3338 return PyUnicode_FromFormat("repeat(%R)", ro->element);
3339 else
3340 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
3341}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003342
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003343static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003344repeat_len(repeatobject *ro)
3345{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003346 if (ro->cnt == -1) {
3347 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3348 return NULL;
3349 }
3350 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003351}
3352
Armin Rigof5b3e362006-02-11 21:32:43 +00003353PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003354
3355static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003356 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3357 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003358};
3359
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003360PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00003361"repeat(object [,times]) -> create an iterator which returns the object\n\
3362for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003363endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003364
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003365static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003366 PyVarObject_HEAD_INIT(NULL, 0)
3367 "itertools.repeat", /* tp_name */
3368 sizeof(repeatobject), /* tp_basicsize */
3369 0, /* tp_itemsize */
3370 /* methods */
3371 (destructor)repeat_dealloc, /* tp_dealloc */
3372 0, /* tp_print */
3373 0, /* tp_getattr */
3374 0, /* tp_setattr */
3375 0, /* tp_reserved */
3376 (reprfunc)repeat_repr, /* tp_repr */
3377 0, /* tp_as_number */
3378 0, /* tp_as_sequence */
3379 0, /* tp_as_mapping */
3380 0, /* tp_hash */
3381 0, /* tp_call */
3382 0, /* tp_str */
3383 PyObject_GenericGetAttr, /* tp_getattro */
3384 0, /* tp_setattro */
3385 0, /* tp_as_buffer */
3386 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3387 Py_TPFLAGS_BASETYPE, /* tp_flags */
3388 repeat_doc, /* tp_doc */
3389 (traverseproc)repeat_traverse, /* tp_traverse */
3390 0, /* tp_clear */
3391 0, /* tp_richcompare */
3392 0, /* tp_weaklistoffset */
3393 PyObject_SelfIter, /* tp_iter */
3394 (iternextfunc)repeat_next, /* tp_iternext */
3395 repeat_methods, /* tp_methods */
3396 0, /* tp_members */
3397 0, /* tp_getset */
3398 0, /* tp_base */
3399 0, /* tp_dict */
3400 0, /* tp_descr_get */
3401 0, /* tp_descr_set */
3402 0, /* tp_dictoffset */
3403 0, /* tp_init */
3404 0, /* tp_alloc */
3405 repeat_new, /* tp_new */
3406 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003407};
3408
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003409/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00003410
3411#include "Python.h"
3412
3413typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003414 PyObject_HEAD
3415 Py_ssize_t tuplesize;
3416 Py_ssize_t numactive;
3417 PyObject *ittuple; /* tuple of iterators */
3418 PyObject *result;
3419 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003420} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003421
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003422static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003423
3424static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003425zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003426{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003427 ziplongestobject *lz;
3428 Py_ssize_t i;
3429 PyObject *ittuple; /* tuple of iterators */
3430 PyObject *result;
3431 PyObject *fillvalue = Py_None;
3432 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003433
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003434 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3435 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3436 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3437 PyErr_SetString(PyExc_TypeError,
3438 "zip_longest() got an unexpected keyword argument");
3439 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003440 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003441 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003442
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003443 /* args must be a tuple */
3444 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00003445
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003446 /* obtain iterators */
3447 ittuple = PyTuple_New(tuplesize);
3448 if (ittuple == NULL)
3449 return NULL;
3450 for (i=0; i < tuplesize; ++i) {
3451 PyObject *item = PyTuple_GET_ITEM(args, i);
3452 PyObject *it = PyObject_GetIter(item);
3453 if (it == NULL) {
3454 if (PyErr_ExceptionMatches(PyExc_TypeError))
3455 PyErr_Format(PyExc_TypeError,
3456 "zip_longest argument #%zd must support iteration",
3457 i+1);
3458 Py_DECREF(ittuple);
3459 return NULL;
3460 }
3461 PyTuple_SET_ITEM(ittuple, i, it);
3462 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003463
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003464 /* create a result holder */
3465 result = PyTuple_New(tuplesize);
3466 if (result == NULL) {
3467 Py_DECREF(ittuple);
3468 return NULL;
3469 }
3470 for (i=0 ; i < tuplesize ; i++) {
3471 Py_INCREF(Py_None);
3472 PyTuple_SET_ITEM(result, i, Py_None);
3473 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003474
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003475 /* create ziplongestobject structure */
3476 lz = (ziplongestobject *)type->tp_alloc(type, 0);
3477 if (lz == NULL) {
3478 Py_DECREF(ittuple);
3479 Py_DECREF(result);
3480 return NULL;
3481 }
3482 lz->ittuple = ittuple;
3483 lz->tuplesize = tuplesize;
3484 lz->numactive = tuplesize;
3485 lz->result = result;
3486 Py_INCREF(fillvalue);
3487 lz->fillvalue = fillvalue;
3488 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003489}
3490
3491static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003492zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003493{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003494 PyObject_GC_UnTrack(lz);
3495 Py_XDECREF(lz->ittuple);
3496 Py_XDECREF(lz->result);
3497 Py_XDECREF(lz->fillvalue);
3498 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003499}
3500
3501static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003502zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003503{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003504 Py_VISIT(lz->ittuple);
3505 Py_VISIT(lz->result);
3506 Py_VISIT(lz->fillvalue);
3507 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003508}
3509
3510static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003511zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003512{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003513 Py_ssize_t i;
3514 Py_ssize_t tuplesize = lz->tuplesize;
3515 PyObject *result = lz->result;
3516 PyObject *it;
3517 PyObject *item;
3518 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003519
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003520 if (tuplesize == 0)
3521 return NULL;
3522 if (lz->numactive == 0)
3523 return NULL;
3524 if (Py_REFCNT(result) == 1) {
3525 Py_INCREF(result);
3526 for (i=0 ; i < tuplesize ; i++) {
3527 it = PyTuple_GET_ITEM(lz->ittuple, i);
3528 if (it == NULL) {
3529 Py_INCREF(lz->fillvalue);
3530 item = lz->fillvalue;
3531 } else {
3532 item = PyIter_Next(it);
3533 if (item == NULL) {
3534 lz->numactive -= 1;
3535 if (lz->numactive == 0 || PyErr_Occurred()) {
3536 lz->numactive = 0;
3537 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00003538 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003539 } else {
3540 Py_INCREF(lz->fillvalue);
3541 item = lz->fillvalue;
3542 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3543 Py_DECREF(it);
3544 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00003545 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003546 }
3547 olditem = PyTuple_GET_ITEM(result, i);
3548 PyTuple_SET_ITEM(result, i, item);
3549 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00003550 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003551 } else {
3552 result = PyTuple_New(tuplesize);
3553 if (result == NULL)
3554 return NULL;
3555 for (i=0 ; i < tuplesize ; i++) {
3556 it = PyTuple_GET_ITEM(lz->ittuple, i);
3557 if (it == NULL) {
3558 Py_INCREF(lz->fillvalue);
3559 item = lz->fillvalue;
3560 } else {
3561 item = PyIter_Next(it);
3562 if (item == NULL) {
3563 lz->numactive -= 1;
3564 if (lz->numactive == 0 || PyErr_Occurred()) {
3565 lz->numactive = 0;
3566 Py_DECREF(result);
3567 return NULL;
3568 } else {
3569 Py_INCREF(lz->fillvalue);
3570 item = lz->fillvalue;
3571 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3572 Py_DECREF(it);
3573 }
3574 }
3575 }
3576 PyTuple_SET_ITEM(result, i, item);
3577 }
3578 }
3579 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003580}
3581
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003582PyDoc_STRVAR(zip_longest_doc,
3583"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003584\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003585Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00003586the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003587method continues until the longest iterable in the argument sequence\n\
3588is exhausted and then it raises StopIteration. When the shorter iterables\n\
3589are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3590defaults to None or can be specified by a keyword argument.\n\
3591");
3592
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003593static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003594 PyVarObject_HEAD_INIT(NULL, 0)
3595 "itertools.zip_longest", /* tp_name */
3596 sizeof(ziplongestobject), /* tp_basicsize */
3597 0, /* tp_itemsize */
3598 /* methods */
3599 (destructor)zip_longest_dealloc, /* tp_dealloc */
3600 0, /* tp_print */
3601 0, /* tp_getattr */
3602 0, /* tp_setattr */
3603 0, /* tp_reserved */
3604 0, /* tp_repr */
3605 0, /* tp_as_number */
3606 0, /* tp_as_sequence */
3607 0, /* tp_as_mapping */
3608 0, /* tp_hash */
3609 0, /* tp_call */
3610 0, /* tp_str */
3611 PyObject_GenericGetAttr, /* tp_getattro */
3612 0, /* tp_setattro */
3613 0, /* tp_as_buffer */
3614 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3615 Py_TPFLAGS_BASETYPE, /* tp_flags */
3616 zip_longest_doc, /* tp_doc */
3617 (traverseproc)zip_longest_traverse, /* tp_traverse */
3618 0, /* tp_clear */
3619 0, /* tp_richcompare */
3620 0, /* tp_weaklistoffset */
3621 PyObject_SelfIter, /* tp_iter */
3622 (iternextfunc)zip_longest_next, /* tp_iternext */
3623 0, /* tp_methods */
3624 0, /* tp_members */
3625 0, /* tp_getset */
3626 0, /* tp_base */
3627 0, /* tp_dict */
3628 0, /* tp_descr_get */
3629 0, /* tp_descr_set */
3630 0, /* tp_dictoffset */
3631 0, /* tp_init */
3632 0, /* tp_alloc */
3633 zip_longest_new, /* tp_new */
3634 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003635};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003636
3637/* module level code ********************************************************/
3638
3639PyDoc_STRVAR(module_doc,
3640"Functional tools for creating and using iterators.\n\
3641\n\
3642Infinite iterators:\n\
3643count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003644cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003645repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003646\n\
3647Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003648accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003649chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
3650compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
3651dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
3652groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003653filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003654islice(seq, [start,] stop [, step]) --> elements from\n\
3655 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003656starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003657tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003658takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003659zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00003660\n\
3661Combinatoric generators:\n\
3662product(p, q, ... [repeat=1]) --> cartesian product\n\
3663permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00003664combinations(p, r)\n\
3665combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003666");
3667
3668
Raymond Hettingerad983e72003-11-12 14:32:26 +00003669static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003670 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3671 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00003672};
3673
Martin v. Löwis1a214512008-06-11 05:26:20 +00003674
3675static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003676 PyModuleDef_HEAD_INIT,
3677 "itertools",
3678 module_doc,
3679 -1,
3680 module_methods,
3681 NULL,
3682 NULL,
3683 NULL,
3684 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00003685};
3686
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003687PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00003688PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003689{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003690 int i;
3691 PyObject *m;
3692 char *name;
3693 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00003694 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003695 &combinations_type,
3696 &cwr_type,
3697 &cycle_type,
3698 &dropwhile_type,
3699 &takewhile_type,
3700 &islice_type,
3701 &starmap_type,
3702 &chain_type,
3703 &compress_type,
3704 &filterfalse_type,
3705 &count_type,
3706 &ziplongest_type,
3707 &permutations_type,
3708 &product_type,
3709 &repeat_type,
3710 &groupby_type,
3711 NULL
3712 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00003713
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003714 Py_TYPE(&teedataobject_type) = &PyType_Type;
3715 m = PyModule_Create(&itertoolsmodule);
3716 if (m == NULL)
3717 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003718
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003719 for (i=0 ; typelist[i] != NULL ; i++) {
3720 if (PyType_Ready(typelist[i]) < 0)
3721 return NULL;
3722 name = strchr(typelist[i]->tp_name, '.');
3723 assert (name != NULL);
3724 Py_INCREF(typelist[i]);
3725 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
3726 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003727
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003728 if (PyType_Ready(&teedataobject_type) < 0)
3729 return NULL;
3730 if (PyType_Ready(&tee_type) < 0)
3731 return NULL;
3732 if (PyType_Ready(&_grouper_type) < 0)
3733 return NULL;
3734 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003735}