blob: 77e76fe588042df8727ad3e9fb96b2c1405898c6 [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);
Antoine Pitrou6f430e42012-08-15 23:18:25 +0200906 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000907 lz->start = 1;
908 return item;
909 }
910 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +0200911 if (ok < 0)
912 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000913 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000914}
915
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000916PyDoc_STRVAR(dropwhile_doc,
917"dropwhile(predicate, iterable) --> dropwhile object\n\
918\n\
919Drop items from the iterable while predicate(item) is true.\n\
920Afterwards, return every element until the iterable is exhausted.");
921
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000922static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000923 PyVarObject_HEAD_INIT(NULL, 0)
924 "itertools.dropwhile", /* tp_name */
925 sizeof(dropwhileobject), /* tp_basicsize */
926 0, /* tp_itemsize */
927 /* methods */
928 (destructor)dropwhile_dealloc, /* tp_dealloc */
929 0, /* tp_print */
930 0, /* tp_getattr */
931 0, /* tp_setattr */
932 0, /* tp_reserved */
933 0, /* tp_repr */
934 0, /* tp_as_number */
935 0, /* tp_as_sequence */
936 0, /* tp_as_mapping */
937 0, /* tp_hash */
938 0, /* tp_call */
939 0, /* tp_str */
940 PyObject_GenericGetAttr, /* tp_getattro */
941 0, /* tp_setattro */
942 0, /* tp_as_buffer */
943 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
944 Py_TPFLAGS_BASETYPE, /* tp_flags */
945 dropwhile_doc, /* tp_doc */
946 (traverseproc)dropwhile_traverse, /* tp_traverse */
947 0, /* tp_clear */
948 0, /* tp_richcompare */
949 0, /* tp_weaklistoffset */
950 PyObject_SelfIter, /* tp_iter */
951 (iternextfunc)dropwhile_next, /* tp_iternext */
952 0, /* tp_methods */
953 0, /* tp_members */
954 0, /* tp_getset */
955 0, /* tp_base */
956 0, /* tp_dict */
957 0, /* tp_descr_get */
958 0, /* tp_descr_set */
959 0, /* tp_dictoffset */
960 0, /* tp_init */
961 0, /* tp_alloc */
962 dropwhile_new, /* tp_new */
963 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000964};
965
966
967/* takewhile object **********************************************************/
968
969typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000970 PyObject_HEAD
971 PyObject *func;
972 PyObject *it;
973 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000974} takewhileobject;
975
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000976static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000977
978static PyObject *
979takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
980{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000981 PyObject *func, *seq;
982 PyObject *it;
983 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000984
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000985 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
986 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000987
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000988 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
989 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000990
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000991 /* Get iterator. */
992 it = PyObject_GetIter(seq);
993 if (it == NULL)
994 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000995
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000996 /* create takewhileobject structure */
997 lz = (takewhileobject *)type->tp_alloc(type, 0);
998 if (lz == NULL) {
999 Py_DECREF(it);
1000 return NULL;
1001 }
1002 Py_INCREF(func);
1003 lz->func = func;
1004 lz->it = it;
1005 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001006
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001007 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001008}
1009
1010static void
1011takewhile_dealloc(takewhileobject *lz)
1012{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001013 PyObject_GC_UnTrack(lz);
1014 Py_XDECREF(lz->func);
1015 Py_XDECREF(lz->it);
1016 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001017}
1018
1019static int
1020takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1021{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001022 Py_VISIT(lz->it);
1023 Py_VISIT(lz->func);
1024 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001025}
1026
1027static PyObject *
1028takewhile_next(takewhileobject *lz)
1029{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001030 PyObject *item, *good;
1031 PyObject *it = lz->it;
1032 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001033
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001034 if (lz->stop == 1)
1035 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001036
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001037 item = (*Py_TYPE(it)->tp_iternext)(it);
1038 if (item == NULL)
1039 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001040
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001041 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1042 if (good == NULL) {
1043 Py_DECREF(item);
1044 return NULL;
1045 }
1046 ok = PyObject_IsTrue(good);
1047 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001048 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001049 return item;
1050 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001051 if (ok == 0)
1052 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001053 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001054}
1055
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001056PyDoc_STRVAR(takewhile_doc,
1057"takewhile(predicate, iterable) --> takewhile object\n\
1058\n\
1059Return successive entries from an iterable as long as the \n\
1060predicate evaluates to true for each entry.");
1061
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001062static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001063 PyVarObject_HEAD_INIT(NULL, 0)
1064 "itertools.takewhile", /* tp_name */
1065 sizeof(takewhileobject), /* tp_basicsize */
1066 0, /* tp_itemsize */
1067 /* methods */
1068 (destructor)takewhile_dealloc, /* tp_dealloc */
1069 0, /* tp_print */
1070 0, /* tp_getattr */
1071 0, /* tp_setattr */
1072 0, /* tp_reserved */
1073 0, /* tp_repr */
1074 0, /* tp_as_number */
1075 0, /* tp_as_sequence */
1076 0, /* tp_as_mapping */
1077 0, /* tp_hash */
1078 0, /* tp_call */
1079 0, /* tp_str */
1080 PyObject_GenericGetAttr, /* tp_getattro */
1081 0, /* tp_setattro */
1082 0, /* tp_as_buffer */
1083 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1084 Py_TPFLAGS_BASETYPE, /* tp_flags */
1085 takewhile_doc, /* tp_doc */
1086 (traverseproc)takewhile_traverse, /* tp_traverse */
1087 0, /* tp_clear */
1088 0, /* tp_richcompare */
1089 0, /* tp_weaklistoffset */
1090 PyObject_SelfIter, /* tp_iter */
1091 (iternextfunc)takewhile_next, /* tp_iternext */
1092 0, /* tp_methods */
1093 0, /* tp_members */
1094 0, /* tp_getset */
1095 0, /* tp_base */
1096 0, /* tp_dict */
1097 0, /* tp_descr_get */
1098 0, /* tp_descr_set */
1099 0, /* tp_dictoffset */
1100 0, /* tp_init */
1101 0, /* tp_alloc */
1102 takewhile_new, /* tp_new */
1103 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001104};
1105
1106
1107/* islice object ************************************************************/
1108
1109typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001110 PyObject_HEAD
1111 PyObject *it;
1112 Py_ssize_t next;
1113 Py_ssize_t stop;
1114 Py_ssize_t step;
1115 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116} isliceobject;
1117
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001118static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001119
1120static PyObject *
1121islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1122{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001123 PyObject *seq;
1124 Py_ssize_t start=0, stop=-1, step=1;
1125 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1126 Py_ssize_t numargs;
1127 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001128
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001129 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1130 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001131
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001132 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1133 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001134
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001135 numargs = PyTuple_Size(args);
1136 if (numargs == 2) {
1137 if (a1 != Py_None) {
1138 stop = PyLong_AsSsize_t(a1);
1139 if (stop == -1) {
1140 if (PyErr_Occurred())
1141 PyErr_Clear();
1142 PyErr_SetString(PyExc_ValueError,
1143 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1144 return NULL;
1145 }
1146 }
1147 } else {
1148 if (a1 != Py_None)
1149 start = PyLong_AsSsize_t(a1);
1150 if (start == -1 && PyErr_Occurred())
1151 PyErr_Clear();
1152 if (a2 != Py_None) {
1153 stop = PyLong_AsSsize_t(a2);
1154 if (stop == -1) {
1155 if (PyErr_Occurred())
1156 PyErr_Clear();
1157 PyErr_SetString(PyExc_ValueError,
1158 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1159 return NULL;
1160 }
1161 }
1162 }
1163 if (start<0 || stop<-1) {
1164 PyErr_SetString(PyExc_ValueError,
1165 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1166 return NULL;
1167 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001168
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001169 if (a3 != NULL) {
1170 if (a3 != Py_None)
1171 step = PyLong_AsSsize_t(a3);
1172 if (step == -1 && PyErr_Occurred())
1173 PyErr_Clear();
1174 }
1175 if (step<1) {
1176 PyErr_SetString(PyExc_ValueError,
1177 "Step for islice() must be a positive integer or None.");
1178 return NULL;
1179 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001180
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001181 /* Get iterator. */
1182 it = PyObject_GetIter(seq);
1183 if (it == NULL)
1184 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001185
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001186 /* create isliceobject structure */
1187 lz = (isliceobject *)type->tp_alloc(type, 0);
1188 if (lz == NULL) {
1189 Py_DECREF(it);
1190 return NULL;
1191 }
1192 lz->it = it;
1193 lz->next = start;
1194 lz->stop = stop;
1195 lz->step = step;
1196 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001197
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001198 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001199}
1200
1201static void
1202islice_dealloc(isliceobject *lz)
1203{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001204 PyObject_GC_UnTrack(lz);
1205 Py_XDECREF(lz->it);
1206 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001207}
1208
1209static int
1210islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1211{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001212 Py_VISIT(lz->it);
1213 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001214}
1215
1216static PyObject *
1217islice_next(isliceobject *lz)
1218{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001219 PyObject *item;
1220 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001221 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001222 Py_ssize_t oldnext;
1223 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001224
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001225 iternext = *Py_TYPE(it)->tp_iternext;
1226 while (lz->cnt < lz->next) {
1227 item = iternext(it);
1228 if (item == NULL)
1229 return NULL;
1230 Py_DECREF(item);
1231 lz->cnt++;
1232 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001233 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001234 return NULL;
1235 item = iternext(it);
1236 if (item == NULL)
1237 return NULL;
1238 lz->cnt++;
1239 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001240 /* The (size_t) cast below avoids the danger of undefined
1241 behaviour from signed integer overflow. */
1242 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001243 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1244 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001245 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001246}
1247
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001248PyDoc_STRVAR(islice_doc,
1249"islice(iterable, [start,] stop [, step]) --> islice object\n\
1250\n\
1251Return an iterator whose next() method returns selected values from an\n\
1252iterable. If start is specified, will skip all preceding elements;\n\
1253otherwise, start defaults to zero. Step defaults to one. If\n\
1254specified as another value, step determines how many values are \n\
1255skipped between successive calls. Works like a slice() on a list\n\
1256but returns an iterator.");
1257
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001258static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001259 PyVarObject_HEAD_INIT(NULL, 0)
1260 "itertools.islice", /* tp_name */
1261 sizeof(isliceobject), /* tp_basicsize */
1262 0, /* tp_itemsize */
1263 /* methods */
1264 (destructor)islice_dealloc, /* tp_dealloc */
1265 0, /* tp_print */
1266 0, /* tp_getattr */
1267 0, /* tp_setattr */
1268 0, /* tp_reserved */
1269 0, /* tp_repr */
1270 0, /* tp_as_number */
1271 0, /* tp_as_sequence */
1272 0, /* tp_as_mapping */
1273 0, /* tp_hash */
1274 0, /* tp_call */
1275 0, /* tp_str */
1276 PyObject_GenericGetAttr, /* tp_getattro */
1277 0, /* tp_setattro */
1278 0, /* tp_as_buffer */
1279 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1280 Py_TPFLAGS_BASETYPE, /* tp_flags */
1281 islice_doc, /* tp_doc */
1282 (traverseproc)islice_traverse, /* tp_traverse */
1283 0, /* tp_clear */
1284 0, /* tp_richcompare */
1285 0, /* tp_weaklistoffset */
1286 PyObject_SelfIter, /* tp_iter */
1287 (iternextfunc)islice_next, /* tp_iternext */
1288 0, /* tp_methods */
1289 0, /* tp_members */
1290 0, /* tp_getset */
1291 0, /* tp_base */
1292 0, /* tp_dict */
1293 0, /* tp_descr_get */
1294 0, /* tp_descr_set */
1295 0, /* tp_dictoffset */
1296 0, /* tp_init */
1297 0, /* tp_alloc */
1298 islice_new, /* tp_new */
1299 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001300};
1301
1302
1303/* starmap object ************************************************************/
1304
1305typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001306 PyObject_HEAD
1307 PyObject *func;
1308 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001309} starmapobject;
1310
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001311static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001312
1313static PyObject *
1314starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1315{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001316 PyObject *func, *seq;
1317 PyObject *it;
1318 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001319
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001320 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1321 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001322
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001323 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1324 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001325
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001326 /* Get iterator. */
1327 it = PyObject_GetIter(seq);
1328 if (it == NULL)
1329 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001330
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001331 /* create starmapobject structure */
1332 lz = (starmapobject *)type->tp_alloc(type, 0);
1333 if (lz == NULL) {
1334 Py_DECREF(it);
1335 return NULL;
1336 }
1337 Py_INCREF(func);
1338 lz->func = func;
1339 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001340
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001341 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001342}
1343
1344static void
1345starmap_dealloc(starmapobject *lz)
1346{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001347 PyObject_GC_UnTrack(lz);
1348 Py_XDECREF(lz->func);
1349 Py_XDECREF(lz->it);
1350 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001351}
1352
1353static int
1354starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1355{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001356 Py_VISIT(lz->it);
1357 Py_VISIT(lz->func);
1358 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001359}
1360
1361static PyObject *
1362starmap_next(starmapobject *lz)
1363{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001364 PyObject *args;
1365 PyObject *result;
1366 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001367
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001368 args = (*Py_TYPE(it)->tp_iternext)(it);
1369 if (args == NULL)
1370 return NULL;
1371 if (!PyTuple_CheckExact(args)) {
1372 PyObject *newargs = PySequence_Tuple(args);
1373 Py_DECREF(args);
1374 if (newargs == NULL)
1375 return NULL;
1376 args = newargs;
1377 }
1378 result = PyObject_Call(lz->func, args, NULL);
1379 Py_DECREF(args);
1380 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001381}
1382
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001383PyDoc_STRVAR(starmap_doc,
1384"starmap(function, sequence) --> starmap object\n\
1385\n\
1386Return an iterator whose values are returned from the function evaluated\n\
1387with a argument tuple taken from the given sequence.");
1388
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001389static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001390 PyVarObject_HEAD_INIT(NULL, 0)
1391 "itertools.starmap", /* tp_name */
1392 sizeof(starmapobject), /* tp_basicsize */
1393 0, /* tp_itemsize */
1394 /* methods */
1395 (destructor)starmap_dealloc, /* tp_dealloc */
1396 0, /* tp_print */
1397 0, /* tp_getattr */
1398 0, /* tp_setattr */
1399 0, /* tp_reserved */
1400 0, /* tp_repr */
1401 0, /* tp_as_number */
1402 0, /* tp_as_sequence */
1403 0, /* tp_as_mapping */
1404 0, /* tp_hash */
1405 0, /* tp_call */
1406 0, /* tp_str */
1407 PyObject_GenericGetAttr, /* tp_getattro */
1408 0, /* tp_setattro */
1409 0, /* tp_as_buffer */
1410 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1411 Py_TPFLAGS_BASETYPE, /* tp_flags */
1412 starmap_doc, /* tp_doc */
1413 (traverseproc)starmap_traverse, /* tp_traverse */
1414 0, /* tp_clear */
1415 0, /* tp_richcompare */
1416 0, /* tp_weaklistoffset */
1417 PyObject_SelfIter, /* tp_iter */
1418 (iternextfunc)starmap_next, /* tp_iternext */
1419 0, /* tp_methods */
1420 0, /* tp_members */
1421 0, /* tp_getset */
1422 0, /* tp_base */
1423 0, /* tp_dict */
1424 0, /* tp_descr_get */
1425 0, /* tp_descr_set */
1426 0, /* tp_dictoffset */
1427 0, /* tp_init */
1428 0, /* tp_alloc */
1429 starmap_new, /* tp_new */
1430 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001431};
1432
1433
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001434/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001435
1436typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001437 PyObject_HEAD
1438 PyObject *source; /* Iterator over input iterables */
1439 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001440} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001441
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001442static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001443
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001444static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001445chain_new_internal(PyTypeObject *type, PyObject *source)
1446{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001447 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001448
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001449 lz = (chainobject *)type->tp_alloc(type, 0);
1450 if (lz == NULL) {
1451 Py_DECREF(source);
1452 return NULL;
1453 }
1454
1455 lz->source = source;
1456 lz->active = NULL;
1457 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001458}
1459
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001460static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001461chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001462{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001463 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001464
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001465 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1466 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001467
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001468 source = PyObject_GetIter(args);
1469 if (source == NULL)
1470 return NULL;
1471
1472 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001473}
1474
1475static PyObject *
1476chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1477{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001478 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001479
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001480 source = PyObject_GetIter(arg);
1481 if (source == NULL)
1482 return NULL;
1483
1484 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001485}
1486
1487static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001488chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001489{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001490 PyObject_GC_UnTrack(lz);
1491 Py_XDECREF(lz->active);
1492 Py_XDECREF(lz->source);
1493 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001494}
1495
Raymond Hettinger2012f172003-02-07 05:32:58 +00001496static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001497chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001498{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001499 Py_VISIT(lz->source);
1500 Py_VISIT(lz->active);
1501 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001502}
1503
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001504static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001505chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001506{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001507 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001508
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001509 if (lz->source == NULL)
1510 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001511
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001512 if (lz->active == NULL) {
1513 PyObject *iterable = PyIter_Next(lz->source);
1514 if (iterable == NULL) {
1515 Py_CLEAR(lz->source);
1516 return NULL; /* no more input sources */
1517 }
1518 lz->active = PyObject_GetIter(iterable);
1519 Py_DECREF(iterable);
1520 if (lz->active == NULL) {
1521 Py_CLEAR(lz->source);
1522 return NULL; /* input not iterable */
1523 }
1524 }
1525 item = PyIter_Next(lz->active);
1526 if (item != NULL)
1527 return item;
1528 if (PyErr_Occurred()) {
1529 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1530 PyErr_Clear();
1531 else
1532 return NULL; /* input raised an exception */
1533 }
1534 Py_CLEAR(lz->active);
1535 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001536}
1537
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001538PyDoc_STRVAR(chain_doc,
1539"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001540\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001541Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001542first iterable until it is exhausted, then elements from the next\n\
1543iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001544
Christian Heimesf16baeb2008-02-29 14:57:44 +00001545PyDoc_STRVAR(chain_from_iterable_doc,
1546"chain.from_iterable(iterable) --> chain object\n\
1547\n\
1548Alternate chain() contructor taking a single iterable argument\n\
1549that evaluates lazily.");
1550
1551static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001552 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1553 chain_from_iterable_doc},
1554 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001555};
1556
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001557static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001558 PyVarObject_HEAD_INIT(NULL, 0)
1559 "itertools.chain", /* tp_name */
1560 sizeof(chainobject), /* tp_basicsize */
1561 0, /* tp_itemsize */
1562 /* methods */
1563 (destructor)chain_dealloc, /* tp_dealloc */
1564 0, /* tp_print */
1565 0, /* tp_getattr */
1566 0, /* tp_setattr */
1567 0, /* tp_reserved */
1568 0, /* tp_repr */
1569 0, /* tp_as_number */
1570 0, /* tp_as_sequence */
1571 0, /* tp_as_mapping */
1572 0, /* tp_hash */
1573 0, /* tp_call */
1574 0, /* tp_str */
1575 PyObject_GenericGetAttr, /* tp_getattro */
1576 0, /* tp_setattro */
1577 0, /* tp_as_buffer */
1578 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1579 Py_TPFLAGS_BASETYPE, /* tp_flags */
1580 chain_doc, /* tp_doc */
1581 (traverseproc)chain_traverse, /* tp_traverse */
1582 0, /* tp_clear */
1583 0, /* tp_richcompare */
1584 0, /* tp_weaklistoffset */
1585 PyObject_SelfIter, /* tp_iter */
1586 (iternextfunc)chain_next, /* tp_iternext */
1587 chain_methods, /* tp_methods */
1588 0, /* tp_members */
1589 0, /* tp_getset */
1590 0, /* tp_base */
1591 0, /* tp_dict */
1592 0, /* tp_descr_get */
1593 0, /* tp_descr_set */
1594 0, /* tp_dictoffset */
1595 0, /* tp_init */
1596 0, /* tp_alloc */
1597 chain_new, /* tp_new */
1598 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001599};
1600
1601
Christian Heimesc3f30c42008-02-22 16:37:40 +00001602/* product object ************************************************************/
1603
1604typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001605 PyObject_HEAD
1606 PyObject *pools; /* tuple of pool tuples */
1607 Py_ssize_t *indices; /* one index per pool */
1608 PyObject *result; /* most recently returned result tuple */
1609 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001610} productobject;
1611
1612static PyTypeObject product_type;
1613
1614static PyObject *
1615product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1616{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001617 productobject *lz;
1618 Py_ssize_t nargs, npools, repeat=1;
1619 PyObject *pools = NULL;
1620 Py_ssize_t *indices = NULL;
1621 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001622
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001623 if (kwds != NULL) {
1624 char *kwlist[] = {"repeat", 0};
1625 PyObject *tmpargs = PyTuple_New(0);
1626 if (tmpargs == NULL)
1627 return NULL;
1628 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1629 Py_DECREF(tmpargs);
1630 return NULL;
1631 }
1632 Py_DECREF(tmpargs);
1633 if (repeat < 0) {
1634 PyErr_SetString(PyExc_ValueError,
1635 "repeat argument cannot be negative");
1636 return NULL;
1637 }
1638 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001639
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001640 assert(PyTuple_Check(args));
1641 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1642 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001643
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001644 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1645 if (indices == NULL) {
1646 PyErr_NoMemory();
1647 goto error;
1648 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001649
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001650 pools = PyTuple_New(npools);
1651 if (pools == NULL)
1652 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001653
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001654 for (i=0; i < nargs ; ++i) {
1655 PyObject *item = PyTuple_GET_ITEM(args, i);
1656 PyObject *pool = PySequence_Tuple(item);
1657 if (pool == NULL)
1658 goto error;
1659 PyTuple_SET_ITEM(pools, i, pool);
1660 indices[i] = 0;
1661 }
1662 for ( ; i < npools; ++i) {
1663 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1664 Py_INCREF(pool);
1665 PyTuple_SET_ITEM(pools, i, pool);
1666 indices[i] = 0;
1667 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001668
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001669 /* create productobject structure */
1670 lz = (productobject *)type->tp_alloc(type, 0);
1671 if (lz == NULL)
1672 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001673
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001674 lz->pools = pools;
1675 lz->indices = indices;
1676 lz->result = NULL;
1677 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001678
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001679 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001680
1681error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001682 if (indices != NULL)
1683 PyMem_Free(indices);
1684 Py_XDECREF(pools);
1685 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001686}
1687
1688static void
1689product_dealloc(productobject *lz)
1690{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001691 PyObject_GC_UnTrack(lz);
1692 Py_XDECREF(lz->pools);
1693 Py_XDECREF(lz->result);
1694 if (lz->indices != NULL)
1695 PyMem_Free(lz->indices);
1696 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001697}
1698
1699static int
1700product_traverse(productobject *lz, visitproc visit, void *arg)
1701{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001702 Py_VISIT(lz->pools);
1703 Py_VISIT(lz->result);
1704 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001705}
1706
1707static PyObject *
1708product_next(productobject *lz)
1709{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001710 PyObject *pool;
1711 PyObject *elem;
1712 PyObject *oldelem;
1713 PyObject *pools = lz->pools;
1714 PyObject *result = lz->result;
1715 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1716 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001717
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001718 if (lz->stopped)
1719 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001720
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001721 if (result == NULL) {
1722 /* On the first pass, return an initial tuple filled with the
1723 first element from each pool. */
1724 result = PyTuple_New(npools);
1725 if (result == NULL)
1726 goto empty;
1727 lz->result = result;
1728 for (i=0; i < npools; i++) {
1729 pool = PyTuple_GET_ITEM(pools, i);
1730 if (PyTuple_GET_SIZE(pool) == 0)
1731 goto empty;
1732 elem = PyTuple_GET_ITEM(pool, 0);
1733 Py_INCREF(elem);
1734 PyTuple_SET_ITEM(result, i, elem);
1735 }
1736 } else {
1737 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001738
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001739 /* Copy the previous result tuple or re-use it if available */
1740 if (Py_REFCNT(result) > 1) {
1741 PyObject *old_result = result;
1742 result = PyTuple_New(npools);
1743 if (result == NULL)
1744 goto empty;
1745 lz->result = result;
1746 for (i=0; i < npools; i++) {
1747 elem = PyTuple_GET_ITEM(old_result, i);
1748 Py_INCREF(elem);
1749 PyTuple_SET_ITEM(result, i, elem);
1750 }
1751 Py_DECREF(old_result);
1752 }
1753 /* Now, we've got the only copy so we can update it in-place */
1754 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001755
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001756 /* Update the pool indices right-to-left. Only advance to the
1757 next pool when the previous one rolls-over */
1758 for (i=npools-1 ; i >= 0 ; i--) {
1759 pool = PyTuple_GET_ITEM(pools, i);
1760 indices[i]++;
1761 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1762 /* Roll-over and advance to next pool */
1763 indices[i] = 0;
1764 elem = PyTuple_GET_ITEM(pool, 0);
1765 Py_INCREF(elem);
1766 oldelem = PyTuple_GET_ITEM(result, i);
1767 PyTuple_SET_ITEM(result, i, elem);
1768 Py_DECREF(oldelem);
1769 } else {
1770 /* No rollover. Just increment and stop here. */
1771 elem = PyTuple_GET_ITEM(pool, indices[i]);
1772 Py_INCREF(elem);
1773 oldelem = PyTuple_GET_ITEM(result, i);
1774 PyTuple_SET_ITEM(result, i, elem);
1775 Py_DECREF(oldelem);
1776 break;
1777 }
1778 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001779
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001780 /* If i is negative, then the indices have all rolled-over
1781 and we're done. */
1782 if (i < 0)
1783 goto empty;
1784 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001785
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001786 Py_INCREF(result);
1787 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001788
1789empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001790 lz->stopped = 1;
1791 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001792}
1793
1794PyDoc_STRVAR(product_doc,
1795"product(*iterables) --> product object\n\
1796\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001797Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001798For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1799The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1800cycle in a manner similar to an odometer (with the rightmost element changing\n\
1801on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00001802To compute the product of an iterable with itself, specify the number\n\
1803of repetitions with the optional repeat keyword argument. For example,\n\
1804product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001805product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1806product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1807
1808static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001809 PyVarObject_HEAD_INIT(NULL, 0)
1810 "itertools.product", /* tp_name */
1811 sizeof(productobject), /* tp_basicsize */
1812 0, /* tp_itemsize */
1813 /* methods */
1814 (destructor)product_dealloc, /* tp_dealloc */
1815 0, /* tp_print */
1816 0, /* tp_getattr */
1817 0, /* tp_setattr */
1818 0, /* tp_reserved */
1819 0, /* tp_repr */
1820 0, /* tp_as_number */
1821 0, /* tp_as_sequence */
1822 0, /* tp_as_mapping */
1823 0, /* tp_hash */
1824 0, /* tp_call */
1825 0, /* tp_str */
1826 PyObject_GenericGetAttr, /* tp_getattro */
1827 0, /* tp_setattro */
1828 0, /* tp_as_buffer */
1829 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1830 Py_TPFLAGS_BASETYPE, /* tp_flags */
1831 product_doc, /* tp_doc */
1832 (traverseproc)product_traverse, /* tp_traverse */
1833 0, /* tp_clear */
1834 0, /* tp_richcompare */
1835 0, /* tp_weaklistoffset */
1836 PyObject_SelfIter, /* tp_iter */
1837 (iternextfunc)product_next, /* tp_iternext */
1838 0, /* tp_methods */
1839 0, /* tp_members */
1840 0, /* tp_getset */
1841 0, /* tp_base */
1842 0, /* tp_dict */
1843 0, /* tp_descr_get */
1844 0, /* tp_descr_set */
1845 0, /* tp_dictoffset */
1846 0, /* tp_init */
1847 0, /* tp_alloc */
1848 product_new, /* tp_new */
1849 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001850};
1851
1852
Christian Heimes380f7f22008-02-28 11:19:05 +00001853/* combinations object ************************************************************/
1854
1855typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001856 PyObject_HEAD
1857 PyObject *pool; /* input converted to a tuple */
1858 Py_ssize_t *indices; /* one index per result element */
1859 PyObject *result; /* most recently returned result tuple */
1860 Py_ssize_t r; /* size of result tuple */
1861 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00001862} combinationsobject;
1863
1864static PyTypeObject combinations_type;
1865
1866static PyObject *
1867combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1868{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001869 combinationsobject *co;
1870 Py_ssize_t n;
1871 Py_ssize_t r;
1872 PyObject *pool = NULL;
1873 PyObject *iterable = NULL;
1874 Py_ssize_t *indices = NULL;
1875 Py_ssize_t i;
1876 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00001877
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001878 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
1879 &iterable, &r))
1880 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001881
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001882 pool = PySequence_Tuple(iterable);
1883 if (pool == NULL)
1884 goto error;
1885 n = PyTuple_GET_SIZE(pool);
1886 if (r < 0) {
1887 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
1888 goto error;
1889 }
Christian Heimes380f7f22008-02-28 11:19:05 +00001890
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001891 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
1892 if (indices == NULL) {
1893 PyErr_NoMemory();
1894 goto error;
1895 }
Christian Heimes380f7f22008-02-28 11:19:05 +00001896
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001897 for (i=0 ; i<r ; i++)
1898 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00001899
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001900 /* create combinationsobject structure */
1901 co = (combinationsobject *)type->tp_alloc(type, 0);
1902 if (co == NULL)
1903 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00001904
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001905 co->pool = pool;
1906 co->indices = indices;
1907 co->result = NULL;
1908 co->r = r;
1909 co->stopped = r > n ? 1 : 0;
1910
1911 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00001912
1913error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001914 if (indices != NULL)
1915 PyMem_Free(indices);
1916 Py_XDECREF(pool);
1917 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001918}
1919
1920static void
1921combinations_dealloc(combinationsobject *co)
1922{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001923 PyObject_GC_UnTrack(co);
1924 Py_XDECREF(co->pool);
1925 Py_XDECREF(co->result);
1926 if (co->indices != NULL)
1927 PyMem_Free(co->indices);
1928 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00001929}
1930
1931static int
1932combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
1933{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001934 Py_VISIT(co->pool);
1935 Py_VISIT(co->result);
1936 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00001937}
1938
1939static PyObject *
1940combinations_next(combinationsobject *co)
1941{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001942 PyObject *elem;
1943 PyObject *oldelem;
1944 PyObject *pool = co->pool;
1945 Py_ssize_t *indices = co->indices;
1946 PyObject *result = co->result;
1947 Py_ssize_t n = PyTuple_GET_SIZE(pool);
1948 Py_ssize_t r = co->r;
1949 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00001950
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001951 if (co->stopped)
1952 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001953
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001954 if (result == NULL) {
1955 /* On the first pass, initialize result tuple using the indices */
1956 result = PyTuple_New(r);
1957 if (result == NULL)
1958 goto empty;
1959 co->result = result;
1960 for (i=0; i<r ; i++) {
1961 index = indices[i];
1962 elem = PyTuple_GET_ITEM(pool, index);
1963 Py_INCREF(elem);
1964 PyTuple_SET_ITEM(result, i, elem);
1965 }
1966 } else {
1967 /* Copy the previous result tuple or re-use it if available */
1968 if (Py_REFCNT(result) > 1) {
1969 PyObject *old_result = result;
1970 result = PyTuple_New(r);
1971 if (result == NULL)
1972 goto empty;
1973 co->result = result;
1974 for (i=0; i<r ; i++) {
1975 elem = PyTuple_GET_ITEM(old_result, i);
1976 Py_INCREF(elem);
1977 PyTuple_SET_ITEM(result, i, elem);
1978 }
1979 Py_DECREF(old_result);
1980 }
1981 /* Now, we've got the only copy so we can update it in-place
1982 * CPython's empty tuple is a singleton and cached in
1983 * PyTuple's freelist.
1984 */
1985 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00001986
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001987 /* Scan indices right-to-left until finding one that is not
1988 at its maximum (i + n - r). */
1989 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
1990 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00001991
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001992 /* If i is negative, then the indices are all at
1993 their maximum value and we're done. */
1994 if (i < 0)
1995 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00001996
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001997 /* Increment the current index which we know is not at its
1998 maximum. Then move back to the right setting each index
1999 to its lowest possible value (one higher than the index
2000 to its left -- this maintains the sort order invariant). */
2001 indices[i]++;
2002 for (j=i+1 ; j<r ; j++)
2003 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002005 /* Update the result tuple for the new indices
2006 starting with i, the leftmost index that changed */
2007 for ( ; i<r ; i++) {
2008 index = indices[i];
2009 elem = PyTuple_GET_ITEM(pool, index);
2010 Py_INCREF(elem);
2011 oldelem = PyTuple_GET_ITEM(result, i);
2012 PyTuple_SET_ITEM(result, i, elem);
2013 Py_DECREF(oldelem);
2014 }
2015 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002016
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002017 Py_INCREF(result);
2018 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002019
2020empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002021 co->stopped = 1;
2022 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002023}
2024
2025PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002026"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002027\n\
2028Return successive r-length combinations of elements in the iterable.\n\n\
2029combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2030
2031static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002032 PyVarObject_HEAD_INIT(NULL, 0)
2033 "itertools.combinations", /* tp_name */
2034 sizeof(combinationsobject), /* tp_basicsize */
2035 0, /* tp_itemsize */
2036 /* methods */
2037 (destructor)combinations_dealloc, /* tp_dealloc */
2038 0, /* tp_print */
2039 0, /* tp_getattr */
2040 0, /* tp_setattr */
2041 0, /* tp_reserved */
2042 0, /* tp_repr */
2043 0, /* tp_as_number */
2044 0, /* tp_as_sequence */
2045 0, /* tp_as_mapping */
2046 0, /* tp_hash */
2047 0, /* tp_call */
2048 0, /* tp_str */
2049 PyObject_GenericGetAttr, /* tp_getattro */
2050 0, /* tp_setattro */
2051 0, /* tp_as_buffer */
2052 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2053 Py_TPFLAGS_BASETYPE, /* tp_flags */
2054 combinations_doc, /* tp_doc */
2055 (traverseproc)combinations_traverse, /* tp_traverse */
2056 0, /* tp_clear */
2057 0, /* tp_richcompare */
2058 0, /* tp_weaklistoffset */
2059 PyObject_SelfIter, /* tp_iter */
2060 (iternextfunc)combinations_next, /* tp_iternext */
2061 0, /* tp_methods */
2062 0, /* tp_members */
2063 0, /* tp_getset */
2064 0, /* tp_base */
2065 0, /* tp_dict */
2066 0, /* tp_descr_get */
2067 0, /* tp_descr_set */
2068 0, /* tp_dictoffset */
2069 0, /* tp_init */
2070 0, /* tp_alloc */
2071 combinations_new, /* tp_new */
2072 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002073};
2074
2075
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002076/* combinations with replacement object *******************************************/
2077
2078/* Equivalent to:
2079
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002080 def combinations_with_replacement(iterable, r):
2081 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2082 # number items returned: (n+r-1)! / r! / (n-1)!
2083 pool = tuple(iterable)
2084 n = len(pool)
2085 indices = [0] * r
2086 yield tuple(pool[i] for i in indices)
2087 while 1:
2088 for i in reversed(range(r)):
2089 if indices[i] != n - 1:
2090 break
2091 else:
2092 return
2093 indices[i:] = [indices[i] + 1] * (r - i)
2094 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002095
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002096 def combinations_with_replacement2(iterable, r):
2097 'Alternate version that filters from product()'
2098 pool = tuple(iterable)
2099 n = len(pool)
2100 for indices in product(range(n), repeat=r):
2101 if sorted(indices) == list(indices):
2102 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002103*/
2104typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002105 PyObject_HEAD
2106 PyObject *pool; /* input converted to a tuple */
2107 Py_ssize_t *indices; /* one index per result element */
2108 PyObject *result; /* most recently returned result tuple */
2109 Py_ssize_t r; /* size of result tuple */
2110 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002111} cwrobject;
2112
2113static PyTypeObject cwr_type;
2114
2115static PyObject *
2116cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2117{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002118 cwrobject *co;
2119 Py_ssize_t n;
2120 Py_ssize_t r;
2121 PyObject *pool = NULL;
2122 PyObject *iterable = NULL;
2123 Py_ssize_t *indices = NULL;
2124 Py_ssize_t i;
2125 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002126
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002127 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2128 &iterable, &r))
2129 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002130
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002131 pool = PySequence_Tuple(iterable);
2132 if (pool == NULL)
2133 goto error;
2134 n = PyTuple_GET_SIZE(pool);
2135 if (r < 0) {
2136 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2137 goto error;
2138 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002139
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002140 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2141 if (indices == NULL) {
2142 PyErr_NoMemory();
2143 goto error;
2144 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002145
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002146 for (i=0 ; i<r ; i++)
2147 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002148
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002149 /* create cwrobject structure */
2150 co = (cwrobject *)type->tp_alloc(type, 0);
2151 if (co == NULL)
2152 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002153
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002154 co->pool = pool;
2155 co->indices = indices;
2156 co->result = NULL;
2157 co->r = r;
2158 co->stopped = !n && r;
2159
2160 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002161
2162error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002163 if (indices != NULL)
2164 PyMem_Free(indices);
2165 Py_XDECREF(pool);
2166 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002167}
2168
2169static void
2170cwr_dealloc(cwrobject *co)
2171{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002172 PyObject_GC_UnTrack(co);
2173 Py_XDECREF(co->pool);
2174 Py_XDECREF(co->result);
2175 if (co->indices != NULL)
2176 PyMem_Free(co->indices);
2177 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002178}
2179
2180static int
2181cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2182{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002183 Py_VISIT(co->pool);
2184 Py_VISIT(co->result);
2185 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002186}
2187
2188static PyObject *
2189cwr_next(cwrobject *co)
2190{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002191 PyObject *elem;
2192 PyObject *oldelem;
2193 PyObject *pool = co->pool;
2194 Py_ssize_t *indices = co->indices;
2195 PyObject *result = co->result;
2196 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2197 Py_ssize_t r = co->r;
2198 Py_ssize_t i, j, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002199
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002200 if (co->stopped)
2201 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002202
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002203 if (result == NULL) {
2204 /* On the first pass, initialize result tuple using the indices */
2205 result = PyTuple_New(r);
2206 if (result == NULL)
2207 goto empty;
2208 co->result = result;
2209 for (i=0; i<r ; i++) {
2210 index = indices[i];
2211 elem = PyTuple_GET_ITEM(pool, index);
2212 Py_INCREF(elem);
2213 PyTuple_SET_ITEM(result, i, elem);
2214 }
2215 } else {
2216 /* Copy the previous result tuple or re-use it if available */
2217 if (Py_REFCNT(result) > 1) {
2218 PyObject *old_result = result;
2219 result = PyTuple_New(r);
2220 if (result == NULL)
2221 goto empty;
2222 co->result = result;
2223 for (i=0; i<r ; i++) {
2224 elem = PyTuple_GET_ITEM(old_result, i);
2225 Py_INCREF(elem);
2226 PyTuple_SET_ITEM(result, i, elem);
2227 }
2228 Py_DECREF(old_result);
2229 }
2230 /* Now, we've got the only copy so we can update it in-place CPython's
2231 empty tuple is a singleton and cached in PyTuple's freelist. */
2232 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002233
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002234 /* Scan indices right-to-left until finding one that is not
2235 * at its maximum (n-1). */
2236 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2237 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002238
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002239 /* If i is negative, then the indices are all at
2240 their maximum value and we're done. */
2241 if (i < 0)
2242 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002243
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002244 /* Increment the current index which we know is not at its
2245 maximum. Then set all to the right to the same value. */
2246 indices[i]++;
2247 for (j=i+1 ; j<r ; j++)
2248 indices[j] = indices[j-1];
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002249
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002250 /* Update the result tuple for the new indices
2251 starting with i, the leftmost index that changed */
2252 for ( ; i<r ; i++) {
2253 index = indices[i];
2254 elem = PyTuple_GET_ITEM(pool, index);
2255 Py_INCREF(elem);
2256 oldelem = PyTuple_GET_ITEM(result, i);
2257 PyTuple_SET_ITEM(result, i, elem);
2258 Py_DECREF(oldelem);
2259 }
2260 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002261
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002262 Py_INCREF(result);
2263 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002264
2265empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002266 co->stopped = 1;
2267 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002268}
2269
2270PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002271"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002272\n\
2273Return successive r-length combinations of elements in the iterable\n\
2274allowing individual elements to have successive repeats.\n\
2275combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2276
2277static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002278 PyVarObject_HEAD_INIT(NULL, 0)
2279 "itertools.combinations_with_replacement", /* tp_name */
2280 sizeof(cwrobject), /* tp_basicsize */
2281 0, /* tp_itemsize */
2282 /* methods */
2283 (destructor)cwr_dealloc, /* tp_dealloc */
2284 0, /* tp_print */
2285 0, /* tp_getattr */
2286 0, /* tp_setattr */
2287 0, /* tp_reserved */
2288 0, /* tp_repr */
2289 0, /* tp_as_number */
2290 0, /* tp_as_sequence */
2291 0, /* tp_as_mapping */
2292 0, /* tp_hash */
2293 0, /* tp_call */
2294 0, /* tp_str */
2295 PyObject_GenericGetAttr, /* tp_getattro */
2296 0, /* tp_setattro */
2297 0, /* tp_as_buffer */
2298 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2299 Py_TPFLAGS_BASETYPE, /* tp_flags */
2300 cwr_doc, /* tp_doc */
2301 (traverseproc)cwr_traverse, /* tp_traverse */
2302 0, /* tp_clear */
2303 0, /* tp_richcompare */
2304 0, /* tp_weaklistoffset */
2305 PyObject_SelfIter, /* tp_iter */
2306 (iternextfunc)cwr_next, /* tp_iternext */
2307 0, /* tp_methods */
2308 0, /* tp_members */
2309 0, /* tp_getset */
2310 0, /* tp_base */
2311 0, /* tp_dict */
2312 0, /* tp_descr_get */
2313 0, /* tp_descr_set */
2314 0, /* tp_dictoffset */
2315 0, /* tp_init */
2316 0, /* tp_alloc */
2317 cwr_new, /* tp_new */
2318 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002319};
2320
2321
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002322/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002323
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002324def permutations(iterable, r=None):
2325 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2326 pool = tuple(iterable)
2327 n = len(pool)
2328 r = n if r is None else r
2329 indices = range(n)
2330 cycles = range(n-r+1, n+1)[::-1]
2331 yield tuple(pool[i] for i in indices[:r])
2332 while n:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002333 for i in reversed(range(r)):
2334 cycles[i] -= 1
2335 if cycles[i] == 0:
2336 indices[i:] = indices[i+1:] + indices[i:i+1]
2337 cycles[i] = n - i
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002338 else:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002339 j = cycles[i]
2340 indices[i], indices[-j] = indices[-j], indices[i]
2341 yield tuple(pool[i] for i in indices[:r])
2342 break
2343 else:
2344 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002345*/
2346
2347typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002348 PyObject_HEAD
2349 PyObject *pool; /* input converted to a tuple */
2350 Py_ssize_t *indices; /* one index per element in the pool */
2351 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2352 PyObject *result; /* most recently returned result tuple */
2353 Py_ssize_t r; /* size of result tuple */
2354 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002355} permutationsobject;
2356
2357static PyTypeObject permutations_type;
2358
2359static PyObject *
2360permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2361{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002362 permutationsobject *po;
2363 Py_ssize_t n;
2364 Py_ssize_t r;
2365 PyObject *robj = Py_None;
2366 PyObject *pool = NULL;
2367 PyObject *iterable = NULL;
2368 Py_ssize_t *indices = NULL;
2369 Py_ssize_t *cycles = NULL;
2370 Py_ssize_t i;
2371 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002372
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002373 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2374 &iterable, &robj))
2375 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002376
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002377 pool = PySequence_Tuple(iterable);
2378 if (pool == NULL)
2379 goto error;
2380 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002381
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002382 r = n;
2383 if (robj != Py_None) {
2384 if (!PyLong_Check(robj)) {
2385 PyErr_SetString(PyExc_TypeError, "Expected int as r");
2386 goto error;
2387 }
2388 r = PyLong_AsSsize_t(robj);
2389 if (r == -1 && PyErr_Occurred())
2390 goto error;
2391 }
2392 if (r < 0) {
2393 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2394 goto error;
2395 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002396
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002397 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2398 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2399 if (indices == NULL || cycles == NULL) {
2400 PyErr_NoMemory();
2401 goto error;
2402 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002403
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002404 for (i=0 ; i<n ; i++)
2405 indices[i] = i;
2406 for (i=0 ; i<r ; i++)
2407 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002408
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002409 /* create permutationsobject structure */
2410 po = (permutationsobject *)type->tp_alloc(type, 0);
2411 if (po == NULL)
2412 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002413
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002414 po->pool = pool;
2415 po->indices = indices;
2416 po->cycles = cycles;
2417 po->result = NULL;
2418 po->r = r;
2419 po->stopped = r > n ? 1 : 0;
2420
2421 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002422
2423error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002424 if (indices != NULL)
2425 PyMem_Free(indices);
2426 if (cycles != NULL)
2427 PyMem_Free(cycles);
2428 Py_XDECREF(pool);
2429 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002430}
2431
2432static void
2433permutations_dealloc(permutationsobject *po)
2434{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002435 PyObject_GC_UnTrack(po);
2436 Py_XDECREF(po->pool);
2437 Py_XDECREF(po->result);
2438 PyMem_Free(po->indices);
2439 PyMem_Free(po->cycles);
2440 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002441}
2442
2443static int
2444permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2445{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002446 Py_VISIT(po->pool);
2447 Py_VISIT(po->result);
2448 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002449}
2450
2451static PyObject *
2452permutations_next(permutationsobject *po)
2453{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002454 PyObject *elem;
2455 PyObject *oldelem;
2456 PyObject *pool = po->pool;
2457 Py_ssize_t *indices = po->indices;
2458 Py_ssize_t *cycles = po->cycles;
2459 PyObject *result = po->result;
2460 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2461 Py_ssize_t r = po->r;
2462 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002463
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002464 if (po->stopped)
2465 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002466
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002467 if (result == NULL) {
2468 /* On the first pass, initialize result tuple using the indices */
2469 result = PyTuple_New(r);
2470 if (result == NULL)
2471 goto empty;
2472 po->result = result;
2473 for (i=0; i<r ; i++) {
2474 index = indices[i];
2475 elem = PyTuple_GET_ITEM(pool, index);
2476 Py_INCREF(elem);
2477 PyTuple_SET_ITEM(result, i, elem);
2478 }
2479 } else {
2480 if (n == 0)
2481 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002482
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002483 /* Copy the previous result tuple or re-use it if available */
2484 if (Py_REFCNT(result) > 1) {
2485 PyObject *old_result = result;
2486 result = PyTuple_New(r);
2487 if (result == NULL)
2488 goto empty;
2489 po->result = result;
2490 for (i=0; i<r ; i++) {
2491 elem = PyTuple_GET_ITEM(old_result, i);
2492 Py_INCREF(elem);
2493 PyTuple_SET_ITEM(result, i, elem);
2494 }
2495 Py_DECREF(old_result);
2496 }
2497 /* Now, we've got the only copy so we can update it in-place */
2498 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002499
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002500 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2501 for (i=r-1 ; i>=0 ; i--) {
2502 cycles[i] -= 1;
2503 if (cycles[i] == 0) {
2504 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2505 index = indices[i];
2506 for (j=i ; j<n-1 ; j++)
2507 indices[j] = indices[j+1];
2508 indices[n-1] = index;
2509 cycles[i] = n - i;
2510 } else {
2511 j = cycles[i];
2512 index = indices[i];
2513 indices[i] = indices[n-j];
2514 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002515
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002516 for (k=i; k<r ; k++) {
2517 /* start with i, the leftmost element that changed */
2518 /* yield tuple(pool[k] for k in indices[:r]) */
2519 index = indices[k];
2520 elem = PyTuple_GET_ITEM(pool, index);
2521 Py_INCREF(elem);
2522 oldelem = PyTuple_GET_ITEM(result, k);
2523 PyTuple_SET_ITEM(result, k, elem);
2524 Py_DECREF(oldelem);
2525 }
2526 break;
2527 }
2528 }
2529 /* If i is negative, then the cycles have all
2530 rolled-over and we're done. */
2531 if (i < 0)
2532 goto empty;
2533 }
2534 Py_INCREF(result);
2535 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002536
2537empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002538 po->stopped = 1;
2539 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002540}
2541
2542PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00002543"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002544\n\
2545Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00002546permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002547
2548static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002549 PyVarObject_HEAD_INIT(NULL, 0)
2550 "itertools.permutations", /* tp_name */
2551 sizeof(permutationsobject), /* tp_basicsize */
2552 0, /* tp_itemsize */
2553 /* methods */
2554 (destructor)permutations_dealloc, /* tp_dealloc */
2555 0, /* tp_print */
2556 0, /* tp_getattr */
2557 0, /* tp_setattr */
2558 0, /* tp_reserved */
2559 0, /* tp_repr */
2560 0, /* tp_as_number */
2561 0, /* tp_as_sequence */
2562 0, /* tp_as_mapping */
2563 0, /* tp_hash */
2564 0, /* tp_call */
2565 0, /* tp_str */
2566 PyObject_GenericGetAttr, /* tp_getattro */
2567 0, /* tp_setattro */
2568 0, /* tp_as_buffer */
2569 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2570 Py_TPFLAGS_BASETYPE, /* tp_flags */
2571 permutations_doc, /* tp_doc */
2572 (traverseproc)permutations_traverse, /* tp_traverse */
2573 0, /* tp_clear */
2574 0, /* tp_richcompare */
2575 0, /* tp_weaklistoffset */
2576 PyObject_SelfIter, /* tp_iter */
2577 (iternextfunc)permutations_next, /* tp_iternext */
2578 0, /* tp_methods */
2579 0, /* tp_members */
2580 0, /* tp_getset */
2581 0, /* tp_base */
2582 0, /* tp_dict */
2583 0, /* tp_descr_get */
2584 0, /* tp_descr_set */
2585 0, /* tp_dictoffset */
2586 0, /* tp_init */
2587 0, /* tp_alloc */
2588 permutations_new, /* tp_new */
2589 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002590};
2591
Raymond Hettinger482ba772010-12-01 22:48:00 +00002592/* accumulate object ************************************************************/
2593
2594typedef struct {
2595 PyObject_HEAD
2596 PyObject *total;
2597 PyObject *it;
2598} accumulateobject;
2599
2600static PyTypeObject accumulate_type;
2601
2602static PyObject *
2603accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2604{
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002605 static char *kwargs[] = {"iterable", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00002606 PyObject *iterable;
2607 PyObject *it;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002608 accumulateobject *lz;
2609
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002610 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O:accumulate", kwargs, &iterable))
2611 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002612
2613 /* Get iterator. */
2614 it = PyObject_GetIter(iterable);
2615 if (it == NULL)
2616 return NULL;
2617
Raymond Hettinger482ba772010-12-01 22:48:00 +00002618 /* create accumulateobject structure */
2619 lz = (accumulateobject *)type->tp_alloc(type, 0);
2620 if (lz == NULL) {
2621 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002622 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002623 }
2624
Raymond 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);
2634 Py_XDECREF(lz->total);
2635 Py_XDECREF(lz->it);
2636 Py_TYPE(lz)->tp_free(lz);
2637}
2638
2639static int
2640accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
2641{
2642 Py_VISIT(lz->it);
2643 Py_VISIT(lz->total);
2644 return 0;
2645}
2646
2647static PyObject *
2648accumulate_next(accumulateobject *lz)
2649{
2650 PyObject *val, *oldtotal, *newtotal;
2651
2652 val = PyIter_Next(lz->it);
2653 if (val == NULL)
2654 return NULL;
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002655
2656 if (lz->total == NULL) {
2657 Py_INCREF(val);
2658 lz->total = val;
2659 return lz->total;
2660 }
2661
Raymond Hettinger482ba772010-12-01 22:48:00 +00002662 newtotal = PyNumber_Add(lz->total, val);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002663 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00002664 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002665 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002666
2667 oldtotal = lz->total;
2668 lz->total = newtotal;
2669 Py_DECREF(oldtotal);
2670
2671 Py_INCREF(newtotal);
2672 return newtotal;
2673}
2674
2675PyDoc_STRVAR(accumulate_doc,
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002676"accumulate(iterable) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00002677\n\
2678Return series of accumulated sums.");
2679
2680static PyTypeObject accumulate_type = {
2681 PyVarObject_HEAD_INIT(NULL, 0)
2682 "itertools.accumulate", /* tp_name */
2683 sizeof(accumulateobject), /* tp_basicsize */
2684 0, /* tp_itemsize */
2685 /* methods */
2686 (destructor)accumulate_dealloc, /* tp_dealloc */
2687 0, /* tp_print */
2688 0, /* tp_getattr */
2689 0, /* tp_setattr */
2690 0, /* tp_reserved */
2691 0, /* tp_repr */
2692 0, /* tp_as_number */
2693 0, /* tp_as_sequence */
2694 0, /* tp_as_mapping */
2695 0, /* tp_hash */
2696 0, /* tp_call */
2697 0, /* tp_str */
2698 PyObject_GenericGetAttr, /* tp_getattro */
2699 0, /* tp_setattro */
2700 0, /* tp_as_buffer */
2701 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2702 Py_TPFLAGS_BASETYPE, /* tp_flags */
2703 accumulate_doc, /* tp_doc */
2704 (traverseproc)accumulate_traverse, /* tp_traverse */
2705 0, /* tp_clear */
2706 0, /* tp_richcompare */
2707 0, /* tp_weaklistoffset */
2708 PyObject_SelfIter, /* tp_iter */
2709 (iternextfunc)accumulate_next, /* tp_iternext */
2710 0, /* tp_methods */
2711 0, /* tp_members */
2712 0, /* tp_getset */
2713 0, /* tp_base */
2714 0, /* tp_dict */
2715 0, /* tp_descr_get */
2716 0, /* tp_descr_set */
2717 0, /* tp_dictoffset */
2718 0, /* tp_init */
2719 0, /* tp_alloc */
2720 accumulate_new, /* tp_new */
2721 PyObject_GC_Del, /* tp_free */
2722};
2723
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002724
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002725/* compress object ************************************************************/
2726
2727/* Equivalent to:
2728
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002729 def compress(data, selectors):
2730 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2731 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002732*/
2733
2734typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002735 PyObject_HEAD
2736 PyObject *data;
2737 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002738} compressobject;
2739
2740static PyTypeObject compress_type;
2741
2742static PyObject *
2743compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2744{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002745 PyObject *seq1, *seq2;
2746 PyObject *data=NULL, *selectors=NULL;
2747 compressobject *lz;
2748 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002749
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002750 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2751 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002752
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002753 data = PyObject_GetIter(seq1);
2754 if (data == NULL)
2755 goto fail;
2756 selectors = PyObject_GetIter(seq2);
2757 if (selectors == NULL)
2758 goto fail;
2759
2760 /* create compressobject structure */
2761 lz = (compressobject *)type->tp_alloc(type, 0);
2762 if (lz == NULL)
2763 goto fail;
2764 lz->data = data;
2765 lz->selectors = selectors;
2766 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002767
2768fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002769 Py_XDECREF(data);
2770 Py_XDECREF(selectors);
2771 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002772}
2773
2774static void
2775compress_dealloc(compressobject *lz)
2776{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002777 PyObject_GC_UnTrack(lz);
2778 Py_XDECREF(lz->data);
2779 Py_XDECREF(lz->selectors);
2780 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002781}
2782
2783static int
2784compress_traverse(compressobject *lz, visitproc visit, void *arg)
2785{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002786 Py_VISIT(lz->data);
2787 Py_VISIT(lz->selectors);
2788 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002789}
2790
2791static PyObject *
2792compress_next(compressobject *lz)
2793{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002794 PyObject *data = lz->data, *selectors = lz->selectors;
2795 PyObject *datum, *selector;
2796 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2797 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2798 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002799
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002800 while (1) {
2801 /* Steps: get datum, get selector, evaluate selector.
2802 Order is important (to match the pure python version
2803 in terms of which input gets a chance to raise an
2804 exception first).
2805 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002806
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002807 datum = datanext(data);
2808 if (datum == NULL)
2809 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002810
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002811 selector = selectornext(selectors);
2812 if (selector == NULL) {
2813 Py_DECREF(datum);
2814 return NULL;
2815 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002816
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002817 ok = PyObject_IsTrue(selector);
2818 Py_DECREF(selector);
2819 if (ok == 1)
2820 return datum;
2821 Py_DECREF(datum);
2822 if (ok == -1)
2823 return NULL;
2824 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002825}
2826
2827PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00002828"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002829\n\
2830Return data elements corresponding to true selector elements.\n\
2831Forms a shorter iterator from selected data elements using the\n\
2832selectors to choose the data elements.");
2833
2834static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002835 PyVarObject_HEAD_INIT(NULL, 0)
2836 "itertools.compress", /* tp_name */
2837 sizeof(compressobject), /* tp_basicsize */
2838 0, /* tp_itemsize */
2839 /* methods */
2840 (destructor)compress_dealloc, /* tp_dealloc */
2841 0, /* tp_print */
2842 0, /* tp_getattr */
2843 0, /* tp_setattr */
2844 0, /* tp_reserved */
2845 0, /* tp_repr */
2846 0, /* tp_as_number */
2847 0, /* tp_as_sequence */
2848 0, /* tp_as_mapping */
2849 0, /* tp_hash */
2850 0, /* tp_call */
2851 0, /* tp_str */
2852 PyObject_GenericGetAttr, /* tp_getattro */
2853 0, /* tp_setattro */
2854 0, /* tp_as_buffer */
2855 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2856 Py_TPFLAGS_BASETYPE, /* tp_flags */
2857 compress_doc, /* tp_doc */
2858 (traverseproc)compress_traverse, /* tp_traverse */
2859 0, /* tp_clear */
2860 0, /* tp_richcompare */
2861 0, /* tp_weaklistoffset */
2862 PyObject_SelfIter, /* tp_iter */
2863 (iternextfunc)compress_next, /* tp_iternext */
2864 0, /* tp_methods */
2865 0, /* tp_members */
2866 0, /* tp_getset */
2867 0, /* tp_base */
2868 0, /* tp_dict */
2869 0, /* tp_descr_get */
2870 0, /* tp_descr_set */
2871 0, /* tp_dictoffset */
2872 0, /* tp_init */
2873 0, /* tp_alloc */
2874 compress_new, /* tp_new */
2875 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002876};
2877
2878
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002879/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00002880
2881typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002882 PyObject_HEAD
2883 PyObject *func;
2884 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002885} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002886
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002887static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002888
2889static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002890filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002891{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002892 PyObject *func, *seq;
2893 PyObject *it;
2894 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002895
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002896 if (type == &filterfalse_type &&
2897 !_PyArg_NoKeywords("filterfalse()", kwds))
2898 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002899
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002900 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
2901 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002902
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002903 /* Get iterator. */
2904 it = PyObject_GetIter(seq);
2905 if (it == NULL)
2906 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002907
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002908 /* create filterfalseobject structure */
2909 lz = (filterfalseobject *)type->tp_alloc(type, 0);
2910 if (lz == NULL) {
2911 Py_DECREF(it);
2912 return NULL;
2913 }
2914 Py_INCREF(func);
2915 lz->func = func;
2916 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002917
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002918 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002919}
2920
2921static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002922filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002923{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002924 PyObject_GC_UnTrack(lz);
2925 Py_XDECREF(lz->func);
2926 Py_XDECREF(lz->it);
2927 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002928}
2929
2930static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002931filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002932{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002933 Py_VISIT(lz->it);
2934 Py_VISIT(lz->func);
2935 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002936}
2937
2938static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002939filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002940{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002941 PyObject *item;
2942 PyObject *it = lz->it;
2943 long ok;
2944 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002945
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002946 iternext = *Py_TYPE(it)->tp_iternext;
2947 for (;;) {
2948 item = iternext(it);
2949 if (item == NULL)
2950 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002951
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002952 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2953 ok = PyObject_IsTrue(item);
2954 } else {
2955 PyObject *good;
2956 good = PyObject_CallFunctionObjArgs(lz->func,
2957 item, NULL);
2958 if (good == NULL) {
2959 Py_DECREF(item);
2960 return NULL;
2961 }
2962 ok = PyObject_IsTrue(good);
2963 Py_DECREF(good);
2964 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02002965 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002966 return item;
2967 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02002968 if (ok < 0)
2969 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002970 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00002971}
2972
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002973PyDoc_STRVAR(filterfalse_doc,
2974"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002975\n\
2976Return those items of sequence for which function(item) is false.\n\
2977If function is None, return the items that are false.");
2978
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002979static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002980 PyVarObject_HEAD_INIT(NULL, 0)
2981 "itertools.filterfalse", /* tp_name */
2982 sizeof(filterfalseobject), /* tp_basicsize */
2983 0, /* tp_itemsize */
2984 /* methods */
2985 (destructor)filterfalse_dealloc, /* tp_dealloc */
2986 0, /* tp_print */
2987 0, /* tp_getattr */
2988 0, /* tp_setattr */
2989 0, /* tp_reserved */
2990 0, /* tp_repr */
2991 0, /* tp_as_number */
2992 0, /* tp_as_sequence */
2993 0, /* tp_as_mapping */
2994 0, /* tp_hash */
2995 0, /* tp_call */
2996 0, /* tp_str */
2997 PyObject_GenericGetAttr, /* tp_getattro */
2998 0, /* tp_setattro */
2999 0, /* tp_as_buffer */
3000 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3001 Py_TPFLAGS_BASETYPE, /* tp_flags */
3002 filterfalse_doc, /* tp_doc */
3003 (traverseproc)filterfalse_traverse, /* tp_traverse */
3004 0, /* tp_clear */
3005 0, /* tp_richcompare */
3006 0, /* tp_weaklistoffset */
3007 PyObject_SelfIter, /* tp_iter */
3008 (iternextfunc)filterfalse_next, /* tp_iternext */
3009 0, /* tp_methods */
3010 0, /* tp_members */
3011 0, /* tp_getset */
3012 0, /* tp_base */
3013 0, /* tp_dict */
3014 0, /* tp_descr_get */
3015 0, /* tp_descr_set */
3016 0, /* tp_dictoffset */
3017 0, /* tp_init */
3018 0, /* tp_alloc */
3019 filterfalse_new, /* tp_new */
3020 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003021};
3022
3023
3024/* count object ************************************************************/
3025
3026typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003027 PyObject_HEAD
3028 Py_ssize_t cnt;
3029 PyObject *long_cnt;
3030 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003031} countobject;
3032
Raymond Hettinger30729212009-02-12 06:28:27 +00003033/* Counting logic and invariants:
3034
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003035fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003036
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003037 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3038 Advances with: cnt += 1
3039 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003040
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003041slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003042
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003043 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3044 All counting is done with python objects (no overflows or underflows).
3045 Advances with: long_cnt += long_step
3046 Step may be zero -- effectively a slow version of repeat(cnt).
3047 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003048*/
3049
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003050static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003051
3052static PyObject *
3053count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3054{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003055 countobject *lz;
3056 int slow_mode = 0;
3057 Py_ssize_t cnt = 0;
3058 PyObject *long_cnt = NULL;
3059 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003060 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003061 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003062
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003063 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3064 kwlist, &long_cnt, &long_step))
3065 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003066
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003067 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3068 (long_step != NULL && !PyNumber_Check(long_step))) {
3069 PyErr_SetString(PyExc_TypeError, "a number is required");
3070 return NULL;
3071 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003072
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003073 if (long_cnt != NULL) {
3074 cnt = PyLong_AsSsize_t(long_cnt);
3075 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3076 PyErr_Clear();
3077 slow_mode = 1;
3078 }
3079 Py_INCREF(long_cnt);
3080 } else {
3081 cnt = 0;
3082 long_cnt = PyLong_FromLong(0);
3083 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003084
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003085 /* If not specified, step defaults to 1 */
3086 if (long_step == NULL) {
3087 long_step = PyLong_FromLong(1);
3088 if (long_step == NULL) {
3089 Py_DECREF(long_cnt);
3090 return NULL;
3091 }
3092 } else
3093 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003094
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003095 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003096
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003097 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003098 step = PyLong_AsLong(long_step);
3099 if (step != 1) {
3100 slow_mode = 1;
3101 if (step == -1 && PyErr_Occurred())
3102 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003103 }
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}