blob: d0897c3eb956f3313413af804b63fa12b1f9f174 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
2#include "Python.h"
Fred Drake08ebfec2004-10-17 19:36:57 +00003#include "structmember.h"
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00005/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00006 by Raymond D. Hettinger <python@rcn.com>
7 Copyright (c) 2003 Python Software Foundation.
8 All rights reserved.
9*/
10
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000011
12/* groupby object ***********************************************************/
13
14typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000015 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000021} groupbyobject;
22
23static PyTypeObject groupby_type;
24static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26static PyObject *
27groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000029 static char *kwargs[] = {"iterable", "key", NULL};
30 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000032
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000033 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
34 &it, &keyfunc))
35 return NULL;
36
37 gbo = (groupbyobject *)type->tp_alloc(type, 0);
38 if (gbo == NULL)
39 return NULL;
40 gbo->tgtkey = NULL;
41 gbo->currkey = NULL;
42 gbo->currvalue = NULL;
43 gbo->keyfunc = keyfunc;
44 Py_INCREF(keyfunc);
45 gbo->it = PyObject_GetIter(it);
46 if (gbo->it == NULL) {
47 Py_DECREF(gbo);
48 return NULL;
49 }
50 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000051}
52
53static void
54groupby_dealloc(groupbyobject *gbo)
55{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000056 PyObject_GC_UnTrack(gbo);
57 Py_XDECREF(gbo->it);
58 Py_XDECREF(gbo->keyfunc);
59 Py_XDECREF(gbo->tgtkey);
60 Py_XDECREF(gbo->currkey);
61 Py_XDECREF(gbo->currvalue);
62 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000063}
64
65static int
66groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
67{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000068 Py_VISIT(gbo->it);
69 Py_VISIT(gbo->keyfunc);
70 Py_VISIT(gbo->tgtkey);
71 Py_VISIT(gbo->currkey);
72 Py_VISIT(gbo->currvalue);
73 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000074}
75
76static PyObject *
77groupby_next(groupbyobject *gbo)
78{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000081 /* skip to next iteration group */
82 for (;;) {
83 if (gbo->currkey == NULL)
84 /* pass */;
85 else if (gbo->tgtkey == NULL)
86 break;
87 else {
88 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000089
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000090 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
91 gbo->currkey, Py_EQ);
92 if (rcmp == -1)
93 return NULL;
94 else if (rcmp == 0)
95 break;
96 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000097
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000098 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000101
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000102 if (gbo->keyfunc == Py_None) {
103 newkey = newvalue;
104 Py_INCREF(newvalue);
105 } else {
106 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
107 newvalue, NULL);
108 if (newkey == NULL) {
109 Py_DECREF(newvalue);
110 return NULL;
111 }
112 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000113
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000114 tmp = gbo->currkey;
115 gbo->currkey = newkey;
116 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000117
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000118 tmp = gbo->currvalue;
119 gbo->currvalue = newvalue;
120 Py_XDECREF(tmp);
121 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000123 Py_INCREF(gbo->currkey);
124 tmp = gbo->tgtkey;
125 gbo->tgtkey = gbo->currkey;
126 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000127
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000131
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000135}
136
137PyDoc_STRVAR(groupby_doc,
138"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
139(key, sub-iterator) grouped by each value of key(value).\n");
140
141static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000142 PyVarObject_HEAD_INIT(NULL, 0)
143 "itertools.groupby", /* tp_name */
144 sizeof(groupbyobject), /* tp_basicsize */
145 0, /* tp_itemsize */
146 /* methods */
147 (destructor)groupby_dealloc, /* tp_dealloc */
148 0, /* tp_print */
149 0, /* tp_getattr */
150 0, /* tp_setattr */
151 0, /* tp_reserved */
152 0, /* tp_repr */
153 0, /* tp_as_number */
154 0, /* tp_as_sequence */
155 0, /* tp_as_mapping */
156 0, /* tp_hash */
157 0, /* tp_call */
158 0, /* tp_str */
159 PyObject_GenericGetAttr, /* tp_getattro */
160 0, /* tp_setattro */
161 0, /* tp_as_buffer */
162 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
163 Py_TPFLAGS_BASETYPE, /* tp_flags */
164 groupby_doc, /* tp_doc */
165 (traverseproc)groupby_traverse, /* tp_traverse */
166 0, /* tp_clear */
167 0, /* tp_richcompare */
168 0, /* tp_weaklistoffset */
169 PyObject_SelfIter, /* tp_iter */
170 (iternextfunc)groupby_next, /* tp_iternext */
171 0, /* tp_methods */
172 0, /* tp_members */
173 0, /* tp_getset */
174 0, /* tp_base */
175 0, /* tp_dict */
176 0, /* tp_descr_get */
177 0, /* tp_descr_set */
178 0, /* tp_dictoffset */
179 0, /* tp_init */
180 0, /* tp_alloc */
181 groupby_new, /* tp_new */
182 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000183};
184
185
186/* _grouper object (internal) ************************************************/
187
188typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000189 PyObject_HEAD
190 PyObject *parent;
191 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000192} _grouperobject;
193
194static PyTypeObject _grouper_type;
195
196static PyObject *
197_grouper_create(groupbyobject *parent, PyObject *tgtkey)
198{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000199 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000200
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000201 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
202 if (igo == NULL)
203 return NULL;
204 igo->parent = (PyObject *)parent;
205 Py_INCREF(parent);
206 igo->tgtkey = tgtkey;
207 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000208
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000209 PyObject_GC_Track(igo);
210 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000211}
212
213static void
214_grouper_dealloc(_grouperobject *igo)
215{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000216 PyObject_GC_UnTrack(igo);
217 Py_DECREF(igo->parent);
218 Py_DECREF(igo->tgtkey);
219 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000220}
221
222static int
223_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
224{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000225 Py_VISIT(igo->parent);
226 Py_VISIT(igo->tgtkey);
227 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000228}
229
230static PyObject *
231_grouper_next(_grouperobject *igo)
232{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000233 groupbyobject *gbo = (groupbyobject *)igo->parent;
234 PyObject *newvalue, *newkey, *r;
235 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000236
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000237 if (gbo->currvalue == NULL) {
238 newvalue = PyIter_Next(gbo->it);
239 if (newvalue == NULL)
240 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000241
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000242 if (gbo->keyfunc == Py_None) {
243 newkey = newvalue;
244 Py_INCREF(newvalue);
245 } else {
246 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
247 newvalue, NULL);
248 if (newkey == NULL) {
249 Py_DECREF(newvalue);
250 return NULL;
251 }
252 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000253
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000254 assert(gbo->currkey == NULL);
255 gbo->currkey = newkey;
256 gbo->currvalue = newvalue;
257 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000259 assert(gbo->currkey != NULL);
260 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
261 if (rcmp <= 0)
262 /* got any error or current group is end */
263 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000264
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000265 r = gbo->currvalue;
266 gbo->currvalue = NULL;
267 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000268
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000269 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000270}
271
272static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000273 PyVarObject_HEAD_INIT(NULL, 0)
274 "itertools._grouper", /* tp_name */
275 sizeof(_grouperobject), /* tp_basicsize */
276 0, /* tp_itemsize */
277 /* methods */
278 (destructor)_grouper_dealloc, /* tp_dealloc */
279 0, /* tp_print */
280 0, /* tp_getattr */
281 0, /* tp_setattr */
282 0, /* tp_reserved */
283 0, /* tp_repr */
284 0, /* tp_as_number */
285 0, /* tp_as_sequence */
286 0, /* tp_as_mapping */
287 0, /* tp_hash */
288 0, /* tp_call */
289 0, /* tp_str */
290 PyObject_GenericGetAttr, /* tp_getattro */
291 0, /* tp_setattro */
292 0, /* tp_as_buffer */
293 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
294 0, /* tp_doc */
295 (traverseproc)_grouper_traverse,/* tp_traverse */
296 0, /* tp_clear */
297 0, /* tp_richcompare */
298 0, /* tp_weaklistoffset */
299 PyObject_SelfIter, /* tp_iter */
300 (iternextfunc)_grouper_next, /* tp_iternext */
301 0, /* tp_methods */
302 0, /* tp_members */
303 0, /* tp_getset */
304 0, /* tp_base */
305 0, /* tp_dict */
306 0, /* tp_descr_get */
307 0, /* tp_descr_set */
308 0, /* tp_dictoffset */
309 0, /* tp_init */
310 0, /* tp_alloc */
311 0, /* tp_new */
312 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000313};
314
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000315
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000316
Raymond Hettingerad983e72003-11-12 14:32:26 +0000317/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000318
Raymond Hettingerad983e72003-11-12 14:32:26 +0000319/* The teedataobject pre-allocates space for LINKCELLS number of objects.
320 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000321 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000322 the other structure members including PyHEAD overhead. The larger the
323 value, the less memory overhead per object and the less time spent
324 allocating/deallocating new links. The smaller the number, the less
325 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000326*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000327#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000328
329typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000330 PyObject_HEAD
331 PyObject *it;
332 int numread;
333 PyObject *nextlink;
334 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000335} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000336
337typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000338 PyObject_HEAD
339 teedataobject *dataobj;
340 int index;
341 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000342} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000343
Raymond Hettingerad983e72003-11-12 14:32:26 +0000344static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000345
346static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000347teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000348{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000349 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000350
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000351 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
352 if (tdo == NULL)
353 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000354
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000355 tdo->numread = 0;
356 tdo->nextlink = NULL;
357 Py_INCREF(it);
358 tdo->it = it;
359 PyObject_GC_Track(tdo);
360 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000361}
362
363static PyObject *
364teedataobject_jumplink(teedataobject *tdo)
365{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000366 if (tdo->nextlink == NULL)
367 tdo->nextlink = teedataobject_new(tdo->it);
368 Py_XINCREF(tdo->nextlink);
369 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000370}
371
372static PyObject *
373teedataobject_getitem(teedataobject *tdo, int i)
374{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000375 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000376
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000377 assert(i < LINKCELLS);
378 if (i < tdo->numread)
379 value = tdo->values[i];
380 else {
381 /* this is the lead iterator, so fetch more data */
382 assert(i == tdo->numread);
383 value = PyIter_Next(tdo->it);
384 if (value == NULL)
385 return NULL;
386 tdo->numread++;
387 tdo->values[i] = value;
388 }
389 Py_INCREF(value);
390 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000391}
392
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000393static int
394teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
395{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000396 int i;
397 Py_VISIT(tdo->it);
398 for (i = 0; i < tdo->numread; i++)
399 Py_VISIT(tdo->values[i]);
400 Py_VISIT(tdo->nextlink);
401 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000402}
403
404static int
405teedataobject_clear(teedataobject *tdo)
406{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000407 int i;
408 Py_CLEAR(tdo->it);
409 for (i=0 ; i<tdo->numread ; i++)
410 Py_CLEAR(tdo->values[i]);
411 Py_CLEAR(tdo->nextlink);
412 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000413}
414
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000416teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000418 PyObject_GC_UnTrack(tdo);
419 teedataobject_clear(tdo);
420 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000421}
422
Raymond Hettingerad983e72003-11-12 14:32:26 +0000423PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000424
Raymond Hettingerad983e72003-11-12 14:32:26 +0000425static PyTypeObject teedataobject_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000426 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
427 "itertools.tee_dataobject", /* tp_name */
428 sizeof(teedataobject), /* tp_basicsize */
429 0, /* tp_itemsize */
430 /* methods */
431 (destructor)teedataobject_dealloc, /* tp_dealloc */
432 0, /* tp_print */
433 0, /* tp_getattr */
434 0, /* tp_setattr */
435 0, /* tp_reserved */
436 0, /* tp_repr */
437 0, /* tp_as_number */
438 0, /* tp_as_sequence */
439 0, /* tp_as_mapping */
440 0, /* tp_hash */
441 0, /* tp_call */
442 0, /* tp_str */
443 PyObject_GenericGetAttr, /* tp_getattro */
444 0, /* tp_setattro */
445 0, /* tp_as_buffer */
446 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
447 teedataobject_doc, /* tp_doc */
448 (traverseproc)teedataobject_traverse, /* tp_traverse */
449 (inquiry)teedataobject_clear, /* tp_clear */
450 0, /* tp_richcompare */
451 0, /* tp_weaklistoffset */
452 0, /* tp_iter */
453 0, /* tp_iternext */
454 0, /* tp_methods */
455 0, /* tp_members */
456 0, /* tp_getset */
457 0, /* tp_base */
458 0, /* tp_dict */
459 0, /* tp_descr_get */
460 0, /* tp_descr_set */
461 0, /* tp_dictoffset */
462 0, /* tp_init */
463 0, /* tp_alloc */
464 0, /* tp_new */
465 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000466};
467
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000468
469static PyTypeObject tee_type;
470
471static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000472tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000473{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000474 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000475
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000476 if (to->index >= LINKCELLS) {
477 link = teedataobject_jumplink(to->dataobj);
478 Py_DECREF(to->dataobj);
479 to->dataobj = (teedataobject *)link;
480 to->index = 0;
481 }
482 value = teedataobject_getitem(to->dataobj, to->index);
483 if (value == NULL)
484 return NULL;
485 to->index++;
486 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000487}
488
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000489static int
490tee_traverse(teeobject *to, visitproc visit, void *arg)
491{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000492 Py_VISIT((PyObject *)to->dataobj);
493 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000494}
495
Raymond Hettingerad983e72003-11-12 14:32:26 +0000496static PyObject *
497tee_copy(teeobject *to)
498{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000499 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000500
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000501 newto = PyObject_GC_New(teeobject, &tee_type);
502 if (newto == NULL)
503 return NULL;
504 Py_INCREF(to->dataobj);
505 newto->dataobj = to->dataobj;
506 newto->index = to->index;
507 newto->weakreflist = NULL;
508 PyObject_GC_Track(newto);
509 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000510}
511
512PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
513
514static PyObject *
515tee_fromiterable(PyObject *iterable)
516{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000517 teeobject *to;
518 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000519
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000520 it = PyObject_GetIter(iterable);
521 if (it == NULL)
522 return NULL;
523 if (PyObject_TypeCheck(it, &tee_type)) {
524 to = (teeobject *)tee_copy((teeobject *)it);
525 goto done;
526 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000527
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000528 to = PyObject_GC_New(teeobject, &tee_type);
529 if (to == NULL)
530 goto done;
531 to->dataobj = (teedataobject *)teedataobject_new(it);
532 if (!to->dataobj) {
533 PyObject_GC_Del(to);
534 to = NULL;
535 goto done;
536 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000537
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000538 to->index = 0;
539 to->weakreflist = NULL;
540 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000541done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000542 Py_XDECREF(it);
543 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000544}
545
546static PyObject *
547tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
548{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000549 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000550
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000551 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
552 return NULL;
553 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000554}
555
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000556static int
557tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000558{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000559 if (to->weakreflist != NULL)
560 PyObject_ClearWeakRefs((PyObject *) to);
561 Py_CLEAR(to->dataobj);
562 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000563}
564
565static void
566tee_dealloc(teeobject *to)
567{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000568 PyObject_GC_UnTrack(to);
569 tee_clear(to);
570 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000571}
572
Raymond Hettingerad983e72003-11-12 14:32:26 +0000573PyDoc_STRVAR(teeobject_doc,
574"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000575
Raymond Hettingerad983e72003-11-12 14:32:26 +0000576static PyMethodDef tee_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000577 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
578 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000579};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000580
581static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000582 PyVarObject_HEAD_INIT(NULL, 0)
583 "itertools.tee", /* tp_name */
584 sizeof(teeobject), /* tp_basicsize */
585 0, /* tp_itemsize */
586 /* methods */
587 (destructor)tee_dealloc, /* tp_dealloc */
588 0, /* tp_print */
589 0, /* tp_getattr */
590 0, /* tp_setattr */
591 0, /* tp_reserved */
592 0, /* tp_repr */
593 0, /* tp_as_number */
594 0, /* tp_as_sequence */
595 0, /* tp_as_mapping */
596 0, /* tp_hash */
597 0, /* tp_call */
598 0, /* tp_str */
599 0, /* tp_getattro */
600 0, /* tp_setattro */
601 0, /* tp_as_buffer */
602 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
603 teeobject_doc, /* tp_doc */
604 (traverseproc)tee_traverse, /* tp_traverse */
605 (inquiry)tee_clear, /* tp_clear */
606 0, /* tp_richcompare */
607 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
608 PyObject_SelfIter, /* tp_iter */
609 (iternextfunc)tee_next, /* tp_iternext */
610 tee_methods, /* tp_methods */
611 0, /* tp_members */
612 0, /* tp_getset */
613 0, /* tp_base */
614 0, /* tp_dict */
615 0, /* tp_descr_get */
616 0, /* tp_descr_set */
617 0, /* tp_dictoffset */
618 0, /* tp_init */
619 0, /* tp_alloc */
620 tee_new, /* tp_new */
621 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000622};
623
Raymond Hettingerad983e72003-11-12 14:32:26 +0000624static PyObject *
625tee(PyObject *self, PyObject *args)
626{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000627 Py_ssize_t i, n=2;
628 PyObject *it, *iterable, *copyable, *result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000629
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000630 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
631 return NULL;
632 if (n < 0) {
633 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
634 return NULL;
635 }
636 result = PyTuple_New(n);
637 if (result == NULL)
638 return NULL;
639 if (n == 0)
640 return result;
641 it = PyObject_GetIter(iterable);
642 if (it == NULL) {
643 Py_DECREF(result);
644 return NULL;
645 }
646 if (!PyObject_HasAttrString(it, "__copy__")) {
647 copyable = tee_fromiterable(it);
648 Py_DECREF(it);
649 if (copyable == NULL) {
650 Py_DECREF(result);
651 return NULL;
652 }
653 } else
654 copyable = it;
655 PyTuple_SET_ITEM(result, 0, copyable);
656 for (i=1 ; i<n ; i++) {
657 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
658 if (copyable == NULL) {
659 Py_DECREF(result);
660 return NULL;
661 }
662 PyTuple_SET_ITEM(result, i, copyable);
663 }
664 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000665}
666
667PyDoc_STRVAR(tee_doc,
668"tee(iterable, n=2) --> tuple of n independent iterators.");
669
670
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000671/* cycle object **********************************************************/
672
673typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000674 PyObject_HEAD
675 PyObject *it;
676 PyObject *saved;
677 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000678} cycleobject;
679
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000680static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000681
682static PyObject *
683cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
684{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000685 PyObject *it;
686 PyObject *iterable;
687 PyObject *saved;
688 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000689
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000690 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
691 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000692
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000693 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
694 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000695
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000696 /* Get iterator. */
697 it = PyObject_GetIter(iterable);
698 if (it == NULL)
699 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000700
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000701 saved = PyList_New(0);
702 if (saved == NULL) {
703 Py_DECREF(it);
704 return NULL;
705 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000706
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000707 /* create cycleobject structure */
708 lz = (cycleobject *)type->tp_alloc(type, 0);
709 if (lz == NULL) {
710 Py_DECREF(it);
711 Py_DECREF(saved);
712 return NULL;
713 }
714 lz->it = it;
715 lz->saved = saved;
716 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000717
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000718 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000719}
720
721static void
722cycle_dealloc(cycleobject *lz)
723{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000724 PyObject_GC_UnTrack(lz);
725 Py_XDECREF(lz->saved);
726 Py_XDECREF(lz->it);
727 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000728}
729
730static int
731cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
732{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000733 Py_VISIT(lz->it);
734 Py_VISIT(lz->saved);
735 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000736}
737
738static PyObject *
739cycle_next(cycleobject *lz)
740{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000741 PyObject *item;
742 PyObject *it;
743 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000744
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000745 while (1) {
746 item = PyIter_Next(lz->it);
747 if (item != NULL) {
748 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
749 Py_DECREF(item);
750 return NULL;
751 }
752 return item;
753 }
754 if (PyErr_Occurred()) {
755 if (PyErr_ExceptionMatches(PyExc_StopIteration))
756 PyErr_Clear();
757 else
758 return NULL;
759 }
760 if (PyList_Size(lz->saved) == 0)
761 return NULL;
762 it = PyObject_GetIter(lz->saved);
763 if (it == NULL)
764 return NULL;
765 tmp = lz->it;
766 lz->it = it;
767 lz->firstpass = 1;
768 Py_DECREF(tmp);
769 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000770}
771
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000772PyDoc_STRVAR(cycle_doc,
773"cycle(iterable) --> cycle object\n\
774\n\
775Return elements from the iterable until it is exhausted.\n\
776Then repeat the sequence indefinitely.");
777
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000778static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000779 PyVarObject_HEAD_INIT(NULL, 0)
780 "itertools.cycle", /* tp_name */
781 sizeof(cycleobject), /* tp_basicsize */
782 0, /* tp_itemsize */
783 /* methods */
784 (destructor)cycle_dealloc, /* tp_dealloc */
785 0, /* tp_print */
786 0, /* tp_getattr */
787 0, /* tp_setattr */
788 0, /* tp_reserved */
789 0, /* tp_repr */
790 0, /* tp_as_number */
791 0, /* tp_as_sequence */
792 0, /* tp_as_mapping */
793 0, /* tp_hash */
794 0, /* tp_call */
795 0, /* tp_str */
796 PyObject_GenericGetAttr, /* tp_getattro */
797 0, /* tp_setattro */
798 0, /* tp_as_buffer */
799 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
800 Py_TPFLAGS_BASETYPE, /* tp_flags */
801 cycle_doc, /* tp_doc */
802 (traverseproc)cycle_traverse, /* tp_traverse */
803 0, /* tp_clear */
804 0, /* tp_richcompare */
805 0, /* tp_weaklistoffset */
806 PyObject_SelfIter, /* tp_iter */
807 (iternextfunc)cycle_next, /* tp_iternext */
808 0, /* tp_methods */
809 0, /* tp_members */
810 0, /* tp_getset */
811 0, /* tp_base */
812 0, /* tp_dict */
813 0, /* tp_descr_get */
814 0, /* tp_descr_set */
815 0, /* tp_dictoffset */
816 0, /* tp_init */
817 0, /* tp_alloc */
818 cycle_new, /* tp_new */
819 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000820};
821
822
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000823/* dropwhile object **********************************************************/
824
825typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000826 PyObject_HEAD
827 PyObject *func;
828 PyObject *it;
829 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000830} dropwhileobject;
831
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000832static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000833
834static PyObject *
835dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
836{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000837 PyObject *func, *seq;
838 PyObject *it;
839 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000840
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000841 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
842 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000843
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000844 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
845 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000846
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000847 /* Get iterator. */
848 it = PyObject_GetIter(seq);
849 if (it == NULL)
850 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000851
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000852 /* create dropwhileobject structure */
853 lz = (dropwhileobject *)type->tp_alloc(type, 0);
854 if (lz == NULL) {
855 Py_DECREF(it);
856 return NULL;
857 }
858 Py_INCREF(func);
859 lz->func = func;
860 lz->it = it;
861 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000862
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000863 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000864}
865
866static void
867dropwhile_dealloc(dropwhileobject *lz)
868{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000869 PyObject_GC_UnTrack(lz);
870 Py_XDECREF(lz->func);
871 Py_XDECREF(lz->it);
872 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000873}
874
875static int
876dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
877{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000878 Py_VISIT(lz->it);
879 Py_VISIT(lz->func);
880 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000881}
882
883static PyObject *
884dropwhile_next(dropwhileobject *lz)
885{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000886 PyObject *item, *good;
887 PyObject *it = lz->it;
888 long ok;
889 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000890
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000891 iternext = *Py_TYPE(it)->tp_iternext;
892 for (;;) {
893 item = iternext(it);
894 if (item == NULL)
895 return NULL;
896 if (lz->start == 1)
897 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000898
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000899 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
900 if (good == NULL) {
901 Py_DECREF(item);
902 return NULL;
903 }
904 ok = PyObject_IsTrue(good);
905 Py_DECREF(good);
906 if (!ok) {
907 lz->start = 1;
908 return item;
909 }
910 Py_DECREF(item);
911 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000912}
913
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000914PyDoc_STRVAR(dropwhile_doc,
915"dropwhile(predicate, iterable) --> dropwhile object\n\
916\n\
917Drop items from the iterable while predicate(item) is true.\n\
918Afterwards, return every element until the iterable is exhausted.");
919
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000920static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000921 PyVarObject_HEAD_INIT(NULL, 0)
922 "itertools.dropwhile", /* tp_name */
923 sizeof(dropwhileobject), /* tp_basicsize */
924 0, /* tp_itemsize */
925 /* methods */
926 (destructor)dropwhile_dealloc, /* tp_dealloc */
927 0, /* tp_print */
928 0, /* tp_getattr */
929 0, /* tp_setattr */
930 0, /* tp_reserved */
931 0, /* tp_repr */
932 0, /* tp_as_number */
933 0, /* tp_as_sequence */
934 0, /* tp_as_mapping */
935 0, /* tp_hash */
936 0, /* tp_call */
937 0, /* tp_str */
938 PyObject_GenericGetAttr, /* tp_getattro */
939 0, /* tp_setattro */
940 0, /* tp_as_buffer */
941 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
942 Py_TPFLAGS_BASETYPE, /* tp_flags */
943 dropwhile_doc, /* tp_doc */
944 (traverseproc)dropwhile_traverse, /* tp_traverse */
945 0, /* tp_clear */
946 0, /* tp_richcompare */
947 0, /* tp_weaklistoffset */
948 PyObject_SelfIter, /* tp_iter */
949 (iternextfunc)dropwhile_next, /* tp_iternext */
950 0, /* tp_methods */
951 0, /* tp_members */
952 0, /* tp_getset */
953 0, /* tp_base */
954 0, /* tp_dict */
955 0, /* tp_descr_get */
956 0, /* tp_descr_set */
957 0, /* tp_dictoffset */
958 0, /* tp_init */
959 0, /* tp_alloc */
960 dropwhile_new, /* tp_new */
961 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000962};
963
964
965/* takewhile object **********************************************************/
966
967typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000968 PyObject_HEAD
969 PyObject *func;
970 PyObject *it;
971 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000972} takewhileobject;
973
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000974static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000975
976static PyObject *
977takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
978{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000979 PyObject *func, *seq;
980 PyObject *it;
981 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000982
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000983 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
984 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000985
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000986 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
987 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000988
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000989 /* Get iterator. */
990 it = PyObject_GetIter(seq);
991 if (it == NULL)
992 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000993
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000994 /* create takewhileobject structure */
995 lz = (takewhileobject *)type->tp_alloc(type, 0);
996 if (lz == NULL) {
997 Py_DECREF(it);
998 return NULL;
999 }
1000 Py_INCREF(func);
1001 lz->func = func;
1002 lz->it = it;
1003 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001005 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001006}
1007
1008static void
1009takewhile_dealloc(takewhileobject *lz)
1010{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001011 PyObject_GC_UnTrack(lz);
1012 Py_XDECREF(lz->func);
1013 Py_XDECREF(lz->it);
1014 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001015}
1016
1017static int
1018takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1019{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001020 Py_VISIT(lz->it);
1021 Py_VISIT(lz->func);
1022 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001023}
1024
1025static PyObject *
1026takewhile_next(takewhileobject *lz)
1027{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001028 PyObject *item, *good;
1029 PyObject *it = lz->it;
1030 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001031
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001032 if (lz->stop == 1)
1033 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001034
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001035 item = (*Py_TYPE(it)->tp_iternext)(it);
1036 if (item == NULL)
1037 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001038
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001039 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1040 if (good == NULL) {
1041 Py_DECREF(item);
1042 return NULL;
1043 }
1044 ok = PyObject_IsTrue(good);
1045 Py_DECREF(good);
1046 if (ok)
1047 return item;
1048 Py_DECREF(item);
1049 lz->stop = 1;
1050 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001051}
1052
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053PyDoc_STRVAR(takewhile_doc,
1054"takewhile(predicate, iterable) --> takewhile object\n\
1055\n\
1056Return successive entries from an iterable as long as the \n\
1057predicate evaluates to true for each entry.");
1058
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001059static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001060 PyVarObject_HEAD_INIT(NULL, 0)
1061 "itertools.takewhile", /* tp_name */
1062 sizeof(takewhileobject), /* tp_basicsize */
1063 0, /* tp_itemsize */
1064 /* methods */
1065 (destructor)takewhile_dealloc, /* tp_dealloc */
1066 0, /* tp_print */
1067 0, /* tp_getattr */
1068 0, /* tp_setattr */
1069 0, /* tp_reserved */
1070 0, /* tp_repr */
1071 0, /* tp_as_number */
1072 0, /* tp_as_sequence */
1073 0, /* tp_as_mapping */
1074 0, /* tp_hash */
1075 0, /* tp_call */
1076 0, /* tp_str */
1077 PyObject_GenericGetAttr, /* tp_getattro */
1078 0, /* tp_setattro */
1079 0, /* tp_as_buffer */
1080 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1081 Py_TPFLAGS_BASETYPE, /* tp_flags */
1082 takewhile_doc, /* tp_doc */
1083 (traverseproc)takewhile_traverse, /* tp_traverse */
1084 0, /* tp_clear */
1085 0, /* tp_richcompare */
1086 0, /* tp_weaklistoffset */
1087 PyObject_SelfIter, /* tp_iter */
1088 (iternextfunc)takewhile_next, /* tp_iternext */
1089 0, /* tp_methods */
1090 0, /* tp_members */
1091 0, /* tp_getset */
1092 0, /* tp_base */
1093 0, /* tp_dict */
1094 0, /* tp_descr_get */
1095 0, /* tp_descr_set */
1096 0, /* tp_dictoffset */
1097 0, /* tp_init */
1098 0, /* tp_alloc */
1099 takewhile_new, /* tp_new */
1100 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001101};
1102
1103
1104/* islice object ************************************************************/
1105
1106typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001107 PyObject_HEAD
1108 PyObject *it;
1109 Py_ssize_t next;
1110 Py_ssize_t stop;
1111 Py_ssize_t step;
1112 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001113} isliceobject;
1114
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001115static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116
1117static PyObject *
1118islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1119{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001120 PyObject *seq;
1121 Py_ssize_t start=0, stop=-1, step=1;
1122 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1123 Py_ssize_t numargs;
1124 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001125
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001126 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1127 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001128
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001129 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1130 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001131
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001132 numargs = PyTuple_Size(args);
1133 if (numargs == 2) {
1134 if (a1 != Py_None) {
1135 stop = PyLong_AsSsize_t(a1);
1136 if (stop == -1) {
1137 if (PyErr_Occurred())
1138 PyErr_Clear();
1139 PyErr_SetString(PyExc_ValueError,
1140 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1141 return NULL;
1142 }
1143 }
1144 } else {
1145 if (a1 != Py_None)
1146 start = PyLong_AsSsize_t(a1);
1147 if (start == -1 && PyErr_Occurred())
1148 PyErr_Clear();
1149 if (a2 != Py_None) {
1150 stop = PyLong_AsSsize_t(a2);
1151 if (stop == -1) {
1152 if (PyErr_Occurred())
1153 PyErr_Clear();
1154 PyErr_SetString(PyExc_ValueError,
1155 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1156 return NULL;
1157 }
1158 }
1159 }
1160 if (start<0 || stop<-1) {
1161 PyErr_SetString(PyExc_ValueError,
1162 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1163 return NULL;
1164 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001165
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001166 if (a3 != NULL) {
1167 if (a3 != Py_None)
1168 step = PyLong_AsSsize_t(a3);
1169 if (step == -1 && PyErr_Occurred())
1170 PyErr_Clear();
1171 }
1172 if (step<1) {
1173 PyErr_SetString(PyExc_ValueError,
1174 "Step for islice() must be a positive integer or None.");
1175 return NULL;
1176 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001177
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001178 /* Get iterator. */
1179 it = PyObject_GetIter(seq);
1180 if (it == NULL)
1181 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001182
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001183 /* create isliceobject structure */
1184 lz = (isliceobject *)type->tp_alloc(type, 0);
1185 if (lz == NULL) {
1186 Py_DECREF(it);
1187 return NULL;
1188 }
1189 lz->it = it;
1190 lz->next = start;
1191 lz->stop = stop;
1192 lz->step = step;
1193 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001194
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001195 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196}
1197
1198static void
1199islice_dealloc(isliceobject *lz)
1200{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001201 PyObject_GC_UnTrack(lz);
1202 Py_XDECREF(lz->it);
1203 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204}
1205
1206static int
1207islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1208{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001209 Py_VISIT(lz->it);
1210 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001211}
1212
1213static PyObject *
1214islice_next(isliceobject *lz)
1215{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001216 PyObject *item;
1217 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001218 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001219 Py_ssize_t oldnext;
1220 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001221
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001222 iternext = *Py_TYPE(it)->tp_iternext;
1223 while (lz->cnt < lz->next) {
1224 item = iternext(it);
1225 if (item == NULL)
1226 return NULL;
1227 Py_DECREF(item);
1228 lz->cnt++;
1229 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001230 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001231 return NULL;
1232 item = iternext(it);
1233 if (item == NULL)
1234 return NULL;
1235 lz->cnt++;
1236 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001237 /* The (size_t) cast below avoids the danger of undefined
1238 behaviour from signed integer overflow. */
1239 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001240 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1241 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001242 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001243}
1244
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001245PyDoc_STRVAR(islice_doc,
1246"islice(iterable, [start,] stop [, step]) --> islice object\n\
1247\n\
1248Return an iterator whose next() method returns selected values from an\n\
1249iterable. If start is specified, will skip all preceding elements;\n\
1250otherwise, start defaults to zero. Step defaults to one. If\n\
1251specified as another value, step determines how many values are \n\
1252skipped between successive calls. Works like a slice() on a list\n\
1253but returns an iterator.");
1254
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001255static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001256 PyVarObject_HEAD_INIT(NULL, 0)
1257 "itertools.islice", /* tp_name */
1258 sizeof(isliceobject), /* tp_basicsize */
1259 0, /* tp_itemsize */
1260 /* methods */
1261 (destructor)islice_dealloc, /* tp_dealloc */
1262 0, /* tp_print */
1263 0, /* tp_getattr */
1264 0, /* tp_setattr */
1265 0, /* tp_reserved */
1266 0, /* tp_repr */
1267 0, /* tp_as_number */
1268 0, /* tp_as_sequence */
1269 0, /* tp_as_mapping */
1270 0, /* tp_hash */
1271 0, /* tp_call */
1272 0, /* tp_str */
1273 PyObject_GenericGetAttr, /* tp_getattro */
1274 0, /* tp_setattro */
1275 0, /* tp_as_buffer */
1276 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1277 Py_TPFLAGS_BASETYPE, /* tp_flags */
1278 islice_doc, /* tp_doc */
1279 (traverseproc)islice_traverse, /* tp_traverse */
1280 0, /* tp_clear */
1281 0, /* tp_richcompare */
1282 0, /* tp_weaklistoffset */
1283 PyObject_SelfIter, /* tp_iter */
1284 (iternextfunc)islice_next, /* tp_iternext */
1285 0, /* tp_methods */
1286 0, /* tp_members */
1287 0, /* tp_getset */
1288 0, /* tp_base */
1289 0, /* tp_dict */
1290 0, /* tp_descr_get */
1291 0, /* tp_descr_set */
1292 0, /* tp_dictoffset */
1293 0, /* tp_init */
1294 0, /* tp_alloc */
1295 islice_new, /* tp_new */
1296 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001297};
1298
1299
1300/* starmap object ************************************************************/
1301
1302typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001303 PyObject_HEAD
1304 PyObject *func;
1305 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001306} starmapobject;
1307
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001308static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001309
1310static PyObject *
1311starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1312{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001313 PyObject *func, *seq;
1314 PyObject *it;
1315 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001316
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001317 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1318 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001319
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001320 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1321 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001322
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001323 /* Get iterator. */
1324 it = PyObject_GetIter(seq);
1325 if (it == NULL)
1326 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001327
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001328 /* create starmapobject structure */
1329 lz = (starmapobject *)type->tp_alloc(type, 0);
1330 if (lz == NULL) {
1331 Py_DECREF(it);
1332 return NULL;
1333 }
1334 Py_INCREF(func);
1335 lz->func = func;
1336 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001337
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001338 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001339}
1340
1341static void
1342starmap_dealloc(starmapobject *lz)
1343{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001344 PyObject_GC_UnTrack(lz);
1345 Py_XDECREF(lz->func);
1346 Py_XDECREF(lz->it);
1347 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001348}
1349
1350static int
1351starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1352{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001353 Py_VISIT(lz->it);
1354 Py_VISIT(lz->func);
1355 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001356}
1357
1358static PyObject *
1359starmap_next(starmapobject *lz)
1360{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001361 PyObject *args;
1362 PyObject *result;
1363 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001364
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001365 args = (*Py_TYPE(it)->tp_iternext)(it);
1366 if (args == NULL)
1367 return NULL;
1368 if (!PyTuple_CheckExact(args)) {
1369 PyObject *newargs = PySequence_Tuple(args);
1370 Py_DECREF(args);
1371 if (newargs == NULL)
1372 return NULL;
1373 args = newargs;
1374 }
1375 result = PyObject_Call(lz->func, args, NULL);
1376 Py_DECREF(args);
1377 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001378}
1379
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001380PyDoc_STRVAR(starmap_doc,
1381"starmap(function, sequence) --> starmap object\n\
1382\n\
1383Return an iterator whose values are returned from the function evaluated\n\
1384with a argument tuple taken from the given sequence.");
1385
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001386static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001387 PyVarObject_HEAD_INIT(NULL, 0)
1388 "itertools.starmap", /* tp_name */
1389 sizeof(starmapobject), /* tp_basicsize */
1390 0, /* tp_itemsize */
1391 /* methods */
1392 (destructor)starmap_dealloc, /* tp_dealloc */
1393 0, /* tp_print */
1394 0, /* tp_getattr */
1395 0, /* tp_setattr */
1396 0, /* tp_reserved */
1397 0, /* tp_repr */
1398 0, /* tp_as_number */
1399 0, /* tp_as_sequence */
1400 0, /* tp_as_mapping */
1401 0, /* tp_hash */
1402 0, /* tp_call */
1403 0, /* tp_str */
1404 PyObject_GenericGetAttr, /* tp_getattro */
1405 0, /* tp_setattro */
1406 0, /* tp_as_buffer */
1407 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1408 Py_TPFLAGS_BASETYPE, /* tp_flags */
1409 starmap_doc, /* tp_doc */
1410 (traverseproc)starmap_traverse, /* tp_traverse */
1411 0, /* tp_clear */
1412 0, /* tp_richcompare */
1413 0, /* tp_weaklistoffset */
1414 PyObject_SelfIter, /* tp_iter */
1415 (iternextfunc)starmap_next, /* tp_iternext */
1416 0, /* tp_methods */
1417 0, /* tp_members */
1418 0, /* tp_getset */
1419 0, /* tp_base */
1420 0, /* tp_dict */
1421 0, /* tp_descr_get */
1422 0, /* tp_descr_set */
1423 0, /* tp_dictoffset */
1424 0, /* tp_init */
1425 0, /* tp_alloc */
1426 starmap_new, /* tp_new */
1427 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001428};
1429
1430
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001431/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001432
1433typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001434 PyObject_HEAD
1435 PyObject *source; /* Iterator over input iterables */
1436 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001437} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001438
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001439static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001440
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001441static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001442chain_new_internal(PyTypeObject *type, PyObject *source)
1443{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001444 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001445
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001446 lz = (chainobject *)type->tp_alloc(type, 0);
1447 if (lz == NULL) {
1448 Py_DECREF(source);
1449 return NULL;
1450 }
1451
1452 lz->source = source;
1453 lz->active = NULL;
1454 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001455}
1456
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001458chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001460 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001461
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001462 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1463 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001464
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001465 source = PyObject_GetIter(args);
1466 if (source == NULL)
1467 return NULL;
1468
1469 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001470}
1471
1472static PyObject *
1473chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1474{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001475 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001476
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001477 source = PyObject_GetIter(arg);
1478 if (source == NULL)
1479 return NULL;
1480
1481 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001482}
1483
1484static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001485chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001486{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001487 PyObject_GC_UnTrack(lz);
1488 Py_XDECREF(lz->active);
1489 Py_XDECREF(lz->source);
1490 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001491}
1492
Raymond Hettinger2012f172003-02-07 05:32:58 +00001493static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001494chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001495{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001496 Py_VISIT(lz->source);
1497 Py_VISIT(lz->active);
1498 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001499}
1500
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001501static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001502chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001503{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001504 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001505
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001506 if (lz->source == NULL)
1507 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001508
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001509 if (lz->active == NULL) {
1510 PyObject *iterable = PyIter_Next(lz->source);
1511 if (iterable == NULL) {
1512 Py_CLEAR(lz->source);
1513 return NULL; /* no more input sources */
1514 }
1515 lz->active = PyObject_GetIter(iterable);
1516 Py_DECREF(iterable);
1517 if (lz->active == NULL) {
1518 Py_CLEAR(lz->source);
1519 return NULL; /* input not iterable */
1520 }
1521 }
1522 item = PyIter_Next(lz->active);
1523 if (item != NULL)
1524 return item;
1525 if (PyErr_Occurred()) {
1526 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1527 PyErr_Clear();
1528 else
1529 return NULL; /* input raised an exception */
1530 }
1531 Py_CLEAR(lz->active);
1532 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001533}
1534
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001535PyDoc_STRVAR(chain_doc,
1536"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001537\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001538Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001539first iterable until it is exhausted, then elements from the next\n\
1540iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001541
Christian Heimesf16baeb2008-02-29 14:57:44 +00001542PyDoc_STRVAR(chain_from_iterable_doc,
1543"chain.from_iterable(iterable) --> chain object\n\
1544\n\
1545Alternate chain() contructor taking a single iterable argument\n\
1546that evaluates lazily.");
1547
1548static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001549 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1550 chain_from_iterable_doc},
1551 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001552};
1553
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001554static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001555 PyVarObject_HEAD_INIT(NULL, 0)
1556 "itertools.chain", /* tp_name */
1557 sizeof(chainobject), /* tp_basicsize */
1558 0, /* tp_itemsize */
1559 /* methods */
1560 (destructor)chain_dealloc, /* tp_dealloc */
1561 0, /* tp_print */
1562 0, /* tp_getattr */
1563 0, /* tp_setattr */
1564 0, /* tp_reserved */
1565 0, /* tp_repr */
1566 0, /* tp_as_number */
1567 0, /* tp_as_sequence */
1568 0, /* tp_as_mapping */
1569 0, /* tp_hash */
1570 0, /* tp_call */
1571 0, /* tp_str */
1572 PyObject_GenericGetAttr, /* tp_getattro */
1573 0, /* tp_setattro */
1574 0, /* tp_as_buffer */
1575 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1576 Py_TPFLAGS_BASETYPE, /* tp_flags */
1577 chain_doc, /* tp_doc */
1578 (traverseproc)chain_traverse, /* tp_traverse */
1579 0, /* tp_clear */
1580 0, /* tp_richcompare */
1581 0, /* tp_weaklistoffset */
1582 PyObject_SelfIter, /* tp_iter */
1583 (iternextfunc)chain_next, /* tp_iternext */
1584 chain_methods, /* tp_methods */
1585 0, /* tp_members */
1586 0, /* tp_getset */
1587 0, /* tp_base */
1588 0, /* tp_dict */
1589 0, /* tp_descr_get */
1590 0, /* tp_descr_set */
1591 0, /* tp_dictoffset */
1592 0, /* tp_init */
1593 0, /* tp_alloc */
1594 chain_new, /* tp_new */
1595 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001596};
1597
1598
Christian Heimesc3f30c42008-02-22 16:37:40 +00001599/* product object ************************************************************/
1600
1601typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001602 PyObject_HEAD
1603 PyObject *pools; /* tuple of pool tuples */
1604 Py_ssize_t *indices; /* one index per pool */
1605 PyObject *result; /* most recently returned result tuple */
1606 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001607} productobject;
1608
1609static PyTypeObject product_type;
1610
1611static PyObject *
1612product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1613{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001614 productobject *lz;
1615 Py_ssize_t nargs, npools, repeat=1;
1616 PyObject *pools = NULL;
1617 Py_ssize_t *indices = NULL;
1618 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001619
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001620 if (kwds != NULL) {
1621 char *kwlist[] = {"repeat", 0};
1622 PyObject *tmpargs = PyTuple_New(0);
1623 if (tmpargs == NULL)
1624 return NULL;
1625 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1626 Py_DECREF(tmpargs);
1627 return NULL;
1628 }
1629 Py_DECREF(tmpargs);
1630 if (repeat < 0) {
1631 PyErr_SetString(PyExc_ValueError,
1632 "repeat argument cannot be negative");
1633 return NULL;
1634 }
1635 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001636
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001637 assert(PyTuple_Check(args));
1638 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1639 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001640
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001641 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1642 if (indices == NULL) {
1643 PyErr_NoMemory();
1644 goto error;
1645 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001646
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001647 pools = PyTuple_New(npools);
1648 if (pools == NULL)
1649 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001650
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001651 for (i=0; i < nargs ; ++i) {
1652 PyObject *item = PyTuple_GET_ITEM(args, i);
1653 PyObject *pool = PySequence_Tuple(item);
1654 if (pool == NULL)
1655 goto error;
1656 PyTuple_SET_ITEM(pools, i, pool);
1657 indices[i] = 0;
1658 }
1659 for ( ; i < npools; ++i) {
1660 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1661 Py_INCREF(pool);
1662 PyTuple_SET_ITEM(pools, i, pool);
1663 indices[i] = 0;
1664 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001665
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001666 /* create productobject structure */
1667 lz = (productobject *)type->tp_alloc(type, 0);
1668 if (lz == NULL)
1669 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001670
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001671 lz->pools = pools;
1672 lz->indices = indices;
1673 lz->result = NULL;
1674 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001675
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001676 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001677
1678error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001679 if (indices != NULL)
1680 PyMem_Free(indices);
1681 Py_XDECREF(pools);
1682 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001683}
1684
1685static void
1686product_dealloc(productobject *lz)
1687{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001688 PyObject_GC_UnTrack(lz);
1689 Py_XDECREF(lz->pools);
1690 Py_XDECREF(lz->result);
1691 if (lz->indices != NULL)
1692 PyMem_Free(lz->indices);
1693 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001694}
1695
1696static int
1697product_traverse(productobject *lz, visitproc visit, void *arg)
1698{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001699 Py_VISIT(lz->pools);
1700 Py_VISIT(lz->result);
1701 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001702}
1703
1704static PyObject *
1705product_next(productobject *lz)
1706{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001707 PyObject *pool;
1708 PyObject *elem;
1709 PyObject *oldelem;
1710 PyObject *pools = lz->pools;
1711 PyObject *result = lz->result;
1712 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1713 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001714
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001715 if (lz->stopped)
1716 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001717
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001718 if (result == NULL) {
1719 /* On the first pass, return an initial tuple filled with the
1720 first element from each pool. */
1721 result = PyTuple_New(npools);
1722 if (result == NULL)
1723 goto empty;
1724 lz->result = result;
1725 for (i=0; i < npools; i++) {
1726 pool = PyTuple_GET_ITEM(pools, i);
1727 if (PyTuple_GET_SIZE(pool) == 0)
1728 goto empty;
1729 elem = PyTuple_GET_ITEM(pool, 0);
1730 Py_INCREF(elem);
1731 PyTuple_SET_ITEM(result, i, elem);
1732 }
1733 } else {
1734 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001735
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001736 /* Copy the previous result tuple or re-use it if available */
1737 if (Py_REFCNT(result) > 1) {
1738 PyObject *old_result = result;
1739 result = PyTuple_New(npools);
1740 if (result == NULL)
1741 goto empty;
1742 lz->result = result;
1743 for (i=0; i < npools; i++) {
1744 elem = PyTuple_GET_ITEM(old_result, i);
1745 Py_INCREF(elem);
1746 PyTuple_SET_ITEM(result, i, elem);
1747 }
1748 Py_DECREF(old_result);
1749 }
1750 /* Now, we've got the only copy so we can update it in-place */
1751 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001752
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001753 /* Update the pool indices right-to-left. Only advance to the
1754 next pool when the previous one rolls-over */
1755 for (i=npools-1 ; i >= 0 ; i--) {
1756 pool = PyTuple_GET_ITEM(pools, i);
1757 indices[i]++;
1758 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1759 /* Roll-over and advance to next pool */
1760 indices[i] = 0;
1761 elem = PyTuple_GET_ITEM(pool, 0);
1762 Py_INCREF(elem);
1763 oldelem = PyTuple_GET_ITEM(result, i);
1764 PyTuple_SET_ITEM(result, i, elem);
1765 Py_DECREF(oldelem);
1766 } else {
1767 /* No rollover. Just increment and stop here. */
1768 elem = PyTuple_GET_ITEM(pool, indices[i]);
1769 Py_INCREF(elem);
1770 oldelem = PyTuple_GET_ITEM(result, i);
1771 PyTuple_SET_ITEM(result, i, elem);
1772 Py_DECREF(oldelem);
1773 break;
1774 }
1775 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001776
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001777 /* If i is negative, then the indices have all rolled-over
1778 and we're done. */
1779 if (i < 0)
1780 goto empty;
1781 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001782
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001783 Py_INCREF(result);
1784 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001785
1786empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001787 lz->stopped = 1;
1788 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001789}
1790
1791PyDoc_STRVAR(product_doc,
1792"product(*iterables) --> product object\n\
1793\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001794Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001795For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1796The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1797cycle in a manner similar to an odometer (with the rightmost element changing\n\
1798on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00001799To compute the product of an iterable with itself, specify the number\n\
1800of repetitions with the optional repeat keyword argument. For example,\n\
1801product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001802product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1803product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1804
1805static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001806 PyVarObject_HEAD_INIT(NULL, 0)
1807 "itertools.product", /* tp_name */
1808 sizeof(productobject), /* tp_basicsize */
1809 0, /* tp_itemsize */
1810 /* methods */
1811 (destructor)product_dealloc, /* tp_dealloc */
1812 0, /* tp_print */
1813 0, /* tp_getattr */
1814 0, /* tp_setattr */
1815 0, /* tp_reserved */
1816 0, /* tp_repr */
1817 0, /* tp_as_number */
1818 0, /* tp_as_sequence */
1819 0, /* tp_as_mapping */
1820 0, /* tp_hash */
1821 0, /* tp_call */
1822 0, /* tp_str */
1823 PyObject_GenericGetAttr, /* tp_getattro */
1824 0, /* tp_setattro */
1825 0, /* tp_as_buffer */
1826 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1827 Py_TPFLAGS_BASETYPE, /* tp_flags */
1828 product_doc, /* tp_doc */
1829 (traverseproc)product_traverse, /* tp_traverse */
1830 0, /* tp_clear */
1831 0, /* tp_richcompare */
1832 0, /* tp_weaklistoffset */
1833 PyObject_SelfIter, /* tp_iter */
1834 (iternextfunc)product_next, /* tp_iternext */
1835 0, /* tp_methods */
1836 0, /* tp_members */
1837 0, /* tp_getset */
1838 0, /* tp_base */
1839 0, /* tp_dict */
1840 0, /* tp_descr_get */
1841 0, /* tp_descr_set */
1842 0, /* tp_dictoffset */
1843 0, /* tp_init */
1844 0, /* tp_alloc */
1845 product_new, /* tp_new */
1846 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001847};
1848
1849
Christian Heimes380f7f22008-02-28 11:19:05 +00001850/* combinations object ************************************************************/
1851
1852typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001853 PyObject_HEAD
1854 PyObject *pool; /* input converted to a tuple */
1855 Py_ssize_t *indices; /* one index per result element */
1856 PyObject *result; /* most recently returned result tuple */
1857 Py_ssize_t r; /* size of result tuple */
1858 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00001859} combinationsobject;
1860
1861static PyTypeObject combinations_type;
1862
1863static PyObject *
1864combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1865{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001866 combinationsobject *co;
1867 Py_ssize_t n;
1868 Py_ssize_t r;
1869 PyObject *pool = NULL;
1870 PyObject *iterable = NULL;
1871 Py_ssize_t *indices = NULL;
1872 Py_ssize_t i;
1873 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00001874
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001875 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
1876 &iterable, &r))
1877 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001878
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001879 pool = PySequence_Tuple(iterable);
1880 if (pool == NULL)
1881 goto error;
1882 n = PyTuple_GET_SIZE(pool);
1883 if (r < 0) {
1884 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
1885 goto error;
1886 }
Christian Heimes380f7f22008-02-28 11:19:05 +00001887
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001888 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
1889 if (indices == NULL) {
1890 PyErr_NoMemory();
1891 goto error;
1892 }
Christian Heimes380f7f22008-02-28 11:19:05 +00001893
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001894 for (i=0 ; i<r ; i++)
1895 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00001896
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001897 /* create combinationsobject structure */
1898 co = (combinationsobject *)type->tp_alloc(type, 0);
1899 if (co == NULL)
1900 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00001901
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001902 co->pool = pool;
1903 co->indices = indices;
1904 co->result = NULL;
1905 co->r = r;
1906 co->stopped = r > n ? 1 : 0;
1907
1908 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00001909
1910error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001911 if (indices != NULL)
1912 PyMem_Free(indices);
1913 Py_XDECREF(pool);
1914 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001915}
1916
1917static void
1918combinations_dealloc(combinationsobject *co)
1919{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001920 PyObject_GC_UnTrack(co);
1921 Py_XDECREF(co->pool);
1922 Py_XDECREF(co->result);
1923 if (co->indices != NULL)
1924 PyMem_Free(co->indices);
1925 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00001926}
1927
1928static int
1929combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
1930{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001931 Py_VISIT(co->pool);
1932 Py_VISIT(co->result);
1933 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00001934}
1935
1936static PyObject *
1937combinations_next(combinationsobject *co)
1938{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001939 PyObject *elem;
1940 PyObject *oldelem;
1941 PyObject *pool = co->pool;
1942 Py_ssize_t *indices = co->indices;
1943 PyObject *result = co->result;
1944 Py_ssize_t n = PyTuple_GET_SIZE(pool);
1945 Py_ssize_t r = co->r;
1946 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00001947
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001948 if (co->stopped)
1949 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001950
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001951 if (result == NULL) {
1952 /* On the first pass, initialize result tuple using the indices */
1953 result = PyTuple_New(r);
1954 if (result == NULL)
1955 goto empty;
1956 co->result = result;
1957 for (i=0; i<r ; i++) {
1958 index = indices[i];
1959 elem = PyTuple_GET_ITEM(pool, index);
1960 Py_INCREF(elem);
1961 PyTuple_SET_ITEM(result, i, elem);
1962 }
1963 } else {
1964 /* Copy the previous result tuple or re-use it if available */
1965 if (Py_REFCNT(result) > 1) {
1966 PyObject *old_result = result;
1967 result = PyTuple_New(r);
1968 if (result == NULL)
1969 goto empty;
1970 co->result = result;
1971 for (i=0; i<r ; i++) {
1972 elem = PyTuple_GET_ITEM(old_result, i);
1973 Py_INCREF(elem);
1974 PyTuple_SET_ITEM(result, i, elem);
1975 }
1976 Py_DECREF(old_result);
1977 }
1978 /* Now, we've got the only copy so we can update it in-place
1979 * CPython's empty tuple is a singleton and cached in
1980 * PyTuple's freelist.
1981 */
1982 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00001983
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001984 /* Scan indices right-to-left until finding one that is not
1985 at its maximum (i + n - r). */
1986 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
1987 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00001988
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001989 /* If i is negative, then the indices are all at
1990 their maximum value and we're done. */
1991 if (i < 0)
1992 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00001993
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001994 /* Increment the current index which we know is not at its
1995 maximum. Then move back to the right setting each index
1996 to its lowest possible value (one higher than the index
1997 to its left -- this maintains the sort order invariant). */
1998 indices[i]++;
1999 for (j=i+1 ; j<r ; j++)
2000 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002001
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002002 /* Update the result tuple for the new indices
2003 starting with i, the leftmost index that changed */
2004 for ( ; i<r ; i++) {
2005 index = indices[i];
2006 elem = PyTuple_GET_ITEM(pool, index);
2007 Py_INCREF(elem);
2008 oldelem = PyTuple_GET_ITEM(result, i);
2009 PyTuple_SET_ITEM(result, i, elem);
2010 Py_DECREF(oldelem);
2011 }
2012 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002013
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002014 Py_INCREF(result);
2015 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002016
2017empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002018 co->stopped = 1;
2019 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002020}
2021
2022PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002023"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002024\n\
2025Return successive r-length combinations of elements in the iterable.\n\n\
2026combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2027
2028static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002029 PyVarObject_HEAD_INIT(NULL, 0)
2030 "itertools.combinations", /* tp_name */
2031 sizeof(combinationsobject), /* tp_basicsize */
2032 0, /* tp_itemsize */
2033 /* methods */
2034 (destructor)combinations_dealloc, /* tp_dealloc */
2035 0, /* tp_print */
2036 0, /* tp_getattr */
2037 0, /* tp_setattr */
2038 0, /* tp_reserved */
2039 0, /* tp_repr */
2040 0, /* tp_as_number */
2041 0, /* tp_as_sequence */
2042 0, /* tp_as_mapping */
2043 0, /* tp_hash */
2044 0, /* tp_call */
2045 0, /* tp_str */
2046 PyObject_GenericGetAttr, /* tp_getattro */
2047 0, /* tp_setattro */
2048 0, /* tp_as_buffer */
2049 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2050 Py_TPFLAGS_BASETYPE, /* tp_flags */
2051 combinations_doc, /* tp_doc */
2052 (traverseproc)combinations_traverse, /* tp_traverse */
2053 0, /* tp_clear */
2054 0, /* tp_richcompare */
2055 0, /* tp_weaklistoffset */
2056 PyObject_SelfIter, /* tp_iter */
2057 (iternextfunc)combinations_next, /* tp_iternext */
2058 0, /* tp_methods */
2059 0, /* tp_members */
2060 0, /* tp_getset */
2061 0, /* tp_base */
2062 0, /* tp_dict */
2063 0, /* tp_descr_get */
2064 0, /* tp_descr_set */
2065 0, /* tp_dictoffset */
2066 0, /* tp_init */
2067 0, /* tp_alloc */
2068 combinations_new, /* tp_new */
2069 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002070};
2071
2072
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002073/* combinations with replacement object *******************************************/
2074
2075/* Equivalent to:
2076
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002077 def combinations_with_replacement(iterable, r):
2078 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2079 # number items returned: (n+r-1)! / r! / (n-1)!
2080 pool = tuple(iterable)
2081 n = len(pool)
2082 indices = [0] * r
2083 yield tuple(pool[i] for i in indices)
2084 while 1:
2085 for i in reversed(range(r)):
2086 if indices[i] != n - 1:
2087 break
2088 else:
2089 return
2090 indices[i:] = [indices[i] + 1] * (r - i)
2091 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002092
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002093 def combinations_with_replacement2(iterable, r):
2094 'Alternate version that filters from product()'
2095 pool = tuple(iterable)
2096 n = len(pool)
2097 for indices in product(range(n), repeat=r):
2098 if sorted(indices) == list(indices):
2099 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002100*/
2101typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002102 PyObject_HEAD
2103 PyObject *pool; /* input converted to a tuple */
2104 Py_ssize_t *indices; /* one index per result element */
2105 PyObject *result; /* most recently returned result tuple */
2106 Py_ssize_t r; /* size of result tuple */
2107 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002108} cwrobject;
2109
2110static PyTypeObject cwr_type;
2111
2112static PyObject *
2113cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2114{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002115 cwrobject *co;
2116 Py_ssize_t n;
2117 Py_ssize_t r;
2118 PyObject *pool = NULL;
2119 PyObject *iterable = NULL;
2120 Py_ssize_t *indices = NULL;
2121 Py_ssize_t i;
2122 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002123
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002124 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2125 &iterable, &r))
2126 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002127
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002128 pool = PySequence_Tuple(iterable);
2129 if (pool == NULL)
2130 goto error;
2131 n = PyTuple_GET_SIZE(pool);
2132 if (r < 0) {
2133 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2134 goto error;
2135 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002136
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002137 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2138 if (indices == NULL) {
2139 PyErr_NoMemory();
2140 goto error;
2141 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002142
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002143 for (i=0 ; i<r ; i++)
2144 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002145
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002146 /* create cwrobject structure */
2147 co = (cwrobject *)type->tp_alloc(type, 0);
2148 if (co == NULL)
2149 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002150
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002151 co->pool = pool;
2152 co->indices = indices;
2153 co->result = NULL;
2154 co->r = r;
2155 co->stopped = !n && r;
2156
2157 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002158
2159error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002160 if (indices != NULL)
2161 PyMem_Free(indices);
2162 Py_XDECREF(pool);
2163 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002164}
2165
2166static void
2167cwr_dealloc(cwrobject *co)
2168{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002169 PyObject_GC_UnTrack(co);
2170 Py_XDECREF(co->pool);
2171 Py_XDECREF(co->result);
2172 if (co->indices != NULL)
2173 PyMem_Free(co->indices);
2174 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002175}
2176
2177static int
2178cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2179{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002180 Py_VISIT(co->pool);
2181 Py_VISIT(co->result);
2182 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002183}
2184
2185static PyObject *
2186cwr_next(cwrobject *co)
2187{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002188 PyObject *elem;
2189 PyObject *oldelem;
2190 PyObject *pool = co->pool;
2191 Py_ssize_t *indices = co->indices;
2192 PyObject *result = co->result;
2193 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2194 Py_ssize_t r = co->r;
2195 Py_ssize_t i, j, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002196
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002197 if (co->stopped)
2198 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002199
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002200 if (result == NULL) {
2201 /* On the first pass, initialize result tuple using the indices */
2202 result = PyTuple_New(r);
2203 if (result == NULL)
2204 goto empty;
2205 co->result = result;
2206 for (i=0; i<r ; i++) {
2207 index = indices[i];
2208 elem = PyTuple_GET_ITEM(pool, index);
2209 Py_INCREF(elem);
2210 PyTuple_SET_ITEM(result, i, elem);
2211 }
2212 } else {
2213 /* Copy the previous result tuple or re-use it if available */
2214 if (Py_REFCNT(result) > 1) {
2215 PyObject *old_result = result;
2216 result = PyTuple_New(r);
2217 if (result == NULL)
2218 goto empty;
2219 co->result = result;
2220 for (i=0; i<r ; i++) {
2221 elem = PyTuple_GET_ITEM(old_result, i);
2222 Py_INCREF(elem);
2223 PyTuple_SET_ITEM(result, i, elem);
2224 }
2225 Py_DECREF(old_result);
2226 }
2227 /* Now, we've got the only copy so we can update it in-place CPython's
2228 empty tuple is a singleton and cached in PyTuple's freelist. */
2229 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002230
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002231 /* Scan indices right-to-left until finding one that is not
2232 * at its maximum (n-1). */
2233 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2234 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002235
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002236 /* If i is negative, then the indices are all at
2237 their maximum value and we're done. */
2238 if (i < 0)
2239 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002240
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002241 /* Increment the current index which we know is not at its
2242 maximum. Then set all to the right to the same value. */
2243 indices[i]++;
2244 for (j=i+1 ; j<r ; j++)
2245 indices[j] = indices[j-1];
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002246
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002247 /* Update the result tuple for the new indices
2248 starting with i, the leftmost index that changed */
2249 for ( ; i<r ; i++) {
2250 index = indices[i];
2251 elem = PyTuple_GET_ITEM(pool, index);
2252 Py_INCREF(elem);
2253 oldelem = PyTuple_GET_ITEM(result, i);
2254 PyTuple_SET_ITEM(result, i, elem);
2255 Py_DECREF(oldelem);
2256 }
2257 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002258
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002259 Py_INCREF(result);
2260 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002261
2262empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002263 co->stopped = 1;
2264 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002265}
2266
2267PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002268"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002269\n\
2270Return successive r-length combinations of elements in the iterable\n\
2271allowing individual elements to have successive repeats.\n\
2272combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2273
2274static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002275 PyVarObject_HEAD_INIT(NULL, 0)
2276 "itertools.combinations_with_replacement", /* tp_name */
2277 sizeof(cwrobject), /* tp_basicsize */
2278 0, /* tp_itemsize */
2279 /* methods */
2280 (destructor)cwr_dealloc, /* tp_dealloc */
2281 0, /* tp_print */
2282 0, /* tp_getattr */
2283 0, /* tp_setattr */
2284 0, /* tp_reserved */
2285 0, /* tp_repr */
2286 0, /* tp_as_number */
2287 0, /* tp_as_sequence */
2288 0, /* tp_as_mapping */
2289 0, /* tp_hash */
2290 0, /* tp_call */
2291 0, /* tp_str */
2292 PyObject_GenericGetAttr, /* tp_getattro */
2293 0, /* tp_setattro */
2294 0, /* tp_as_buffer */
2295 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2296 Py_TPFLAGS_BASETYPE, /* tp_flags */
2297 cwr_doc, /* tp_doc */
2298 (traverseproc)cwr_traverse, /* tp_traverse */
2299 0, /* tp_clear */
2300 0, /* tp_richcompare */
2301 0, /* tp_weaklistoffset */
2302 PyObject_SelfIter, /* tp_iter */
2303 (iternextfunc)cwr_next, /* tp_iternext */
2304 0, /* tp_methods */
2305 0, /* tp_members */
2306 0, /* tp_getset */
2307 0, /* tp_base */
2308 0, /* tp_dict */
2309 0, /* tp_descr_get */
2310 0, /* tp_descr_set */
2311 0, /* tp_dictoffset */
2312 0, /* tp_init */
2313 0, /* tp_alloc */
2314 cwr_new, /* tp_new */
2315 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002316};
2317
2318
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002319/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002320
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002321def permutations(iterable, r=None):
2322 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2323 pool = tuple(iterable)
2324 n = len(pool)
2325 r = n if r is None else r
2326 indices = range(n)
2327 cycles = range(n-r+1, n+1)[::-1]
2328 yield tuple(pool[i] for i in indices[:r])
2329 while n:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002330 for i in reversed(range(r)):
2331 cycles[i] -= 1
2332 if cycles[i] == 0:
2333 indices[i:] = indices[i+1:] + indices[i:i+1]
2334 cycles[i] = n - i
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002335 else:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002336 j = cycles[i]
2337 indices[i], indices[-j] = indices[-j], indices[i]
2338 yield tuple(pool[i] for i in indices[:r])
2339 break
2340 else:
2341 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002342*/
2343
2344typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002345 PyObject_HEAD
2346 PyObject *pool; /* input converted to a tuple */
2347 Py_ssize_t *indices; /* one index per element in the pool */
2348 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2349 PyObject *result; /* most recently returned result tuple */
2350 Py_ssize_t r; /* size of result tuple */
2351 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002352} permutationsobject;
2353
2354static PyTypeObject permutations_type;
2355
2356static PyObject *
2357permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2358{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002359 permutationsobject *po;
2360 Py_ssize_t n;
2361 Py_ssize_t r;
2362 PyObject *robj = Py_None;
2363 PyObject *pool = NULL;
2364 PyObject *iterable = NULL;
2365 Py_ssize_t *indices = NULL;
2366 Py_ssize_t *cycles = NULL;
2367 Py_ssize_t i;
2368 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002369
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002370 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2371 &iterable, &robj))
2372 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002373
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002374 pool = PySequence_Tuple(iterable);
2375 if (pool == NULL)
2376 goto error;
2377 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002378
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002379 r = n;
2380 if (robj != Py_None) {
2381 if (!PyLong_Check(robj)) {
2382 PyErr_SetString(PyExc_TypeError, "Expected int as r");
2383 goto error;
2384 }
2385 r = PyLong_AsSsize_t(robj);
2386 if (r == -1 && PyErr_Occurred())
2387 goto error;
2388 }
2389 if (r < 0) {
2390 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2391 goto error;
2392 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002393
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002394 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2395 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2396 if (indices == NULL || cycles == NULL) {
2397 PyErr_NoMemory();
2398 goto error;
2399 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002400
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002401 for (i=0 ; i<n ; i++)
2402 indices[i] = i;
2403 for (i=0 ; i<r ; i++)
2404 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002405
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002406 /* create permutationsobject structure */
2407 po = (permutationsobject *)type->tp_alloc(type, 0);
2408 if (po == NULL)
2409 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002410
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002411 po->pool = pool;
2412 po->indices = indices;
2413 po->cycles = cycles;
2414 po->result = NULL;
2415 po->r = r;
2416 po->stopped = r > n ? 1 : 0;
2417
2418 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002419
2420error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002421 if (indices != NULL)
2422 PyMem_Free(indices);
2423 if (cycles != NULL)
2424 PyMem_Free(cycles);
2425 Py_XDECREF(pool);
2426 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002427}
2428
2429static void
2430permutations_dealloc(permutationsobject *po)
2431{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002432 PyObject_GC_UnTrack(po);
2433 Py_XDECREF(po->pool);
2434 Py_XDECREF(po->result);
2435 PyMem_Free(po->indices);
2436 PyMem_Free(po->cycles);
2437 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002438}
2439
2440static int
2441permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2442{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002443 Py_VISIT(po->pool);
2444 Py_VISIT(po->result);
2445 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002446}
2447
2448static PyObject *
2449permutations_next(permutationsobject *po)
2450{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002451 PyObject *elem;
2452 PyObject *oldelem;
2453 PyObject *pool = po->pool;
2454 Py_ssize_t *indices = po->indices;
2455 Py_ssize_t *cycles = po->cycles;
2456 PyObject *result = po->result;
2457 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2458 Py_ssize_t r = po->r;
2459 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002460
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002461 if (po->stopped)
2462 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002463
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002464 if (result == NULL) {
2465 /* On the first pass, initialize result tuple using the indices */
2466 result = PyTuple_New(r);
2467 if (result == NULL)
2468 goto empty;
2469 po->result = result;
2470 for (i=0; i<r ; i++) {
2471 index = indices[i];
2472 elem = PyTuple_GET_ITEM(pool, index);
2473 Py_INCREF(elem);
2474 PyTuple_SET_ITEM(result, i, elem);
2475 }
2476 } else {
2477 if (n == 0)
2478 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002479
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002480 /* Copy the previous result tuple or re-use it if available */
2481 if (Py_REFCNT(result) > 1) {
2482 PyObject *old_result = result;
2483 result = PyTuple_New(r);
2484 if (result == NULL)
2485 goto empty;
2486 po->result = result;
2487 for (i=0; i<r ; i++) {
2488 elem = PyTuple_GET_ITEM(old_result, i);
2489 Py_INCREF(elem);
2490 PyTuple_SET_ITEM(result, i, elem);
2491 }
2492 Py_DECREF(old_result);
2493 }
2494 /* Now, we've got the only copy so we can update it in-place */
2495 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002496
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002497 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2498 for (i=r-1 ; i>=0 ; i--) {
2499 cycles[i] -= 1;
2500 if (cycles[i] == 0) {
2501 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2502 index = indices[i];
2503 for (j=i ; j<n-1 ; j++)
2504 indices[j] = indices[j+1];
2505 indices[n-1] = index;
2506 cycles[i] = n - i;
2507 } else {
2508 j = cycles[i];
2509 index = indices[i];
2510 indices[i] = indices[n-j];
2511 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002512
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002513 for (k=i; k<r ; k++) {
2514 /* start with i, the leftmost element that changed */
2515 /* yield tuple(pool[k] for k in indices[:r]) */
2516 index = indices[k];
2517 elem = PyTuple_GET_ITEM(pool, index);
2518 Py_INCREF(elem);
2519 oldelem = PyTuple_GET_ITEM(result, k);
2520 PyTuple_SET_ITEM(result, k, elem);
2521 Py_DECREF(oldelem);
2522 }
2523 break;
2524 }
2525 }
2526 /* If i is negative, then the cycles have all
2527 rolled-over and we're done. */
2528 if (i < 0)
2529 goto empty;
2530 }
2531 Py_INCREF(result);
2532 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002533
2534empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002535 po->stopped = 1;
2536 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002537}
2538
2539PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00002540"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002541\n\
2542Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00002543permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002544
2545static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002546 PyVarObject_HEAD_INIT(NULL, 0)
2547 "itertools.permutations", /* tp_name */
2548 sizeof(permutationsobject), /* tp_basicsize */
2549 0, /* tp_itemsize */
2550 /* methods */
2551 (destructor)permutations_dealloc, /* tp_dealloc */
2552 0, /* tp_print */
2553 0, /* tp_getattr */
2554 0, /* tp_setattr */
2555 0, /* tp_reserved */
2556 0, /* tp_repr */
2557 0, /* tp_as_number */
2558 0, /* tp_as_sequence */
2559 0, /* tp_as_mapping */
2560 0, /* tp_hash */
2561 0, /* tp_call */
2562 0, /* tp_str */
2563 PyObject_GenericGetAttr, /* tp_getattro */
2564 0, /* tp_setattro */
2565 0, /* tp_as_buffer */
2566 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2567 Py_TPFLAGS_BASETYPE, /* tp_flags */
2568 permutations_doc, /* tp_doc */
2569 (traverseproc)permutations_traverse, /* tp_traverse */
2570 0, /* tp_clear */
2571 0, /* tp_richcompare */
2572 0, /* tp_weaklistoffset */
2573 PyObject_SelfIter, /* tp_iter */
2574 (iternextfunc)permutations_next, /* tp_iternext */
2575 0, /* tp_methods */
2576 0, /* tp_members */
2577 0, /* tp_getset */
2578 0, /* tp_base */
2579 0, /* tp_dict */
2580 0, /* tp_descr_get */
2581 0, /* tp_descr_set */
2582 0, /* tp_dictoffset */
2583 0, /* tp_init */
2584 0, /* tp_alloc */
2585 permutations_new, /* tp_new */
2586 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002587};
2588
Raymond Hettinger482ba772010-12-01 22:48:00 +00002589/* accumulate object ************************************************************/
2590
2591typedef struct {
2592 PyObject_HEAD
2593 PyObject *total;
2594 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07002595 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002596} accumulateobject;
2597
2598static PyTypeObject accumulate_type;
2599
2600static PyObject *
2601accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2602{
Raymond Hettinger5d446132011-03-27 18:52:10 -07002603 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00002604 PyObject *iterable;
2605 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07002606 PyObject *binop = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002607 accumulateobject *lz;
2608
Raymond Hettinger5d446132011-03-27 18:52:10 -07002609 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
2610 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002611 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002612
2613 /* Get iterator. */
2614 it = PyObject_GetIter(iterable);
2615 if (it == NULL)
2616 return NULL;
2617
Raymond Hettinger482ba772010-12-01 22:48:00 +00002618 /* create accumulateobject structure */
2619 lz = (accumulateobject *)type->tp_alloc(type, 0);
2620 if (lz == NULL) {
2621 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002622 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002623 }
2624
Raymond Hettinger5d446132011-03-27 18:52:10 -07002625 Py_XINCREF(binop);
2626 lz->binop = binop;
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002627 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002628 lz->it = it;
2629 return (PyObject *)lz;
2630}
2631
2632static void
2633accumulate_dealloc(accumulateobject *lz)
2634{
2635 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07002636 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00002637 Py_XDECREF(lz->total);
2638 Py_XDECREF(lz->it);
2639 Py_TYPE(lz)->tp_free(lz);
2640}
2641
2642static int
2643accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
2644{
Raymond Hettinger5d446132011-03-27 18:52:10 -07002645 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00002646 Py_VISIT(lz->it);
2647 Py_VISIT(lz->total);
2648 return 0;
2649}
2650
2651static PyObject *
2652accumulate_next(accumulateobject *lz)
2653{
2654 PyObject *val, *oldtotal, *newtotal;
2655
2656 val = PyIter_Next(lz->it);
2657 if (val == NULL)
2658 return NULL;
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002659
2660 if (lz->total == NULL) {
2661 Py_INCREF(val);
2662 lz->total = val;
2663 return lz->total;
2664 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07002665
2666 if (lz->binop == NULL)
2667 newtotal = PyNumber_Add(lz->total, val);
2668 else
2669 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002670 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00002671 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002672 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002673
2674 oldtotal = lz->total;
2675 lz->total = newtotal;
2676 Py_DECREF(oldtotal);
2677
2678 Py_INCREF(newtotal);
2679 return newtotal;
2680}
2681
2682PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07002683"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00002684\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07002685Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00002686
2687static PyTypeObject accumulate_type = {
2688 PyVarObject_HEAD_INIT(NULL, 0)
2689 "itertools.accumulate", /* tp_name */
2690 sizeof(accumulateobject), /* tp_basicsize */
2691 0, /* tp_itemsize */
2692 /* methods */
2693 (destructor)accumulate_dealloc, /* tp_dealloc */
2694 0, /* tp_print */
2695 0, /* tp_getattr */
2696 0, /* tp_setattr */
2697 0, /* tp_reserved */
2698 0, /* tp_repr */
2699 0, /* tp_as_number */
2700 0, /* tp_as_sequence */
2701 0, /* tp_as_mapping */
2702 0, /* tp_hash */
2703 0, /* tp_call */
2704 0, /* tp_str */
2705 PyObject_GenericGetAttr, /* tp_getattro */
2706 0, /* tp_setattro */
2707 0, /* tp_as_buffer */
2708 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2709 Py_TPFLAGS_BASETYPE, /* tp_flags */
2710 accumulate_doc, /* tp_doc */
2711 (traverseproc)accumulate_traverse, /* tp_traverse */
2712 0, /* tp_clear */
2713 0, /* tp_richcompare */
2714 0, /* tp_weaklistoffset */
2715 PyObject_SelfIter, /* tp_iter */
2716 (iternextfunc)accumulate_next, /* tp_iternext */
2717 0, /* tp_methods */
2718 0, /* tp_members */
2719 0, /* tp_getset */
2720 0, /* tp_base */
2721 0, /* tp_dict */
2722 0, /* tp_descr_get */
2723 0, /* tp_descr_set */
2724 0, /* tp_dictoffset */
2725 0, /* tp_init */
2726 0, /* tp_alloc */
2727 accumulate_new, /* tp_new */
2728 PyObject_GC_Del, /* tp_free */
2729};
2730
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002731
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002732/* compress object ************************************************************/
2733
2734/* Equivalent to:
2735
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002736 def compress(data, selectors):
2737 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2738 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002739*/
2740
2741typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002742 PyObject_HEAD
2743 PyObject *data;
2744 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002745} compressobject;
2746
2747static PyTypeObject compress_type;
2748
2749static PyObject *
2750compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2751{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002752 PyObject *seq1, *seq2;
2753 PyObject *data=NULL, *selectors=NULL;
2754 compressobject *lz;
2755 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002756
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002757 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2758 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002759
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002760 data = PyObject_GetIter(seq1);
2761 if (data == NULL)
2762 goto fail;
2763 selectors = PyObject_GetIter(seq2);
2764 if (selectors == NULL)
2765 goto fail;
2766
2767 /* create compressobject structure */
2768 lz = (compressobject *)type->tp_alloc(type, 0);
2769 if (lz == NULL)
2770 goto fail;
2771 lz->data = data;
2772 lz->selectors = selectors;
2773 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002774
2775fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002776 Py_XDECREF(data);
2777 Py_XDECREF(selectors);
2778 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002779}
2780
2781static void
2782compress_dealloc(compressobject *lz)
2783{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002784 PyObject_GC_UnTrack(lz);
2785 Py_XDECREF(lz->data);
2786 Py_XDECREF(lz->selectors);
2787 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002788}
2789
2790static int
2791compress_traverse(compressobject *lz, visitproc visit, void *arg)
2792{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002793 Py_VISIT(lz->data);
2794 Py_VISIT(lz->selectors);
2795 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002796}
2797
2798static PyObject *
2799compress_next(compressobject *lz)
2800{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002801 PyObject *data = lz->data, *selectors = lz->selectors;
2802 PyObject *datum, *selector;
2803 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2804 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2805 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002806
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002807 while (1) {
2808 /* Steps: get datum, get selector, evaluate selector.
2809 Order is important (to match the pure python version
2810 in terms of which input gets a chance to raise an
2811 exception first).
2812 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002813
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002814 datum = datanext(data);
2815 if (datum == NULL)
2816 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002817
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002818 selector = selectornext(selectors);
2819 if (selector == NULL) {
2820 Py_DECREF(datum);
2821 return NULL;
2822 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002823
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002824 ok = PyObject_IsTrue(selector);
2825 Py_DECREF(selector);
2826 if (ok == 1)
2827 return datum;
2828 Py_DECREF(datum);
2829 if (ok == -1)
2830 return NULL;
2831 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002832}
2833
2834PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00002835"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002836\n\
2837Return data elements corresponding to true selector elements.\n\
2838Forms a shorter iterator from selected data elements using the\n\
2839selectors to choose the data elements.");
2840
2841static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002842 PyVarObject_HEAD_INIT(NULL, 0)
2843 "itertools.compress", /* tp_name */
2844 sizeof(compressobject), /* tp_basicsize */
2845 0, /* tp_itemsize */
2846 /* methods */
2847 (destructor)compress_dealloc, /* tp_dealloc */
2848 0, /* tp_print */
2849 0, /* tp_getattr */
2850 0, /* tp_setattr */
2851 0, /* tp_reserved */
2852 0, /* tp_repr */
2853 0, /* tp_as_number */
2854 0, /* tp_as_sequence */
2855 0, /* tp_as_mapping */
2856 0, /* tp_hash */
2857 0, /* tp_call */
2858 0, /* tp_str */
2859 PyObject_GenericGetAttr, /* tp_getattro */
2860 0, /* tp_setattro */
2861 0, /* tp_as_buffer */
2862 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2863 Py_TPFLAGS_BASETYPE, /* tp_flags */
2864 compress_doc, /* tp_doc */
2865 (traverseproc)compress_traverse, /* tp_traverse */
2866 0, /* tp_clear */
2867 0, /* tp_richcompare */
2868 0, /* tp_weaklistoffset */
2869 PyObject_SelfIter, /* tp_iter */
2870 (iternextfunc)compress_next, /* tp_iternext */
2871 0, /* tp_methods */
2872 0, /* tp_members */
2873 0, /* tp_getset */
2874 0, /* tp_base */
2875 0, /* tp_dict */
2876 0, /* tp_descr_get */
2877 0, /* tp_descr_set */
2878 0, /* tp_dictoffset */
2879 0, /* tp_init */
2880 0, /* tp_alloc */
2881 compress_new, /* tp_new */
2882 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002883};
2884
2885
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002886/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00002887
2888typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002889 PyObject_HEAD
2890 PyObject *func;
2891 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002892} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002893
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002894static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002895
2896static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002897filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002898{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002899 PyObject *func, *seq;
2900 PyObject *it;
2901 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002902
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002903 if (type == &filterfalse_type &&
2904 !_PyArg_NoKeywords("filterfalse()", kwds))
2905 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002906
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002907 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
2908 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002909
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002910 /* Get iterator. */
2911 it = PyObject_GetIter(seq);
2912 if (it == NULL)
2913 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002914
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002915 /* create filterfalseobject structure */
2916 lz = (filterfalseobject *)type->tp_alloc(type, 0);
2917 if (lz == NULL) {
2918 Py_DECREF(it);
2919 return NULL;
2920 }
2921 Py_INCREF(func);
2922 lz->func = func;
2923 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002924
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002925 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002926}
2927
2928static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002929filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002930{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002931 PyObject_GC_UnTrack(lz);
2932 Py_XDECREF(lz->func);
2933 Py_XDECREF(lz->it);
2934 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002935}
2936
2937static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002938filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002939{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002940 Py_VISIT(lz->it);
2941 Py_VISIT(lz->func);
2942 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002943}
2944
2945static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002946filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002947{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002948 PyObject *item;
2949 PyObject *it = lz->it;
2950 long ok;
2951 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002952
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002953 iternext = *Py_TYPE(it)->tp_iternext;
2954 for (;;) {
2955 item = iternext(it);
2956 if (item == NULL)
2957 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002958
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002959 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2960 ok = PyObject_IsTrue(item);
2961 } else {
2962 PyObject *good;
2963 good = PyObject_CallFunctionObjArgs(lz->func,
2964 item, NULL);
2965 if (good == NULL) {
2966 Py_DECREF(item);
2967 return NULL;
2968 }
2969 ok = PyObject_IsTrue(good);
2970 Py_DECREF(good);
2971 }
2972 if (!ok)
2973 return item;
2974 Py_DECREF(item);
2975 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00002976}
2977
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002978PyDoc_STRVAR(filterfalse_doc,
2979"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002980\n\
2981Return those items of sequence for which function(item) is false.\n\
2982If function is None, return the items that are false.");
2983
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002984static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002985 PyVarObject_HEAD_INIT(NULL, 0)
2986 "itertools.filterfalse", /* tp_name */
2987 sizeof(filterfalseobject), /* tp_basicsize */
2988 0, /* tp_itemsize */
2989 /* methods */
2990 (destructor)filterfalse_dealloc, /* tp_dealloc */
2991 0, /* tp_print */
2992 0, /* tp_getattr */
2993 0, /* tp_setattr */
2994 0, /* tp_reserved */
2995 0, /* tp_repr */
2996 0, /* tp_as_number */
2997 0, /* tp_as_sequence */
2998 0, /* tp_as_mapping */
2999 0, /* tp_hash */
3000 0, /* tp_call */
3001 0, /* tp_str */
3002 PyObject_GenericGetAttr, /* tp_getattro */
3003 0, /* tp_setattro */
3004 0, /* tp_as_buffer */
3005 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3006 Py_TPFLAGS_BASETYPE, /* tp_flags */
3007 filterfalse_doc, /* tp_doc */
3008 (traverseproc)filterfalse_traverse, /* tp_traverse */
3009 0, /* tp_clear */
3010 0, /* tp_richcompare */
3011 0, /* tp_weaklistoffset */
3012 PyObject_SelfIter, /* tp_iter */
3013 (iternextfunc)filterfalse_next, /* tp_iternext */
3014 0, /* tp_methods */
3015 0, /* tp_members */
3016 0, /* tp_getset */
3017 0, /* tp_base */
3018 0, /* tp_dict */
3019 0, /* tp_descr_get */
3020 0, /* tp_descr_set */
3021 0, /* tp_dictoffset */
3022 0, /* tp_init */
3023 0, /* tp_alloc */
3024 filterfalse_new, /* tp_new */
3025 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003026};
3027
3028
3029/* count object ************************************************************/
3030
3031typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003032 PyObject_HEAD
3033 Py_ssize_t cnt;
3034 PyObject *long_cnt;
3035 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003036} countobject;
3037
Raymond Hettinger30729212009-02-12 06:28:27 +00003038/* Counting logic and invariants:
3039
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003040fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003041
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003042 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3043 Advances with: cnt += 1
3044 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003045
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003046slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003047
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003048 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3049 All counting is done with python objects (no overflows or underflows).
3050 Advances with: long_cnt += long_step
3051 Step may be zero -- effectively a slow version of repeat(cnt).
3052 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003053*/
3054
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003055static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003056
3057static PyObject *
3058count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3059{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003060 countobject *lz;
3061 int slow_mode = 0;
3062 Py_ssize_t cnt = 0;
3063 PyObject *long_cnt = NULL;
3064 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003065 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003066 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003067
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003068 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3069 kwlist, &long_cnt, &long_step))
3070 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003071
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003072 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3073 (long_step != NULL && !PyNumber_Check(long_step))) {
3074 PyErr_SetString(PyExc_TypeError, "a number is required");
3075 return NULL;
3076 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003077
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003078 if (long_cnt != NULL) {
3079 cnt = PyLong_AsSsize_t(long_cnt);
3080 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3081 PyErr_Clear();
3082 slow_mode = 1;
3083 }
3084 Py_INCREF(long_cnt);
3085 } else {
3086 cnt = 0;
3087 long_cnt = PyLong_FromLong(0);
3088 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003089
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003090 /* If not specified, step defaults to 1 */
3091 if (long_step == NULL) {
3092 long_step = PyLong_FromLong(1);
3093 if (long_step == NULL) {
3094 Py_DECREF(long_cnt);
3095 return NULL;
3096 }
3097 } else
3098 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003099
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003100 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003101
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003102 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003103 step = PyLong_AsLong(long_step);
3104 if (step != 1) {
3105 slow_mode = 1;
3106 if (step == -1 && PyErr_Occurred())
3107 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003108 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003109
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003110 if (slow_mode)
3111 cnt = PY_SSIZE_T_MAX;
3112 else
3113 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003114
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003115 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3116 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3117 assert(slow_mode ||
3118 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003119
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003120 /* create countobject structure */
3121 lz = (countobject *)type->tp_alloc(type, 0);
3122 if (lz == NULL) {
3123 Py_XDECREF(long_cnt);
3124 return NULL;
3125 }
3126 lz->cnt = cnt;
3127 lz->long_cnt = long_cnt;
3128 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003129
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003130 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003131}
3132
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003133static void
3134count_dealloc(countobject *lz)
3135{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003136 PyObject_GC_UnTrack(lz);
3137 Py_XDECREF(lz->long_cnt);
3138 Py_XDECREF(lz->long_step);
3139 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003140}
3141
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003142static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003143count_traverse(countobject *lz, visitproc visit, void *arg)
3144{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003145 Py_VISIT(lz->long_cnt);
3146 Py_VISIT(lz->long_step);
3147 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003148}
3149
3150static PyObject *
3151count_nextlong(countobject *lz)
3152{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003153 PyObject *long_cnt;
3154 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003155
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003156 long_cnt = lz->long_cnt;
3157 if (long_cnt == NULL) {
3158 /* Switch to slow_mode */
3159 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3160 if (long_cnt == NULL)
3161 return NULL;
3162 }
3163 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00003164
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003165 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3166 if (stepped_up == NULL)
3167 return NULL;
3168 lz->long_cnt = stepped_up;
3169 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003170}
3171
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003172static PyObject *
3173count_next(countobject *lz)
3174{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003175 if (lz->cnt == PY_SSIZE_T_MAX)
3176 return count_nextlong(lz);
3177 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003178}
3179
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003180static PyObject *
3181count_repr(countobject *lz)
3182{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003183 if (lz->cnt != PY_SSIZE_T_MAX)
3184 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003185
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003186 if (PyLong_Check(lz->long_step)) {
3187 long step = PyLong_AsLong(lz->long_step);
3188 if (step == -1 && PyErr_Occurred()) {
3189 PyErr_Clear();
3190 }
3191 if (step == 1) {
3192 /* Don't display step when it is an integer equal to 1 */
3193 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
3194 }
3195 }
3196 return PyUnicode_FromFormat("count(%R, %R)",
3197 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003198}
3199
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003200static PyObject *
3201count_reduce(countobject *lz)
3202{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003203 if (lz->cnt == PY_SSIZE_T_MAX)
3204 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3205 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003206}
3207
3208PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3209
3210static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003211 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3212 count_reduce_doc},
3213 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003214};
3215
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003216PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00003217 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003218\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003219Return a count object whose .__next__() method returns consecutive values.\n\
3220Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00003221 def count(firstval=0, step=1):\n\
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003222 x = firstval\n\
3223 while 1:\n\
3224 yield x\n\
3225 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003226
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003227static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003228 PyVarObject_HEAD_INIT(NULL, 0)
3229 "itertools.count", /* tp_name */
3230 sizeof(countobject), /* tp_basicsize */
3231 0, /* tp_itemsize */
3232 /* methods */
3233 (destructor)count_dealloc, /* tp_dealloc */
3234 0, /* tp_print */
3235 0, /* tp_getattr */
3236 0, /* tp_setattr */
3237 0, /* tp_reserved */
3238 (reprfunc)count_repr, /* tp_repr */
3239 0, /* tp_as_number */
3240 0, /* tp_as_sequence */
3241 0, /* tp_as_mapping */
3242 0, /* tp_hash */
3243 0, /* tp_call */
3244 0, /* tp_str */
3245 PyObject_GenericGetAttr, /* tp_getattro */
3246 0, /* tp_setattro */
3247 0, /* tp_as_buffer */
3248 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3249 Py_TPFLAGS_BASETYPE, /* tp_flags */
3250 count_doc, /* tp_doc */
3251 (traverseproc)count_traverse, /* tp_traverse */
3252 0, /* tp_clear */
3253 0, /* tp_richcompare */
3254 0, /* tp_weaklistoffset */
3255 PyObject_SelfIter, /* tp_iter */
3256 (iternextfunc)count_next, /* tp_iternext */
3257 count_methods, /* tp_methods */
3258 0, /* tp_members */
3259 0, /* tp_getset */
3260 0, /* tp_base */
3261 0, /* tp_dict */
3262 0, /* tp_descr_get */
3263 0, /* tp_descr_set */
3264 0, /* tp_dictoffset */
3265 0, /* tp_init */
3266 0, /* tp_alloc */
3267 count_new, /* tp_new */
3268 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003269};
3270
3271
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003272/* repeat object ************************************************************/
3273
3274typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003275 PyObject_HEAD
3276 PyObject *element;
3277 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003278} repeatobject;
3279
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003280static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003281
3282static PyObject *
3283repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3284{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003285 repeatobject *ro;
3286 PyObject *element;
3287 Py_ssize_t cnt = -1;
3288 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003289
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003290 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3291 &element, &cnt))
3292 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003293
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003294 if (PyTuple_Size(args) == 2 && cnt < 0)
3295 cnt = 0;
3296
3297 ro = (repeatobject *)type->tp_alloc(type, 0);
3298 if (ro == NULL)
3299 return NULL;
3300 Py_INCREF(element);
3301 ro->element = element;
3302 ro->cnt = cnt;
3303 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003304}
3305
3306static void
3307repeat_dealloc(repeatobject *ro)
3308{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003309 PyObject_GC_UnTrack(ro);
3310 Py_XDECREF(ro->element);
3311 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003312}
3313
3314static int
3315repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3316{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003317 Py_VISIT(ro->element);
3318 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003319}
3320
3321static PyObject *
3322repeat_next(repeatobject *ro)
3323{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003324 if (ro->cnt == 0)
3325 return NULL;
3326 if (ro->cnt > 0)
3327 ro->cnt--;
3328 Py_INCREF(ro->element);
3329 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003330}
3331
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003332static PyObject *
3333repeat_repr(repeatobject *ro)
3334{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003335 if (ro->cnt == -1)
3336 return PyUnicode_FromFormat("repeat(%R)", ro->element);
3337 else
3338 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
3339}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003340
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003341static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003342repeat_len(repeatobject *ro)
3343{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003344 if (ro->cnt == -1) {
3345 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3346 return NULL;
3347 }
3348 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003349}
3350
Armin Rigof5b3e362006-02-11 21:32:43 +00003351PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003352
3353static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003354 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3355 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003356};
3357
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003358PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00003359"repeat(object [,times]) -> create an iterator which returns the object\n\
3360for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003361endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003362
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003363static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003364 PyVarObject_HEAD_INIT(NULL, 0)
3365 "itertools.repeat", /* tp_name */
3366 sizeof(repeatobject), /* tp_basicsize */
3367 0, /* tp_itemsize */
3368 /* methods */
3369 (destructor)repeat_dealloc, /* tp_dealloc */
3370 0, /* tp_print */
3371 0, /* tp_getattr */
3372 0, /* tp_setattr */
3373 0, /* tp_reserved */
3374 (reprfunc)repeat_repr, /* tp_repr */
3375 0, /* tp_as_number */
3376 0, /* tp_as_sequence */
3377 0, /* tp_as_mapping */
3378 0, /* tp_hash */
3379 0, /* tp_call */
3380 0, /* tp_str */
3381 PyObject_GenericGetAttr, /* tp_getattro */
3382 0, /* tp_setattro */
3383 0, /* tp_as_buffer */
3384 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3385 Py_TPFLAGS_BASETYPE, /* tp_flags */
3386 repeat_doc, /* tp_doc */
3387 (traverseproc)repeat_traverse, /* tp_traverse */
3388 0, /* tp_clear */
3389 0, /* tp_richcompare */
3390 0, /* tp_weaklistoffset */
3391 PyObject_SelfIter, /* tp_iter */
3392 (iternextfunc)repeat_next, /* tp_iternext */
3393 repeat_methods, /* tp_methods */
3394 0, /* tp_members */
3395 0, /* tp_getset */
3396 0, /* tp_base */
3397 0, /* tp_dict */
3398 0, /* tp_descr_get */
3399 0, /* tp_descr_set */
3400 0, /* tp_dictoffset */
3401 0, /* tp_init */
3402 0, /* tp_alloc */
3403 repeat_new, /* tp_new */
3404 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003405};
3406
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003407/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00003408
3409#include "Python.h"
3410
3411typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003412 PyObject_HEAD
3413 Py_ssize_t tuplesize;
3414 Py_ssize_t numactive;
3415 PyObject *ittuple; /* tuple of iterators */
3416 PyObject *result;
3417 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003418} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003419
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003420static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003421
3422static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003423zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003424{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003425 ziplongestobject *lz;
3426 Py_ssize_t i;
3427 PyObject *ittuple; /* tuple of iterators */
3428 PyObject *result;
3429 PyObject *fillvalue = Py_None;
3430 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003431
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003432 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3433 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3434 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3435 PyErr_SetString(PyExc_TypeError,
3436 "zip_longest() got an unexpected keyword argument");
3437 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003438 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003439 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003440
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003441 /* args must be a tuple */
3442 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00003443
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003444 /* obtain iterators */
3445 ittuple = PyTuple_New(tuplesize);
3446 if (ittuple == NULL)
3447 return NULL;
3448 for (i=0; i < tuplesize; ++i) {
3449 PyObject *item = PyTuple_GET_ITEM(args, i);
3450 PyObject *it = PyObject_GetIter(item);
3451 if (it == NULL) {
3452 if (PyErr_ExceptionMatches(PyExc_TypeError))
3453 PyErr_Format(PyExc_TypeError,
3454 "zip_longest argument #%zd must support iteration",
3455 i+1);
3456 Py_DECREF(ittuple);
3457 return NULL;
3458 }
3459 PyTuple_SET_ITEM(ittuple, i, it);
3460 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003461
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003462 /* create a result holder */
3463 result = PyTuple_New(tuplesize);
3464 if (result == NULL) {
3465 Py_DECREF(ittuple);
3466 return NULL;
3467 }
3468 for (i=0 ; i < tuplesize ; i++) {
3469 Py_INCREF(Py_None);
3470 PyTuple_SET_ITEM(result, i, Py_None);
3471 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003472
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003473 /* create ziplongestobject structure */
3474 lz = (ziplongestobject *)type->tp_alloc(type, 0);
3475 if (lz == NULL) {
3476 Py_DECREF(ittuple);
3477 Py_DECREF(result);
3478 return NULL;
3479 }
3480 lz->ittuple = ittuple;
3481 lz->tuplesize = tuplesize;
3482 lz->numactive = tuplesize;
3483 lz->result = result;
3484 Py_INCREF(fillvalue);
3485 lz->fillvalue = fillvalue;
3486 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003487}
3488
3489static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003490zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003491{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003492 PyObject_GC_UnTrack(lz);
3493 Py_XDECREF(lz->ittuple);
3494 Py_XDECREF(lz->result);
3495 Py_XDECREF(lz->fillvalue);
3496 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003497}
3498
3499static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003500zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003501{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003502 Py_VISIT(lz->ittuple);
3503 Py_VISIT(lz->result);
3504 Py_VISIT(lz->fillvalue);
3505 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003506}
3507
3508static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003509zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003510{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003511 Py_ssize_t i;
3512 Py_ssize_t tuplesize = lz->tuplesize;
3513 PyObject *result = lz->result;
3514 PyObject *it;
3515 PyObject *item;
3516 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003517
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003518 if (tuplesize == 0)
3519 return NULL;
3520 if (lz->numactive == 0)
3521 return NULL;
3522 if (Py_REFCNT(result) == 1) {
3523 Py_INCREF(result);
3524 for (i=0 ; i < tuplesize ; i++) {
3525 it = PyTuple_GET_ITEM(lz->ittuple, i);
3526 if (it == NULL) {
3527 Py_INCREF(lz->fillvalue);
3528 item = lz->fillvalue;
3529 } else {
3530 item = PyIter_Next(it);
3531 if (item == NULL) {
3532 lz->numactive -= 1;
3533 if (lz->numactive == 0 || PyErr_Occurred()) {
3534 lz->numactive = 0;
3535 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00003536 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003537 } else {
3538 Py_INCREF(lz->fillvalue);
3539 item = lz->fillvalue;
3540 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3541 Py_DECREF(it);
3542 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00003543 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003544 }
3545 olditem = PyTuple_GET_ITEM(result, i);
3546 PyTuple_SET_ITEM(result, i, item);
3547 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00003548 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003549 } else {
3550 result = PyTuple_New(tuplesize);
3551 if (result == NULL)
3552 return NULL;
3553 for (i=0 ; i < tuplesize ; i++) {
3554 it = PyTuple_GET_ITEM(lz->ittuple, i);
3555 if (it == NULL) {
3556 Py_INCREF(lz->fillvalue);
3557 item = lz->fillvalue;
3558 } else {
3559 item = PyIter_Next(it);
3560 if (item == NULL) {
3561 lz->numactive -= 1;
3562 if (lz->numactive == 0 || PyErr_Occurred()) {
3563 lz->numactive = 0;
3564 Py_DECREF(result);
3565 return NULL;
3566 } else {
3567 Py_INCREF(lz->fillvalue);
3568 item = lz->fillvalue;
3569 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3570 Py_DECREF(it);
3571 }
3572 }
3573 }
3574 PyTuple_SET_ITEM(result, i, item);
3575 }
3576 }
3577 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003578}
3579
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003580PyDoc_STRVAR(zip_longest_doc,
3581"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003582\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003583Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00003584the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003585method continues until the longest iterable in the argument sequence\n\
3586is exhausted and then it raises StopIteration. When the shorter iterables\n\
3587are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3588defaults to None or can be specified by a keyword argument.\n\
3589");
3590
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003591static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003592 PyVarObject_HEAD_INIT(NULL, 0)
3593 "itertools.zip_longest", /* tp_name */
3594 sizeof(ziplongestobject), /* tp_basicsize */
3595 0, /* tp_itemsize */
3596 /* methods */
3597 (destructor)zip_longest_dealloc, /* tp_dealloc */
3598 0, /* tp_print */
3599 0, /* tp_getattr */
3600 0, /* tp_setattr */
3601 0, /* tp_reserved */
3602 0, /* tp_repr */
3603 0, /* tp_as_number */
3604 0, /* tp_as_sequence */
3605 0, /* tp_as_mapping */
3606 0, /* tp_hash */
3607 0, /* tp_call */
3608 0, /* tp_str */
3609 PyObject_GenericGetAttr, /* tp_getattro */
3610 0, /* tp_setattro */
3611 0, /* tp_as_buffer */
3612 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3613 Py_TPFLAGS_BASETYPE, /* tp_flags */
3614 zip_longest_doc, /* tp_doc */
3615 (traverseproc)zip_longest_traverse, /* tp_traverse */
3616 0, /* tp_clear */
3617 0, /* tp_richcompare */
3618 0, /* tp_weaklistoffset */
3619 PyObject_SelfIter, /* tp_iter */
3620 (iternextfunc)zip_longest_next, /* tp_iternext */
3621 0, /* tp_methods */
3622 0, /* tp_members */
3623 0, /* tp_getset */
3624 0, /* tp_base */
3625 0, /* tp_dict */
3626 0, /* tp_descr_get */
3627 0, /* tp_descr_set */
3628 0, /* tp_dictoffset */
3629 0, /* tp_init */
3630 0, /* tp_alloc */
3631 zip_longest_new, /* tp_new */
3632 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003633};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003634
3635/* module level code ********************************************************/
3636
3637PyDoc_STRVAR(module_doc,
3638"Functional tools for creating and using iterators.\n\
3639\n\
3640Infinite iterators:\n\
3641count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003642cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003643repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003644\n\
3645Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003646accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003647chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
3648compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
3649dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
3650groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003651filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003652islice(seq, [start,] stop [, step]) --> elements from\n\
3653 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003654starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003655tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003656takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003657zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00003658\n\
3659Combinatoric generators:\n\
3660product(p, q, ... [repeat=1]) --> cartesian product\n\
3661permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00003662combinations(p, r)\n\
3663combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003664");
3665
3666
Raymond Hettingerad983e72003-11-12 14:32:26 +00003667static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003668 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3669 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00003670};
3671
Martin v. Löwis1a214512008-06-11 05:26:20 +00003672
3673static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003674 PyModuleDef_HEAD_INIT,
3675 "itertools",
3676 module_doc,
3677 -1,
3678 module_methods,
3679 NULL,
3680 NULL,
3681 NULL,
3682 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00003683};
3684
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003685PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00003686PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003687{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003688 int i;
3689 PyObject *m;
3690 char *name;
3691 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00003692 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003693 &combinations_type,
3694 &cwr_type,
3695 &cycle_type,
3696 &dropwhile_type,
3697 &takewhile_type,
3698 &islice_type,
3699 &starmap_type,
3700 &chain_type,
3701 &compress_type,
3702 &filterfalse_type,
3703 &count_type,
3704 &ziplongest_type,
3705 &permutations_type,
3706 &product_type,
3707 &repeat_type,
3708 &groupby_type,
3709 NULL
3710 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00003711
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003712 Py_TYPE(&teedataobject_type) = &PyType_Type;
3713 m = PyModule_Create(&itertoolsmodule);
3714 if (m == NULL)
3715 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003716
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003717 for (i=0 ; typelist[i] != NULL ; i++) {
3718 if (PyType_Ready(typelist[i]) < 0)
3719 return NULL;
3720 name = strchr(typelist[i]->tp_name, '.');
3721 assert (name != NULL);
3722 Py_INCREF(typelist[i]);
3723 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
3724 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003725
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003726 if (PyType_Ready(&teedataobject_type) < 0)
3727 return NULL;
3728 if (PyType_Ready(&tee_type) < 0)
3729 return NULL;
3730 if (PyType_Ready(&_grouper_type) < 0)
3731 return NULL;
3732 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003733}