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