blob: ba8d6dfde3d218b6c6af930c81b567b416abb9b1 [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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +000098 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000101
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000114 tmp = gbo->currkey;
115 gbo->currkey = newkey;
116 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000117
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000118 tmp = gbo->currvalue;
119 gbo->currvalue = newvalue;
120 Py_XDECREF(tmp);
121 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000131
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000199 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000200
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000233 groupbyobject *gbo = (groupbyobject *)igo->parent;
234 PyObject *newvalue, *newkey, *r;
235 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000236
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000254 assert(gbo->currkey == NULL);
255 gbo->currkey = newkey;
256 gbo->currvalue = newvalue;
257 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000265 r = gbo->currvalue;
266 gbo->currvalue = NULL;
267 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000268
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000269 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000270}
271
272static PyTypeObject _grouper_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000349 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000350
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000351 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
352 if (tdo == NULL)
353 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000354
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000375 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000376
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000474 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000475
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000499 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000500
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000517 teeobject *to;
518 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000519
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000538 to->index = 0;
539 to->weakreflist = NULL;
540 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000541done:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000549 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000550
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000627 Py_ssize_t i, n=2;
628 PyObject *it, *iterable, *copyable, *result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000629
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000685 PyObject *it;
686 PyObject *iterable;
687 PyObject *saved;
688 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000689
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000690 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
691 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000692
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000693 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
694 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000695
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000696 /* Get iterator. */
697 it = PyObject_GetIter(iterable);
698 if (it == NULL)
699 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000700
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000718 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000719}
720
721static void
722cycle_dealloc(cycleobject *lz)
723{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000741 PyObject *item;
742 PyObject *it;
743 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000744
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000837 PyObject *func, *seq;
838 PyObject *it;
839 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000840
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000841 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
842 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000843
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000844 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
845 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000846
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000847 /* Get iterator. */
848 it = PyObject_GetIter(seq);
849 if (it == NULL)
850 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000851
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000863 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000864}
865
866static void
867dropwhile_dealloc(dropwhileobject *lz)
868{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +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 Pitrou7f14f0d2010-05-09 16:14:21 +0000899 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
900 if (good == NULL) {
901 Py_DECREF(item);
902 return NULL;
903 }
904 ok = PyObject_IsTrue(good);
905 Py_DECREF(good);
906 if (!ok) {
907 lz->start = 1;
908 return item;
909 }
910 Py_DECREF(item);
911 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000912}
913
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000914PyDoc_STRVAR(dropwhile_doc,
915"dropwhile(predicate, iterable) --> dropwhile object\n\
916\n\
917Drop items from the iterable while predicate(item) is true.\n\
918Afterwards, return every element until the iterable is exhausted.");
919
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000920static PyTypeObject dropwhile_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000921 PyVarObject_HEAD_INIT(NULL, 0)
922 "itertools.dropwhile", /* tp_name */
923 sizeof(dropwhileobject), /* tp_basicsize */
924 0, /* tp_itemsize */
925 /* methods */
926 (destructor)dropwhile_dealloc, /* tp_dealloc */
927 0, /* tp_print */
928 0, /* tp_getattr */
929 0, /* tp_setattr */
930 0, /* tp_reserved */
931 0, /* tp_repr */
932 0, /* tp_as_number */
933 0, /* tp_as_sequence */
934 0, /* tp_as_mapping */
935 0, /* tp_hash */
936 0, /* tp_call */
937 0, /* tp_str */
938 PyObject_GenericGetAttr, /* tp_getattro */
939 0, /* tp_setattro */
940 0, /* tp_as_buffer */
941 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
942 Py_TPFLAGS_BASETYPE, /* tp_flags */
943 dropwhile_doc, /* tp_doc */
944 (traverseproc)dropwhile_traverse, /* tp_traverse */
945 0, /* tp_clear */
946 0, /* tp_richcompare */
947 0, /* tp_weaklistoffset */
948 PyObject_SelfIter, /* tp_iter */
949 (iternextfunc)dropwhile_next, /* tp_iternext */
950 0, /* tp_methods */
951 0, /* tp_members */
952 0, /* tp_getset */
953 0, /* tp_base */
954 0, /* tp_dict */
955 0, /* tp_descr_get */
956 0, /* tp_descr_set */
957 0, /* tp_dictoffset */
958 0, /* tp_init */
959 0, /* tp_alloc */
960 dropwhile_new, /* tp_new */
961 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000962};
963
964
965/* takewhile object **********************************************************/
966
967typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000968 PyObject_HEAD
969 PyObject *func;
970 PyObject *it;
971 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000972} takewhileobject;
973
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000974static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000975
976static PyObject *
977takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
978{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000979 PyObject *func, *seq;
980 PyObject *it;
981 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000982
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000983 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
984 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000985
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000986 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
987 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000988
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000989 /* Get iterator. */
990 it = PyObject_GetIter(seq);
991 if (it == NULL)
992 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000993
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000994 /* create takewhileobject structure */
995 lz = (takewhileobject *)type->tp_alloc(type, 0);
996 if (lz == NULL) {
997 Py_DECREF(it);
998 return NULL;
999 }
1000 Py_INCREF(func);
1001 lz->func = func;
1002 lz->it = it;
1003 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001004
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001005 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001006}
1007
1008static void
1009takewhile_dealloc(takewhileobject *lz)
1010{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001011 PyObject_GC_UnTrack(lz);
1012 Py_XDECREF(lz->func);
1013 Py_XDECREF(lz->it);
1014 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001015}
1016
1017static int
1018takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1019{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001020 Py_VISIT(lz->it);
1021 Py_VISIT(lz->func);
1022 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001023}
1024
1025static PyObject *
1026takewhile_next(takewhileobject *lz)
1027{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001028 PyObject *item, *good;
1029 PyObject *it = lz->it;
1030 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001031
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001032 if (lz->stop == 1)
1033 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001034
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001035 item = (*Py_TYPE(it)->tp_iternext)(it);
1036 if (item == NULL)
1037 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001038
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001039 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1040 if (good == NULL) {
1041 Py_DECREF(item);
1042 return NULL;
1043 }
1044 ok = PyObject_IsTrue(good);
1045 Py_DECREF(good);
1046 if (ok)
1047 return item;
1048 Py_DECREF(item);
1049 lz->stop = 1;
1050 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001051}
1052
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053PyDoc_STRVAR(takewhile_doc,
1054"takewhile(predicate, iterable) --> takewhile object\n\
1055\n\
1056Return successive entries from an iterable as long as the \n\
1057predicate evaluates to true for each entry.");
1058
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001059static PyTypeObject takewhile_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001060 PyVarObject_HEAD_INIT(NULL, 0)
1061 "itertools.takewhile", /* tp_name */
1062 sizeof(takewhileobject), /* tp_basicsize */
1063 0, /* tp_itemsize */
1064 /* methods */
1065 (destructor)takewhile_dealloc, /* tp_dealloc */
1066 0, /* tp_print */
1067 0, /* tp_getattr */
1068 0, /* tp_setattr */
1069 0, /* tp_reserved */
1070 0, /* tp_repr */
1071 0, /* tp_as_number */
1072 0, /* tp_as_sequence */
1073 0, /* tp_as_mapping */
1074 0, /* tp_hash */
1075 0, /* tp_call */
1076 0, /* tp_str */
1077 PyObject_GenericGetAttr, /* tp_getattro */
1078 0, /* tp_setattro */
1079 0, /* tp_as_buffer */
1080 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1081 Py_TPFLAGS_BASETYPE, /* tp_flags */
1082 takewhile_doc, /* tp_doc */
1083 (traverseproc)takewhile_traverse, /* tp_traverse */
1084 0, /* tp_clear */
1085 0, /* tp_richcompare */
1086 0, /* tp_weaklistoffset */
1087 PyObject_SelfIter, /* tp_iter */
1088 (iternextfunc)takewhile_next, /* tp_iternext */
1089 0, /* tp_methods */
1090 0, /* tp_members */
1091 0, /* tp_getset */
1092 0, /* tp_base */
1093 0, /* tp_dict */
1094 0, /* tp_descr_get */
1095 0, /* tp_descr_set */
1096 0, /* tp_dictoffset */
1097 0, /* tp_init */
1098 0, /* tp_alloc */
1099 takewhile_new, /* tp_new */
1100 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001101};
1102
1103
1104/* islice object ************************************************************/
1105
1106typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001107 PyObject_HEAD
1108 PyObject *it;
1109 Py_ssize_t next;
1110 Py_ssize_t stop;
1111 Py_ssize_t step;
1112 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001113} isliceobject;
1114
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001115static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116
1117static PyObject *
1118islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1119{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001120 PyObject *seq;
1121 Py_ssize_t start=0, stop=-1, step=1;
1122 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1123 Py_ssize_t numargs;
1124 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001125
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001126 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1127 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001128
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001129 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1130 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001131
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001132 numargs = PyTuple_Size(args);
1133 if (numargs == 2) {
1134 if (a1 != Py_None) {
1135 stop = PyLong_AsSsize_t(a1);
1136 if (stop == -1) {
1137 if (PyErr_Occurred())
1138 PyErr_Clear();
1139 PyErr_SetString(PyExc_ValueError,
1140 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1141 return NULL;
1142 }
1143 }
1144 } else {
1145 if (a1 != Py_None)
1146 start = PyLong_AsSsize_t(a1);
1147 if (start == -1 && PyErr_Occurred())
1148 PyErr_Clear();
1149 if (a2 != Py_None) {
1150 stop = PyLong_AsSsize_t(a2);
1151 if (stop == -1) {
1152 if (PyErr_Occurred())
1153 PyErr_Clear();
1154 PyErr_SetString(PyExc_ValueError,
1155 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1156 return NULL;
1157 }
1158 }
1159 }
1160 if (start<0 || stop<-1) {
1161 PyErr_SetString(PyExc_ValueError,
1162 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1163 return NULL;
1164 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001165
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001166 if (a3 != NULL) {
1167 if (a3 != Py_None)
1168 step = PyLong_AsSsize_t(a3);
1169 if (step == -1 && PyErr_Occurred())
1170 PyErr_Clear();
1171 }
1172 if (step<1) {
1173 PyErr_SetString(PyExc_ValueError,
1174 "Step for islice() must be a positive integer or None.");
1175 return NULL;
1176 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001177
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001178 /* Get iterator. */
1179 it = PyObject_GetIter(seq);
1180 if (it == NULL)
1181 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001182
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001183 /* create isliceobject structure */
1184 lz = (isliceobject *)type->tp_alloc(type, 0);
1185 if (lz == NULL) {
1186 Py_DECREF(it);
1187 return NULL;
1188 }
1189 lz->it = it;
1190 lz->next = start;
1191 lz->stop = stop;
1192 lz->step = step;
1193 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001194
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001195 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196}
1197
1198static void
1199islice_dealloc(isliceobject *lz)
1200{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001201 PyObject_GC_UnTrack(lz);
1202 Py_XDECREF(lz->it);
1203 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204}
1205
1206static int
1207islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1208{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001209 Py_VISIT(lz->it);
1210 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001211}
1212
1213static PyObject *
1214islice_next(isliceobject *lz)
1215{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001216 PyObject *item;
1217 PyObject *it = lz->it;
1218 Py_ssize_t oldnext;
1219 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001220
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001221 iternext = *Py_TYPE(it)->tp_iternext;
1222 while (lz->cnt < lz->next) {
1223 item = iternext(it);
1224 if (item == NULL)
1225 return NULL;
1226 Py_DECREF(item);
1227 lz->cnt++;
1228 }
1229 if (lz->stop != -1 && lz->cnt >= lz->stop)
1230 return NULL;
1231 item = iternext(it);
1232 if (item == NULL)
1233 return NULL;
1234 lz->cnt++;
1235 oldnext = lz->next;
1236 lz->next += lz->step;
1237 if (lz->next < oldnext) /* Check for overflow */
1238 lz->next = lz->stop;
1239 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001240}
1241
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001242PyDoc_STRVAR(islice_doc,
1243"islice(iterable, [start,] stop [, step]) --> islice object\n\
1244\n\
1245Return an iterator whose next() method returns selected values from an\n\
1246iterable. If start is specified, will skip all preceding elements;\n\
1247otherwise, start defaults to zero. Step defaults to one. If\n\
1248specified as another value, step determines how many values are \n\
1249skipped between successive calls. Works like a slice() on a list\n\
1250but returns an iterator.");
1251
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001252static PyTypeObject islice_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001253 PyVarObject_HEAD_INIT(NULL, 0)
1254 "itertools.islice", /* tp_name */
1255 sizeof(isliceobject), /* tp_basicsize */
1256 0, /* tp_itemsize */
1257 /* methods */
1258 (destructor)islice_dealloc, /* tp_dealloc */
1259 0, /* tp_print */
1260 0, /* tp_getattr */
1261 0, /* tp_setattr */
1262 0, /* tp_reserved */
1263 0, /* tp_repr */
1264 0, /* tp_as_number */
1265 0, /* tp_as_sequence */
1266 0, /* tp_as_mapping */
1267 0, /* tp_hash */
1268 0, /* tp_call */
1269 0, /* tp_str */
1270 PyObject_GenericGetAttr, /* tp_getattro */
1271 0, /* tp_setattro */
1272 0, /* tp_as_buffer */
1273 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1274 Py_TPFLAGS_BASETYPE, /* tp_flags */
1275 islice_doc, /* tp_doc */
1276 (traverseproc)islice_traverse, /* tp_traverse */
1277 0, /* tp_clear */
1278 0, /* tp_richcompare */
1279 0, /* tp_weaklistoffset */
1280 PyObject_SelfIter, /* tp_iter */
1281 (iternextfunc)islice_next, /* tp_iternext */
1282 0, /* tp_methods */
1283 0, /* tp_members */
1284 0, /* tp_getset */
1285 0, /* tp_base */
1286 0, /* tp_dict */
1287 0, /* tp_descr_get */
1288 0, /* tp_descr_set */
1289 0, /* tp_dictoffset */
1290 0, /* tp_init */
1291 0, /* tp_alloc */
1292 islice_new, /* tp_new */
1293 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001294};
1295
1296
1297/* starmap object ************************************************************/
1298
1299typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001300 PyObject_HEAD
1301 PyObject *func;
1302 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001303} starmapobject;
1304
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001305static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001306
1307static PyObject *
1308starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1309{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001310 PyObject *func, *seq;
1311 PyObject *it;
1312 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001313
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001314 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1315 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001316
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001317 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1318 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001319
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001320 /* Get iterator. */
1321 it = PyObject_GetIter(seq);
1322 if (it == NULL)
1323 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001324
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001325 /* create starmapobject structure */
1326 lz = (starmapobject *)type->tp_alloc(type, 0);
1327 if (lz == NULL) {
1328 Py_DECREF(it);
1329 return NULL;
1330 }
1331 Py_INCREF(func);
1332 lz->func = func;
1333 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001334
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001335 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001336}
1337
1338static void
1339starmap_dealloc(starmapobject *lz)
1340{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001341 PyObject_GC_UnTrack(lz);
1342 Py_XDECREF(lz->func);
1343 Py_XDECREF(lz->it);
1344 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001345}
1346
1347static int
1348starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1349{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001350 Py_VISIT(lz->it);
1351 Py_VISIT(lz->func);
1352 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001353}
1354
1355static PyObject *
1356starmap_next(starmapobject *lz)
1357{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001358 PyObject *args;
1359 PyObject *result;
1360 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001361
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001362 args = (*Py_TYPE(it)->tp_iternext)(it);
1363 if (args == NULL)
1364 return NULL;
1365 if (!PyTuple_CheckExact(args)) {
1366 PyObject *newargs = PySequence_Tuple(args);
1367 Py_DECREF(args);
1368 if (newargs == NULL)
1369 return NULL;
1370 args = newargs;
1371 }
1372 result = PyObject_Call(lz->func, args, NULL);
1373 Py_DECREF(args);
1374 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001375}
1376
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001377PyDoc_STRVAR(starmap_doc,
1378"starmap(function, sequence) --> starmap object\n\
1379\n\
1380Return an iterator whose values are returned from the function evaluated\n\
1381with a argument tuple taken from the given sequence.");
1382
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001383static PyTypeObject starmap_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001384 PyVarObject_HEAD_INIT(NULL, 0)
1385 "itertools.starmap", /* tp_name */
1386 sizeof(starmapobject), /* tp_basicsize */
1387 0, /* tp_itemsize */
1388 /* methods */
1389 (destructor)starmap_dealloc, /* tp_dealloc */
1390 0, /* tp_print */
1391 0, /* tp_getattr */
1392 0, /* tp_setattr */
1393 0, /* tp_reserved */
1394 0, /* tp_repr */
1395 0, /* tp_as_number */
1396 0, /* tp_as_sequence */
1397 0, /* tp_as_mapping */
1398 0, /* tp_hash */
1399 0, /* tp_call */
1400 0, /* tp_str */
1401 PyObject_GenericGetAttr, /* tp_getattro */
1402 0, /* tp_setattro */
1403 0, /* tp_as_buffer */
1404 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1405 Py_TPFLAGS_BASETYPE, /* tp_flags */
1406 starmap_doc, /* tp_doc */
1407 (traverseproc)starmap_traverse, /* tp_traverse */
1408 0, /* tp_clear */
1409 0, /* tp_richcompare */
1410 0, /* tp_weaklistoffset */
1411 PyObject_SelfIter, /* tp_iter */
1412 (iternextfunc)starmap_next, /* tp_iternext */
1413 0, /* tp_methods */
1414 0, /* tp_members */
1415 0, /* tp_getset */
1416 0, /* tp_base */
1417 0, /* tp_dict */
1418 0, /* tp_descr_get */
1419 0, /* tp_descr_set */
1420 0, /* tp_dictoffset */
1421 0, /* tp_init */
1422 0, /* tp_alloc */
1423 starmap_new, /* tp_new */
1424 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001425};
1426
1427
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001428/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001429
1430typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001431 PyObject_HEAD
1432 PyObject *source; /* Iterator over input iterables */
1433 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001434} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001435
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001436static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001437
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001438static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001439chain_new_internal(PyTypeObject *type, PyObject *source)
1440{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001441 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001442
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001443 lz = (chainobject *)type->tp_alloc(type, 0);
1444 if (lz == NULL) {
1445 Py_DECREF(source);
1446 return NULL;
1447 }
1448
1449 lz->source = source;
1450 lz->active = NULL;
1451 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001452}
1453
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001454static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001455chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001456{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001457 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001458
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001459 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1460 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001461
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001462 source = PyObject_GetIter(args);
1463 if (source == NULL)
1464 return NULL;
1465
1466 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001467}
1468
1469static PyObject *
1470chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1471{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001472 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001473
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001474 source = PyObject_GetIter(arg);
1475 if (source == NULL)
1476 return NULL;
1477
1478 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001479}
1480
1481static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001482chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001483{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001484 PyObject_GC_UnTrack(lz);
1485 Py_XDECREF(lz->active);
1486 Py_XDECREF(lz->source);
1487 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001488}
1489
Raymond Hettinger2012f172003-02-07 05:32:58 +00001490static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001491chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001492{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001493 Py_VISIT(lz->source);
1494 Py_VISIT(lz->active);
1495 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001496}
1497
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001498static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001499chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001500{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001501 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001502
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001503 if (lz->source == NULL)
1504 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001505
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001506 if (lz->active == NULL) {
1507 PyObject *iterable = PyIter_Next(lz->source);
1508 if (iterable == NULL) {
1509 Py_CLEAR(lz->source);
1510 return NULL; /* no more input sources */
1511 }
1512 lz->active = PyObject_GetIter(iterable);
1513 Py_DECREF(iterable);
1514 if (lz->active == NULL) {
1515 Py_CLEAR(lz->source);
1516 return NULL; /* input not iterable */
1517 }
1518 }
1519 item = PyIter_Next(lz->active);
1520 if (item != NULL)
1521 return item;
1522 if (PyErr_Occurred()) {
1523 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1524 PyErr_Clear();
1525 else
1526 return NULL; /* input raised an exception */
1527 }
1528 Py_CLEAR(lz->active);
1529 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001530}
1531
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001532PyDoc_STRVAR(chain_doc,
1533"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001534\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001535Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001536first iterable until it is exhausted, then elements from the next\n\
1537iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001538
Christian Heimesf16baeb2008-02-29 14:57:44 +00001539PyDoc_STRVAR(chain_from_iterable_doc,
1540"chain.from_iterable(iterable) --> chain object\n\
1541\n\
1542Alternate chain() contructor taking a single iterable argument\n\
1543that evaluates lazily.");
1544
1545static PyMethodDef chain_methods[] = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001546 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1547 chain_from_iterable_doc},
1548 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001549};
1550
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001551static PyTypeObject chain_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001552 PyVarObject_HEAD_INIT(NULL, 0)
1553 "itertools.chain", /* tp_name */
1554 sizeof(chainobject), /* tp_basicsize */
1555 0, /* tp_itemsize */
1556 /* methods */
1557 (destructor)chain_dealloc, /* tp_dealloc */
1558 0, /* tp_print */
1559 0, /* tp_getattr */
1560 0, /* tp_setattr */
1561 0, /* tp_reserved */
1562 0, /* tp_repr */
1563 0, /* tp_as_number */
1564 0, /* tp_as_sequence */
1565 0, /* tp_as_mapping */
1566 0, /* tp_hash */
1567 0, /* tp_call */
1568 0, /* tp_str */
1569 PyObject_GenericGetAttr, /* tp_getattro */
1570 0, /* tp_setattro */
1571 0, /* tp_as_buffer */
1572 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1573 Py_TPFLAGS_BASETYPE, /* tp_flags */
1574 chain_doc, /* tp_doc */
1575 (traverseproc)chain_traverse, /* tp_traverse */
1576 0, /* tp_clear */
1577 0, /* tp_richcompare */
1578 0, /* tp_weaklistoffset */
1579 PyObject_SelfIter, /* tp_iter */
1580 (iternextfunc)chain_next, /* tp_iternext */
1581 chain_methods, /* tp_methods */
1582 0, /* tp_members */
1583 0, /* tp_getset */
1584 0, /* tp_base */
1585 0, /* tp_dict */
1586 0, /* tp_descr_get */
1587 0, /* tp_descr_set */
1588 0, /* tp_dictoffset */
1589 0, /* tp_init */
1590 0, /* tp_alloc */
1591 chain_new, /* tp_new */
1592 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001593};
1594
1595
Christian Heimesc3f30c42008-02-22 16:37:40 +00001596/* product object ************************************************************/
1597
1598typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001599 PyObject_HEAD
1600 PyObject *pools; /* tuple of pool tuples */
1601 Py_ssize_t *indices; /* one index per pool */
1602 PyObject *result; /* most recently returned result tuple */
1603 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001604} productobject;
1605
1606static PyTypeObject product_type;
1607
1608static PyObject *
1609product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1610{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001611 productobject *lz;
1612 Py_ssize_t nargs, npools, repeat=1;
1613 PyObject *pools = NULL;
1614 Py_ssize_t *indices = NULL;
1615 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001616
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001617 if (kwds != NULL) {
1618 char *kwlist[] = {"repeat", 0};
1619 PyObject *tmpargs = PyTuple_New(0);
1620 if (tmpargs == NULL)
1621 return NULL;
1622 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1623 Py_DECREF(tmpargs);
1624 return NULL;
1625 }
1626 Py_DECREF(tmpargs);
1627 if (repeat < 0) {
1628 PyErr_SetString(PyExc_ValueError,
1629 "repeat argument cannot be negative");
1630 return NULL;
1631 }
1632 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001633
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001634 assert(PyTuple_Check(args));
1635 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1636 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001637
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001638 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1639 if (indices == NULL) {
1640 PyErr_NoMemory();
1641 goto error;
1642 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001643
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001644 pools = PyTuple_New(npools);
1645 if (pools == NULL)
1646 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001647
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001648 for (i=0; i < nargs ; ++i) {
1649 PyObject *item = PyTuple_GET_ITEM(args, i);
1650 PyObject *pool = PySequence_Tuple(item);
1651 if (pool == NULL)
1652 goto error;
1653 PyTuple_SET_ITEM(pools, i, pool);
1654 indices[i] = 0;
1655 }
1656 for ( ; i < npools; ++i) {
1657 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1658 Py_INCREF(pool);
1659 PyTuple_SET_ITEM(pools, i, pool);
1660 indices[i] = 0;
1661 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001662
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001663 /* create productobject structure */
1664 lz = (productobject *)type->tp_alloc(type, 0);
1665 if (lz == NULL)
1666 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001667
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001668 lz->pools = pools;
1669 lz->indices = indices;
1670 lz->result = NULL;
1671 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001672
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001673 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001674
1675error:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001676 if (indices != NULL)
1677 PyMem_Free(indices);
1678 Py_XDECREF(pools);
1679 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001680}
1681
1682static void
1683product_dealloc(productobject *lz)
1684{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001685 PyObject_GC_UnTrack(lz);
1686 Py_XDECREF(lz->pools);
1687 Py_XDECREF(lz->result);
1688 if (lz->indices != NULL)
1689 PyMem_Free(lz->indices);
1690 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001691}
1692
1693static int
1694product_traverse(productobject *lz, visitproc visit, void *arg)
1695{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001696 Py_VISIT(lz->pools);
1697 Py_VISIT(lz->result);
1698 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001699}
1700
1701static PyObject *
1702product_next(productobject *lz)
1703{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001704 PyObject *pool;
1705 PyObject *elem;
1706 PyObject *oldelem;
1707 PyObject *pools = lz->pools;
1708 PyObject *result = lz->result;
1709 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1710 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001711
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001712 if (lz->stopped)
1713 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001714
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001715 if (result == NULL) {
1716 /* On the first pass, return an initial tuple filled with the
1717 first element from each pool. */
1718 result = PyTuple_New(npools);
1719 if (result == NULL)
1720 goto empty;
1721 lz->result = result;
1722 for (i=0; i < npools; i++) {
1723 pool = PyTuple_GET_ITEM(pools, i);
1724 if (PyTuple_GET_SIZE(pool) == 0)
1725 goto empty;
1726 elem = PyTuple_GET_ITEM(pool, 0);
1727 Py_INCREF(elem);
1728 PyTuple_SET_ITEM(result, i, elem);
1729 }
1730 } else {
1731 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001732
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001733 /* Copy the previous result tuple or re-use it if available */
1734 if (Py_REFCNT(result) > 1) {
1735 PyObject *old_result = result;
1736 result = PyTuple_New(npools);
1737 if (result == NULL)
1738 goto empty;
1739 lz->result = result;
1740 for (i=0; i < npools; i++) {
1741 elem = PyTuple_GET_ITEM(old_result, i);
1742 Py_INCREF(elem);
1743 PyTuple_SET_ITEM(result, i, elem);
1744 }
1745 Py_DECREF(old_result);
1746 }
1747 /* Now, we've got the only copy so we can update it in-place */
1748 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001749
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001750 /* Update the pool indices right-to-left. Only advance to the
1751 next pool when the previous one rolls-over */
1752 for (i=npools-1 ; i >= 0 ; i--) {
1753 pool = PyTuple_GET_ITEM(pools, i);
1754 indices[i]++;
1755 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1756 /* Roll-over and advance to next pool */
1757 indices[i] = 0;
1758 elem = PyTuple_GET_ITEM(pool, 0);
1759 Py_INCREF(elem);
1760 oldelem = PyTuple_GET_ITEM(result, i);
1761 PyTuple_SET_ITEM(result, i, elem);
1762 Py_DECREF(oldelem);
1763 } else {
1764 /* No rollover. Just increment and stop here. */
1765 elem = PyTuple_GET_ITEM(pool, indices[i]);
1766 Py_INCREF(elem);
1767 oldelem = PyTuple_GET_ITEM(result, i);
1768 PyTuple_SET_ITEM(result, i, elem);
1769 Py_DECREF(oldelem);
1770 break;
1771 }
1772 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001773
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001774 /* If i is negative, then the indices have all rolled-over
1775 and we're done. */
1776 if (i < 0)
1777 goto empty;
1778 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001779
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001780 Py_INCREF(result);
1781 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001782
1783empty:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001784 lz->stopped = 1;
1785 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001786}
1787
1788PyDoc_STRVAR(product_doc,
1789"product(*iterables) --> product object\n\
1790\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001791Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001792For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1793The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1794cycle in a manner similar to an odometer (with the rightmost element changing\n\
1795on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00001796To compute the product of an iterable with itself, specify the number\n\
1797of repetitions with the optional repeat keyword argument. For example,\n\
1798product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001799product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1800product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1801
1802static PyTypeObject product_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001803 PyVarObject_HEAD_INIT(NULL, 0)
1804 "itertools.product", /* tp_name */
1805 sizeof(productobject), /* tp_basicsize */
1806 0, /* tp_itemsize */
1807 /* methods */
1808 (destructor)product_dealloc, /* tp_dealloc */
1809 0, /* tp_print */
1810 0, /* tp_getattr */
1811 0, /* tp_setattr */
1812 0, /* tp_reserved */
1813 0, /* tp_repr */
1814 0, /* tp_as_number */
1815 0, /* tp_as_sequence */
1816 0, /* tp_as_mapping */
1817 0, /* tp_hash */
1818 0, /* tp_call */
1819 0, /* tp_str */
1820 PyObject_GenericGetAttr, /* tp_getattro */
1821 0, /* tp_setattro */
1822 0, /* tp_as_buffer */
1823 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1824 Py_TPFLAGS_BASETYPE, /* tp_flags */
1825 product_doc, /* tp_doc */
1826 (traverseproc)product_traverse, /* tp_traverse */
1827 0, /* tp_clear */
1828 0, /* tp_richcompare */
1829 0, /* tp_weaklistoffset */
1830 PyObject_SelfIter, /* tp_iter */
1831 (iternextfunc)product_next, /* tp_iternext */
1832 0, /* tp_methods */
1833 0, /* tp_members */
1834 0, /* tp_getset */
1835 0, /* tp_base */
1836 0, /* tp_dict */
1837 0, /* tp_descr_get */
1838 0, /* tp_descr_set */
1839 0, /* tp_dictoffset */
1840 0, /* tp_init */
1841 0, /* tp_alloc */
1842 product_new, /* tp_new */
1843 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001844};
1845
1846
Christian Heimes380f7f22008-02-28 11:19:05 +00001847/* combinations object ************************************************************/
1848
1849typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001850 PyObject_HEAD
1851 PyObject *pool; /* input converted to a tuple */
1852 Py_ssize_t *indices; /* one index per result element */
1853 PyObject *result; /* most recently returned result tuple */
1854 Py_ssize_t r; /* size of result tuple */
1855 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00001856} combinationsobject;
1857
1858static PyTypeObject combinations_type;
1859
1860static PyObject *
1861combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1862{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001863 combinationsobject *co;
1864 Py_ssize_t n;
1865 Py_ssize_t r;
1866 PyObject *pool = NULL;
1867 PyObject *iterable = NULL;
1868 Py_ssize_t *indices = NULL;
1869 Py_ssize_t i;
1870 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00001871
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001872 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
1873 &iterable, &r))
1874 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001875
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001876 pool = PySequence_Tuple(iterable);
1877 if (pool == NULL)
1878 goto error;
1879 n = PyTuple_GET_SIZE(pool);
1880 if (r < 0) {
1881 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
1882 goto error;
1883 }
Christian Heimes380f7f22008-02-28 11:19:05 +00001884
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001885 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
1886 if (indices == NULL) {
1887 PyErr_NoMemory();
1888 goto error;
1889 }
Christian Heimes380f7f22008-02-28 11:19:05 +00001890
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001891 for (i=0 ; i<r ; i++)
1892 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00001893
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001894 /* create combinationsobject structure */
1895 co = (combinationsobject *)type->tp_alloc(type, 0);
1896 if (co == NULL)
1897 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00001898
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001899 co->pool = pool;
1900 co->indices = indices;
1901 co->result = NULL;
1902 co->r = r;
1903 co->stopped = r > n ? 1 : 0;
1904
1905 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00001906
1907error:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001908 if (indices != NULL)
1909 PyMem_Free(indices);
1910 Py_XDECREF(pool);
1911 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001912}
1913
1914static void
1915combinations_dealloc(combinationsobject *co)
1916{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001917 PyObject_GC_UnTrack(co);
1918 Py_XDECREF(co->pool);
1919 Py_XDECREF(co->result);
1920 if (co->indices != NULL)
1921 PyMem_Free(co->indices);
1922 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00001923}
1924
1925static int
1926combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
1927{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001928 Py_VISIT(co->pool);
1929 Py_VISIT(co->result);
1930 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00001931}
1932
1933static PyObject *
1934combinations_next(combinationsobject *co)
1935{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001936 PyObject *elem;
1937 PyObject *oldelem;
1938 PyObject *pool = co->pool;
1939 Py_ssize_t *indices = co->indices;
1940 PyObject *result = co->result;
1941 Py_ssize_t n = PyTuple_GET_SIZE(pool);
1942 Py_ssize_t r = co->r;
1943 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00001944
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001945 if (co->stopped)
1946 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001947
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001948 if (result == NULL) {
1949 /* On the first pass, initialize result tuple using the indices */
1950 result = PyTuple_New(r);
1951 if (result == NULL)
1952 goto empty;
1953 co->result = result;
1954 for (i=0; i<r ; i++) {
1955 index = indices[i];
1956 elem = PyTuple_GET_ITEM(pool, index);
1957 Py_INCREF(elem);
1958 PyTuple_SET_ITEM(result, i, elem);
1959 }
1960 } else {
1961 /* Copy the previous result tuple or re-use it if available */
1962 if (Py_REFCNT(result) > 1) {
1963 PyObject *old_result = result;
1964 result = PyTuple_New(r);
1965 if (result == NULL)
1966 goto empty;
1967 co->result = result;
1968 for (i=0; i<r ; i++) {
1969 elem = PyTuple_GET_ITEM(old_result, i);
1970 Py_INCREF(elem);
1971 PyTuple_SET_ITEM(result, i, elem);
1972 }
1973 Py_DECREF(old_result);
1974 }
1975 /* Now, we've got the only copy so we can update it in-place
1976 * CPython's empty tuple is a singleton and cached in
1977 * PyTuple's freelist.
1978 */
1979 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00001980
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001981 /* Scan indices right-to-left until finding one that is not
1982 at its maximum (i + n - r). */
1983 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
1984 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00001985
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001986 /* If i is negative, then the indices are all at
1987 their maximum value and we're done. */
1988 if (i < 0)
1989 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00001990
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001991 /* Increment the current index which we know is not at its
1992 maximum. Then move back to the right setting each index
1993 to its lowest possible value (one higher than the index
1994 to its left -- this maintains the sort order invariant). */
1995 indices[i]++;
1996 for (j=i+1 ; j<r ; j++)
1997 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00001998
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001999 /* Update the result tuple for the new indices
2000 starting with i, the leftmost index that changed */
2001 for ( ; i<r ; i++) {
2002 index = indices[i];
2003 elem = PyTuple_GET_ITEM(pool, index);
2004 Py_INCREF(elem);
2005 oldelem = PyTuple_GET_ITEM(result, i);
2006 PyTuple_SET_ITEM(result, i, elem);
2007 Py_DECREF(oldelem);
2008 }
2009 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002010
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002011 Py_INCREF(result);
2012 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002013
2014empty:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002015 co->stopped = 1;
2016 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002017}
2018
2019PyDoc_STRVAR(combinations_doc,
Raymond Hettingerd3a77c02009-11-19 01:23:41 +00002020"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002021\n\
2022Return successive r-length combinations of elements in the iterable.\n\n\
2023combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2024
2025static PyTypeObject combinations_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002026 PyVarObject_HEAD_INIT(NULL, 0)
2027 "itertools.combinations", /* tp_name */
2028 sizeof(combinationsobject), /* tp_basicsize */
2029 0, /* tp_itemsize */
2030 /* methods */
2031 (destructor)combinations_dealloc, /* tp_dealloc */
2032 0, /* tp_print */
2033 0, /* tp_getattr */
2034 0, /* tp_setattr */
2035 0, /* tp_reserved */
2036 0, /* tp_repr */
2037 0, /* tp_as_number */
2038 0, /* tp_as_sequence */
2039 0, /* tp_as_mapping */
2040 0, /* tp_hash */
2041 0, /* tp_call */
2042 0, /* tp_str */
2043 PyObject_GenericGetAttr, /* tp_getattro */
2044 0, /* tp_setattro */
2045 0, /* tp_as_buffer */
2046 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2047 Py_TPFLAGS_BASETYPE, /* tp_flags */
2048 combinations_doc, /* tp_doc */
2049 (traverseproc)combinations_traverse, /* tp_traverse */
2050 0, /* tp_clear */
2051 0, /* tp_richcompare */
2052 0, /* tp_weaklistoffset */
2053 PyObject_SelfIter, /* tp_iter */
2054 (iternextfunc)combinations_next, /* tp_iternext */
2055 0, /* tp_methods */
2056 0, /* tp_members */
2057 0, /* tp_getset */
2058 0, /* tp_base */
2059 0, /* tp_dict */
2060 0, /* tp_descr_get */
2061 0, /* tp_descr_set */
2062 0, /* tp_dictoffset */
2063 0, /* tp_init */
2064 0, /* tp_alloc */
2065 combinations_new, /* tp_new */
2066 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002067};
2068
2069
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002070/* combinations with replacement object *******************************************/
2071
2072/* Equivalent to:
2073
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002074 def combinations_with_replacement(iterable, r):
2075 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2076 # number items returned: (n+r-1)! / r! / (n-1)!
2077 pool = tuple(iterable)
2078 n = len(pool)
2079 indices = [0] * r
2080 yield tuple(pool[i] for i in indices)
2081 while 1:
2082 for i in reversed(range(r)):
2083 if indices[i] != n - 1:
2084 break
2085 else:
2086 return
2087 indices[i:] = [indices[i] + 1] * (r - i)
2088 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002089
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002090 def combinations_with_replacement2(iterable, r):
2091 'Alternate version that filters from product()'
2092 pool = tuple(iterable)
2093 n = len(pool)
2094 for indices in product(range(n), repeat=r):
2095 if sorted(indices) == list(indices):
2096 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002097*/
2098typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002099 PyObject_HEAD
2100 PyObject *pool; /* input converted to a tuple */
2101 Py_ssize_t *indices; /* one index per result element */
2102 PyObject *result; /* most recently returned result tuple */
2103 Py_ssize_t r; /* size of result tuple */
2104 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002105} cwrobject;
2106
2107static PyTypeObject cwr_type;
2108
2109static PyObject *
2110cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2111{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002112 cwrobject *co;
2113 Py_ssize_t n;
2114 Py_ssize_t r;
2115 PyObject *pool = NULL;
2116 PyObject *iterable = NULL;
2117 Py_ssize_t *indices = NULL;
2118 Py_ssize_t i;
2119 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002120
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002121 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2122 &iterable, &r))
2123 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002124
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002125 pool = PySequence_Tuple(iterable);
2126 if (pool == NULL)
2127 goto error;
2128 n = PyTuple_GET_SIZE(pool);
2129 if (r < 0) {
2130 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2131 goto error;
2132 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002133
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002134 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2135 if (indices == NULL) {
2136 PyErr_NoMemory();
2137 goto error;
2138 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002139
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002140 for (i=0 ; i<r ; i++)
2141 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002142
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002143 /* create cwrobject structure */
2144 co = (cwrobject *)type->tp_alloc(type, 0);
2145 if (co == NULL)
2146 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002147
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002148 co->pool = pool;
2149 co->indices = indices;
2150 co->result = NULL;
2151 co->r = r;
2152 co->stopped = !n && r;
2153
2154 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002155
2156error:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002157 if (indices != NULL)
2158 PyMem_Free(indices);
2159 Py_XDECREF(pool);
2160 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002161}
2162
2163static void
2164cwr_dealloc(cwrobject *co)
2165{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002166 PyObject_GC_UnTrack(co);
2167 Py_XDECREF(co->pool);
2168 Py_XDECREF(co->result);
2169 if (co->indices != NULL)
2170 PyMem_Free(co->indices);
2171 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002172}
2173
2174static int
2175cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2176{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002177 Py_VISIT(co->pool);
2178 Py_VISIT(co->result);
2179 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002180}
2181
2182static PyObject *
2183cwr_next(cwrobject *co)
2184{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002185 PyObject *elem;
2186 PyObject *oldelem;
2187 PyObject *pool = co->pool;
2188 Py_ssize_t *indices = co->indices;
2189 PyObject *result = co->result;
2190 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2191 Py_ssize_t r = co->r;
2192 Py_ssize_t i, j, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002193
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002194 if (co->stopped)
2195 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002196
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002197 if (result == NULL) {
2198 /* On the first pass, initialize result tuple using the indices */
2199 result = PyTuple_New(r);
2200 if (result == NULL)
2201 goto empty;
2202 co->result = result;
2203 for (i=0; i<r ; i++) {
2204 index = indices[i];
2205 elem = PyTuple_GET_ITEM(pool, index);
2206 Py_INCREF(elem);
2207 PyTuple_SET_ITEM(result, i, elem);
2208 }
2209 } else {
2210 /* Copy the previous result tuple or re-use it if available */
2211 if (Py_REFCNT(result) > 1) {
2212 PyObject *old_result = result;
2213 result = PyTuple_New(r);
2214 if (result == NULL)
2215 goto empty;
2216 co->result = result;
2217 for (i=0; i<r ; i++) {
2218 elem = PyTuple_GET_ITEM(old_result, i);
2219 Py_INCREF(elem);
2220 PyTuple_SET_ITEM(result, i, elem);
2221 }
2222 Py_DECREF(old_result);
2223 }
2224 /* Now, we've got the only copy so we can update it in-place CPython's
2225 empty tuple is a singleton and cached in PyTuple's freelist. */
2226 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002227
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002228 /* Scan indices right-to-left until finding one that is not
2229 * at its maximum (n-1). */
2230 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2231 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002232
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002233 /* If i is negative, then the indices are all at
2234 their maximum value and we're done. */
2235 if (i < 0)
2236 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002237
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002238 /* Increment the current index which we know is not at its
2239 maximum. Then set all to the right to the same value. */
2240 indices[i]++;
2241 for (j=i+1 ; j<r ; j++)
2242 indices[j] = indices[j-1];
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002243
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002244 /* Update the result tuple for the new indices
2245 starting with i, the leftmost index that changed */
2246 for ( ; i<r ; i++) {
2247 index = indices[i];
2248 elem = PyTuple_GET_ITEM(pool, index);
2249 Py_INCREF(elem);
2250 oldelem = PyTuple_GET_ITEM(result, i);
2251 PyTuple_SET_ITEM(result, i, elem);
2252 Py_DECREF(oldelem);
2253 }
2254 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002255
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002256 Py_INCREF(result);
2257 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002258
2259empty:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002260 co->stopped = 1;
2261 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002262}
2263
2264PyDoc_STRVAR(cwr_doc,
Raymond Hettingerd3a77c02009-11-19 01:23:41 +00002265"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002266\n\
2267Return successive r-length combinations of elements in the iterable\n\
2268allowing individual elements to have successive repeats.\n\
2269combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2270
2271static PyTypeObject cwr_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002272 PyVarObject_HEAD_INIT(NULL, 0)
2273 "itertools.combinations_with_replacement", /* tp_name */
2274 sizeof(cwrobject), /* tp_basicsize */
2275 0, /* tp_itemsize */
2276 /* methods */
2277 (destructor)cwr_dealloc, /* tp_dealloc */
2278 0, /* tp_print */
2279 0, /* tp_getattr */
2280 0, /* tp_setattr */
2281 0, /* tp_reserved */
2282 0, /* tp_repr */
2283 0, /* tp_as_number */
2284 0, /* tp_as_sequence */
2285 0, /* tp_as_mapping */
2286 0, /* tp_hash */
2287 0, /* tp_call */
2288 0, /* tp_str */
2289 PyObject_GenericGetAttr, /* tp_getattro */
2290 0, /* tp_setattro */
2291 0, /* tp_as_buffer */
2292 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2293 Py_TPFLAGS_BASETYPE, /* tp_flags */
2294 cwr_doc, /* tp_doc */
2295 (traverseproc)cwr_traverse, /* tp_traverse */
2296 0, /* tp_clear */
2297 0, /* tp_richcompare */
2298 0, /* tp_weaklistoffset */
2299 PyObject_SelfIter, /* tp_iter */
2300 (iternextfunc)cwr_next, /* tp_iternext */
2301 0, /* tp_methods */
2302 0, /* tp_members */
2303 0, /* tp_getset */
2304 0, /* tp_base */
2305 0, /* tp_dict */
2306 0, /* tp_descr_get */
2307 0, /* tp_descr_set */
2308 0, /* tp_dictoffset */
2309 0, /* tp_init */
2310 0, /* tp_alloc */
2311 cwr_new, /* tp_new */
2312 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002313};
2314
2315
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002316/* permutations object ************************************************************
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002317
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002318def permutations(iterable, r=None):
2319 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2320 pool = tuple(iterable)
2321 n = len(pool)
2322 r = n if r is None else r
2323 indices = range(n)
2324 cycles = range(n-r+1, n+1)[::-1]
2325 yield tuple(pool[i] for i in indices[:r])
2326 while n:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002327 for i in reversed(range(r)):
2328 cycles[i] -= 1
2329 if cycles[i] == 0:
2330 indices[i:] = indices[i+1:] + indices[i:i+1]
2331 cycles[i] = n - i
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002332 else:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002333 j = cycles[i]
2334 indices[i], indices[-j] = indices[-j], indices[i]
2335 yield tuple(pool[i] for i in indices[:r])
2336 break
2337 else:
2338 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002339*/
2340
2341typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002342 PyObject_HEAD
2343 PyObject *pool; /* input converted to a tuple */
2344 Py_ssize_t *indices; /* one index per element in the pool */
2345 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2346 PyObject *result; /* most recently returned result tuple */
2347 Py_ssize_t r; /* size of result tuple */
2348 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002349} permutationsobject;
2350
2351static PyTypeObject permutations_type;
2352
2353static PyObject *
2354permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2355{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002356 permutationsobject *po;
2357 Py_ssize_t n;
2358 Py_ssize_t r;
2359 PyObject *robj = Py_None;
2360 PyObject *pool = NULL;
2361 PyObject *iterable = NULL;
2362 Py_ssize_t *indices = NULL;
2363 Py_ssize_t *cycles = NULL;
2364 Py_ssize_t i;
2365 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002366
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002367 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2368 &iterable, &robj))
2369 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002370
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002371 pool = PySequence_Tuple(iterable);
2372 if (pool == NULL)
2373 goto error;
2374 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002375
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002376 r = n;
2377 if (robj != Py_None) {
2378 if (!PyLong_Check(robj)) {
2379 PyErr_SetString(PyExc_TypeError, "Expected int as r");
2380 goto error;
2381 }
2382 r = PyLong_AsSsize_t(robj);
2383 if (r == -1 && PyErr_Occurred())
2384 goto error;
2385 }
2386 if (r < 0) {
2387 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2388 goto error;
2389 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002390
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002391 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2392 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2393 if (indices == NULL || cycles == NULL) {
2394 PyErr_NoMemory();
2395 goto error;
2396 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002397
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002398 for (i=0 ; i<n ; i++)
2399 indices[i] = i;
2400 for (i=0 ; i<r ; i++)
2401 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002402
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002403 /* create permutationsobject structure */
2404 po = (permutationsobject *)type->tp_alloc(type, 0);
2405 if (po == NULL)
2406 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002407
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002408 po->pool = pool;
2409 po->indices = indices;
2410 po->cycles = cycles;
2411 po->result = NULL;
2412 po->r = r;
2413 po->stopped = r > n ? 1 : 0;
2414
2415 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002416
2417error:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002418 if (indices != NULL)
2419 PyMem_Free(indices);
2420 if (cycles != NULL)
2421 PyMem_Free(cycles);
2422 Py_XDECREF(pool);
2423 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002424}
2425
2426static void
2427permutations_dealloc(permutationsobject *po)
2428{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002429 PyObject_GC_UnTrack(po);
2430 Py_XDECREF(po->pool);
2431 Py_XDECREF(po->result);
2432 PyMem_Free(po->indices);
2433 PyMem_Free(po->cycles);
2434 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002435}
2436
2437static int
2438permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2439{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002440 Py_VISIT(po->pool);
2441 Py_VISIT(po->result);
2442 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002443}
2444
2445static PyObject *
2446permutations_next(permutationsobject *po)
2447{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002448 PyObject *elem;
2449 PyObject *oldelem;
2450 PyObject *pool = po->pool;
2451 Py_ssize_t *indices = po->indices;
2452 Py_ssize_t *cycles = po->cycles;
2453 PyObject *result = po->result;
2454 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2455 Py_ssize_t r = po->r;
2456 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002457
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002458 if (po->stopped)
2459 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002460
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002461 if (result == NULL) {
2462 /* On the first pass, initialize result tuple using the indices */
2463 result = PyTuple_New(r);
2464 if (result == NULL)
2465 goto empty;
2466 po->result = result;
2467 for (i=0; i<r ; i++) {
2468 index = indices[i];
2469 elem = PyTuple_GET_ITEM(pool, index);
2470 Py_INCREF(elem);
2471 PyTuple_SET_ITEM(result, i, elem);
2472 }
2473 } else {
2474 if (n == 0)
2475 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002476
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002477 /* Copy the previous result tuple or re-use it if available */
2478 if (Py_REFCNT(result) > 1) {
2479 PyObject *old_result = result;
2480 result = PyTuple_New(r);
2481 if (result == NULL)
2482 goto empty;
2483 po->result = result;
2484 for (i=0; i<r ; i++) {
2485 elem = PyTuple_GET_ITEM(old_result, i);
2486 Py_INCREF(elem);
2487 PyTuple_SET_ITEM(result, i, elem);
2488 }
2489 Py_DECREF(old_result);
2490 }
2491 /* Now, we've got the only copy so we can update it in-place */
2492 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002493
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002494 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2495 for (i=r-1 ; i>=0 ; i--) {
2496 cycles[i] -= 1;
2497 if (cycles[i] == 0) {
2498 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2499 index = indices[i];
2500 for (j=i ; j<n-1 ; j++)
2501 indices[j] = indices[j+1];
2502 indices[n-1] = index;
2503 cycles[i] = n - i;
2504 } else {
2505 j = cycles[i];
2506 index = indices[i];
2507 indices[i] = indices[n-j];
2508 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002509
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002510 for (k=i; k<r ; k++) {
2511 /* start with i, the leftmost element that changed */
2512 /* yield tuple(pool[k] for k in indices[:r]) */
2513 index = indices[k];
2514 elem = PyTuple_GET_ITEM(pool, index);
2515 Py_INCREF(elem);
2516 oldelem = PyTuple_GET_ITEM(result, k);
2517 PyTuple_SET_ITEM(result, k, elem);
2518 Py_DECREF(oldelem);
2519 }
2520 break;
2521 }
2522 }
2523 /* If i is negative, then the cycles have all
2524 rolled-over and we're done. */
2525 if (i < 0)
2526 goto empty;
2527 }
2528 Py_INCREF(result);
2529 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002530
2531empty:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002532 po->stopped = 1;
2533 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002534}
2535
2536PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00002537"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002538\n\
2539Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00002540permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002541
2542static PyTypeObject permutations_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002543 PyVarObject_HEAD_INIT(NULL, 0)
2544 "itertools.permutations", /* tp_name */
2545 sizeof(permutationsobject), /* tp_basicsize */
2546 0, /* tp_itemsize */
2547 /* methods */
2548 (destructor)permutations_dealloc, /* tp_dealloc */
2549 0, /* tp_print */
2550 0, /* tp_getattr */
2551 0, /* tp_setattr */
2552 0, /* tp_reserved */
2553 0, /* tp_repr */
2554 0, /* tp_as_number */
2555 0, /* tp_as_sequence */
2556 0, /* tp_as_mapping */
2557 0, /* tp_hash */
2558 0, /* tp_call */
2559 0, /* tp_str */
2560 PyObject_GenericGetAttr, /* tp_getattro */
2561 0, /* tp_setattro */
2562 0, /* tp_as_buffer */
2563 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2564 Py_TPFLAGS_BASETYPE, /* tp_flags */
2565 permutations_doc, /* tp_doc */
2566 (traverseproc)permutations_traverse, /* tp_traverse */
2567 0, /* tp_clear */
2568 0, /* tp_richcompare */
2569 0, /* tp_weaklistoffset */
2570 PyObject_SelfIter, /* tp_iter */
2571 (iternextfunc)permutations_next, /* tp_iternext */
2572 0, /* tp_methods */
2573 0, /* tp_members */
2574 0, /* tp_getset */
2575 0, /* tp_base */
2576 0, /* tp_dict */
2577 0, /* tp_descr_get */
2578 0, /* tp_descr_set */
2579 0, /* tp_dictoffset */
2580 0, /* tp_init */
2581 0, /* tp_alloc */
2582 permutations_new, /* tp_new */
2583 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002584};
2585
2586
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002587/* compress object ************************************************************/
2588
2589/* Equivalent to:
2590
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002591 def compress(data, selectors):
2592 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2593 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002594*/
2595
2596typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002597 PyObject_HEAD
2598 PyObject *data;
2599 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002600} compressobject;
2601
2602static PyTypeObject compress_type;
2603
2604static PyObject *
2605compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2606{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002607 PyObject *seq1, *seq2;
2608 PyObject *data=NULL, *selectors=NULL;
2609 compressobject *lz;
2610 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002611
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002612 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2613 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002614
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002615 data = PyObject_GetIter(seq1);
2616 if (data == NULL)
2617 goto fail;
2618 selectors = PyObject_GetIter(seq2);
2619 if (selectors == NULL)
2620 goto fail;
2621
2622 /* create compressobject structure */
2623 lz = (compressobject *)type->tp_alloc(type, 0);
2624 if (lz == NULL)
2625 goto fail;
2626 lz->data = data;
2627 lz->selectors = selectors;
2628 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002629
2630fail:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002631 Py_XDECREF(data);
2632 Py_XDECREF(selectors);
2633 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002634}
2635
2636static void
2637compress_dealloc(compressobject *lz)
2638{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002639 PyObject_GC_UnTrack(lz);
2640 Py_XDECREF(lz->data);
2641 Py_XDECREF(lz->selectors);
2642 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002643}
2644
2645static int
2646compress_traverse(compressobject *lz, visitproc visit, void *arg)
2647{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002648 Py_VISIT(lz->data);
2649 Py_VISIT(lz->selectors);
2650 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002651}
2652
2653static PyObject *
2654compress_next(compressobject *lz)
2655{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002656 PyObject *data = lz->data, *selectors = lz->selectors;
2657 PyObject *datum, *selector;
2658 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2659 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2660 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002661
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002662 while (1) {
2663 /* Steps: get datum, get selector, evaluate selector.
2664 Order is important (to match the pure python version
2665 in terms of which input gets a chance to raise an
2666 exception first).
2667 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002668
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002669 datum = datanext(data);
2670 if (datum == NULL)
2671 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002672
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002673 selector = selectornext(selectors);
2674 if (selector == NULL) {
2675 Py_DECREF(datum);
2676 return NULL;
2677 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002678
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002679 ok = PyObject_IsTrue(selector);
2680 Py_DECREF(selector);
2681 if (ok == 1)
2682 return datum;
2683 Py_DECREF(datum);
2684 if (ok == -1)
2685 return NULL;
2686 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002687}
2688
2689PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00002690"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002691\n\
2692Return data elements corresponding to true selector elements.\n\
2693Forms a shorter iterator from selected data elements using the\n\
2694selectors to choose the data elements.");
2695
2696static PyTypeObject compress_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002697 PyVarObject_HEAD_INIT(NULL, 0)
2698 "itertools.compress", /* tp_name */
2699 sizeof(compressobject), /* tp_basicsize */
2700 0, /* tp_itemsize */
2701 /* methods */
2702 (destructor)compress_dealloc, /* tp_dealloc */
2703 0, /* tp_print */
2704 0, /* tp_getattr */
2705 0, /* tp_setattr */
2706 0, /* tp_reserved */
2707 0, /* tp_repr */
2708 0, /* tp_as_number */
2709 0, /* tp_as_sequence */
2710 0, /* tp_as_mapping */
2711 0, /* tp_hash */
2712 0, /* tp_call */
2713 0, /* tp_str */
2714 PyObject_GenericGetAttr, /* tp_getattro */
2715 0, /* tp_setattro */
2716 0, /* tp_as_buffer */
2717 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2718 Py_TPFLAGS_BASETYPE, /* tp_flags */
2719 compress_doc, /* tp_doc */
2720 (traverseproc)compress_traverse, /* tp_traverse */
2721 0, /* tp_clear */
2722 0, /* tp_richcompare */
2723 0, /* tp_weaklistoffset */
2724 PyObject_SelfIter, /* tp_iter */
2725 (iternextfunc)compress_next, /* tp_iternext */
2726 0, /* tp_methods */
2727 0, /* tp_members */
2728 0, /* tp_getset */
2729 0, /* tp_base */
2730 0, /* tp_dict */
2731 0, /* tp_descr_get */
2732 0, /* tp_descr_set */
2733 0, /* tp_dictoffset */
2734 0, /* tp_init */
2735 0, /* tp_alloc */
2736 compress_new, /* tp_new */
2737 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002738};
2739
2740
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002741/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00002742
2743typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002744 PyObject_HEAD
2745 PyObject *func;
2746 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002747} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002748
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002749static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002750
2751static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002752filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002753{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002754 PyObject *func, *seq;
2755 PyObject *it;
2756 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002757
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002758 if (type == &filterfalse_type &&
2759 !_PyArg_NoKeywords("filterfalse()", kwds))
2760 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002761
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002762 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
2763 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002764
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002765 /* Get iterator. */
2766 it = PyObject_GetIter(seq);
2767 if (it == NULL)
2768 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002769
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002770 /* create filterfalseobject structure */
2771 lz = (filterfalseobject *)type->tp_alloc(type, 0);
2772 if (lz == NULL) {
2773 Py_DECREF(it);
2774 return NULL;
2775 }
2776 Py_INCREF(func);
2777 lz->func = func;
2778 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002779
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002780 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002781}
2782
2783static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002784filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002785{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002786 PyObject_GC_UnTrack(lz);
2787 Py_XDECREF(lz->func);
2788 Py_XDECREF(lz->it);
2789 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002790}
2791
2792static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002793filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002794{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002795 Py_VISIT(lz->it);
2796 Py_VISIT(lz->func);
2797 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002798}
2799
2800static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002801filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002802{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002803 PyObject *item;
2804 PyObject *it = lz->it;
2805 long ok;
2806 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002807
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002808 iternext = *Py_TYPE(it)->tp_iternext;
2809 for (;;) {
2810 item = iternext(it);
2811 if (item == NULL)
2812 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002813
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002814 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2815 ok = PyObject_IsTrue(item);
2816 } else {
2817 PyObject *good;
2818 good = PyObject_CallFunctionObjArgs(lz->func,
2819 item, NULL);
2820 if (good == NULL) {
2821 Py_DECREF(item);
2822 return NULL;
2823 }
2824 ok = PyObject_IsTrue(good);
2825 Py_DECREF(good);
2826 }
2827 if (!ok)
2828 return item;
2829 Py_DECREF(item);
2830 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00002831}
2832
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002833PyDoc_STRVAR(filterfalse_doc,
2834"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002835\n\
2836Return those items of sequence for which function(item) is false.\n\
2837If function is None, return the items that are false.");
2838
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002839static PyTypeObject filterfalse_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002840 PyVarObject_HEAD_INIT(NULL, 0)
2841 "itertools.filterfalse", /* tp_name */
2842 sizeof(filterfalseobject), /* tp_basicsize */
2843 0, /* tp_itemsize */
2844 /* methods */
2845 (destructor)filterfalse_dealloc, /* tp_dealloc */
2846 0, /* tp_print */
2847 0, /* tp_getattr */
2848 0, /* tp_setattr */
2849 0, /* tp_reserved */
2850 0, /* tp_repr */
2851 0, /* tp_as_number */
2852 0, /* tp_as_sequence */
2853 0, /* tp_as_mapping */
2854 0, /* tp_hash */
2855 0, /* tp_call */
2856 0, /* tp_str */
2857 PyObject_GenericGetAttr, /* tp_getattro */
2858 0, /* tp_setattro */
2859 0, /* tp_as_buffer */
2860 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2861 Py_TPFLAGS_BASETYPE, /* tp_flags */
2862 filterfalse_doc, /* tp_doc */
2863 (traverseproc)filterfalse_traverse, /* tp_traverse */
2864 0, /* tp_clear */
2865 0, /* tp_richcompare */
2866 0, /* tp_weaklistoffset */
2867 PyObject_SelfIter, /* tp_iter */
2868 (iternextfunc)filterfalse_next, /* tp_iternext */
2869 0, /* tp_methods */
2870 0, /* tp_members */
2871 0, /* tp_getset */
2872 0, /* tp_base */
2873 0, /* tp_dict */
2874 0, /* tp_descr_get */
2875 0, /* tp_descr_set */
2876 0, /* tp_dictoffset */
2877 0, /* tp_init */
2878 0, /* tp_alloc */
2879 filterfalse_new, /* tp_new */
2880 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002881};
2882
2883
2884/* count object ************************************************************/
2885
2886typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002887 PyObject_HEAD
2888 Py_ssize_t cnt;
2889 PyObject *long_cnt;
2890 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002891} countobject;
2892
Raymond Hettinger30729212009-02-12 06:28:27 +00002893/* Counting logic and invariants:
2894
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002895fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00002896
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002897 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
2898 Advances with: cnt += 1
2899 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00002900
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002901slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00002902
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002903 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
2904 All counting is done with python objects (no overflows or underflows).
2905 Advances with: long_cnt += long_step
2906 Step may be zero -- effectively a slow version of repeat(cnt).
2907 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00002908*/
2909
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002910static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002911
2912static PyObject *
2913count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2914{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002915 countobject *lz;
2916 int slow_mode = 0;
2917 Py_ssize_t cnt = 0;
2918 PyObject *long_cnt = NULL;
2919 PyObject *long_step = NULL;
2920 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002921
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002922 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
2923 kwlist, &long_cnt, &long_step))
2924 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002925
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002926 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
2927 (long_step != NULL && !PyNumber_Check(long_step))) {
2928 PyErr_SetString(PyExc_TypeError, "a number is required");
2929 return NULL;
2930 }
Raymond Hettinger30729212009-02-12 06:28:27 +00002931
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002932 if (long_cnt != NULL) {
2933 cnt = PyLong_AsSsize_t(long_cnt);
2934 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
2935 PyErr_Clear();
2936 slow_mode = 1;
2937 }
2938 Py_INCREF(long_cnt);
2939 } else {
2940 cnt = 0;
2941 long_cnt = PyLong_FromLong(0);
2942 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002943
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002944 /* If not specified, step defaults to 1 */
2945 if (long_step == NULL) {
2946 long_step = PyLong_FromLong(1);
2947 if (long_step == NULL) {
2948 Py_DECREF(long_cnt);
2949 return NULL;
2950 }
2951 } else
2952 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002953
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002954 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002955
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002956 /* Fast mode only works when the step is 1 */
2957 if (!PyLong_Check(long_step) ||
2958 PyLong_AS_LONG(long_step) != 1) {
2959 slow_mode = 1;
2960 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002961
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002962 if (slow_mode)
2963 cnt = PY_SSIZE_T_MAX;
2964 else
2965 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00002966
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002967 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
2968 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
2969 assert(slow_mode ||
2970 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002971
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002972 /* create countobject structure */
2973 lz = (countobject *)type->tp_alloc(type, 0);
2974 if (lz == NULL) {
2975 Py_XDECREF(long_cnt);
2976 return NULL;
2977 }
2978 lz->cnt = cnt;
2979 lz->long_cnt = long_cnt;
2980 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002981
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002982 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002983}
2984
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002985static void
2986count_dealloc(countobject *lz)
2987{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002988 PyObject_GC_UnTrack(lz);
2989 Py_XDECREF(lz->long_cnt);
2990 Py_XDECREF(lz->long_step);
2991 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00002992}
2993
Benjamin Peterson25e26a02009-02-16 21:28:29 +00002994static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00002995count_traverse(countobject *lz, visitproc visit, void *arg)
2996{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002997 Py_VISIT(lz->long_cnt);
2998 Py_VISIT(lz->long_step);
2999 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003000}
3001
3002static PyObject *
3003count_nextlong(countobject *lz)
3004{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003005 PyObject *long_cnt;
3006 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003007
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003008 long_cnt = lz->long_cnt;
3009 if (long_cnt == NULL) {
3010 /* Switch to slow_mode */
3011 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3012 if (long_cnt == NULL)
3013 return NULL;
3014 }
3015 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00003016
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003017 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3018 if (stepped_up == NULL)
3019 return NULL;
3020 lz->long_cnt = stepped_up;
3021 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003022}
3023
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003024static PyObject *
3025count_next(countobject *lz)
3026{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003027 if (lz->cnt == PY_SSIZE_T_MAX)
3028 return count_nextlong(lz);
3029 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003030}
3031
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003032static PyObject *
3033count_repr(countobject *lz)
3034{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003035 if (lz->cnt != PY_SSIZE_T_MAX)
3036 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003037
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003038 if (PyLong_Check(lz->long_step)) {
3039 long step = PyLong_AsLong(lz->long_step);
3040 if (step == -1 && PyErr_Occurred()) {
3041 PyErr_Clear();
3042 }
3043 if (step == 1) {
3044 /* Don't display step when it is an integer equal to 1 */
3045 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
3046 }
3047 }
3048 return PyUnicode_FromFormat("count(%R, %R)",
3049 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003050}
3051
Raymond Hettinger6c8ee7a2009-11-30 21:55:17 +00003052static PyObject *
3053count_reduce(countobject *lz)
3054{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003055 if (lz->cnt == PY_SSIZE_T_MAX)
3056 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3057 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettinger6c8ee7a2009-11-30 21:55:17 +00003058}
3059
3060PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3061
3062static PyMethodDef count_methods[] = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003063 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3064 count_reduce_doc},
3065 {NULL, NULL} /* sentinel */
Raymond Hettinger6c8ee7a2009-11-30 21:55:17 +00003066};
3067
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003068PyDoc_STRVAR(count_doc,
Ezio Melotti7a61e3c2010-06-11 02:28:37 +00003069 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003070\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003071Return a count object whose .__next__() method returns consecutive values.\n\
3072Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00003073 def count(firstval=0, step=1):\n\
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003074 x = firstval\n\
3075 while 1:\n\
3076 yield x\n\
3077 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003078
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003079static PyTypeObject count_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003080 PyVarObject_HEAD_INIT(NULL, 0)
3081 "itertools.count", /* tp_name */
3082 sizeof(countobject), /* tp_basicsize */
3083 0, /* tp_itemsize */
3084 /* methods */
3085 (destructor)count_dealloc, /* tp_dealloc */
3086 0, /* tp_print */
3087 0, /* tp_getattr */
3088 0, /* tp_setattr */
3089 0, /* tp_reserved */
3090 (reprfunc)count_repr, /* tp_repr */
3091 0, /* tp_as_number */
3092 0, /* tp_as_sequence */
3093 0, /* tp_as_mapping */
3094 0, /* tp_hash */
3095 0, /* tp_call */
3096 0, /* tp_str */
3097 PyObject_GenericGetAttr, /* tp_getattro */
3098 0, /* tp_setattro */
3099 0, /* tp_as_buffer */
3100 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3101 Py_TPFLAGS_BASETYPE, /* tp_flags */
3102 count_doc, /* tp_doc */
3103 (traverseproc)count_traverse, /* tp_traverse */
3104 0, /* tp_clear */
3105 0, /* tp_richcompare */
3106 0, /* tp_weaklistoffset */
3107 PyObject_SelfIter, /* tp_iter */
3108 (iternextfunc)count_next, /* tp_iternext */
3109 count_methods, /* tp_methods */
3110 0, /* tp_members */
3111 0, /* tp_getset */
3112 0, /* tp_base */
3113 0, /* tp_dict */
3114 0, /* tp_descr_get */
3115 0, /* tp_descr_set */
3116 0, /* tp_dictoffset */
3117 0, /* tp_init */
3118 0, /* tp_alloc */
3119 count_new, /* tp_new */
3120 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003121};
3122
3123
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003124/* repeat object ************************************************************/
3125
3126typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003127 PyObject_HEAD
3128 PyObject *element;
3129 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003130} repeatobject;
3131
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003132static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003133
3134static PyObject *
3135repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3136{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003137 repeatobject *ro;
3138 PyObject *element;
3139 Py_ssize_t cnt = -1;
3140 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003141
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003142 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3143 &element, &cnt))
3144 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003145
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003146 if (PyTuple_Size(args) == 2 && cnt < 0)
3147 cnt = 0;
3148
3149 ro = (repeatobject *)type->tp_alloc(type, 0);
3150 if (ro == NULL)
3151 return NULL;
3152 Py_INCREF(element);
3153 ro->element = element;
3154 ro->cnt = cnt;
3155 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003156}
3157
3158static void
3159repeat_dealloc(repeatobject *ro)
3160{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003161 PyObject_GC_UnTrack(ro);
3162 Py_XDECREF(ro->element);
3163 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003164}
3165
3166static int
3167repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3168{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003169 Py_VISIT(ro->element);
3170 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003171}
3172
3173static PyObject *
3174repeat_next(repeatobject *ro)
3175{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003176 if (ro->cnt == 0)
3177 return NULL;
3178 if (ro->cnt > 0)
3179 ro->cnt--;
3180 Py_INCREF(ro->element);
3181 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003182}
3183
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003184static PyObject *
3185repeat_repr(repeatobject *ro)
3186{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003187 if (ro->cnt == -1)
3188 return PyUnicode_FromFormat("repeat(%R)", ro->element);
3189 else
3190 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
3191}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003192
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003193static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003194repeat_len(repeatobject *ro)
3195{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003196 if (ro->cnt == -1) {
3197 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3198 return NULL;
3199 }
3200 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003201}
3202
Armin Rigof5b3e362006-02-11 21:32:43 +00003203PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003204
3205static PyMethodDef repeat_methods[] = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003206 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3207 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003208};
3209
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003210PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00003211"repeat(object [,times]) -> create an iterator which returns the object\n\
3212for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003213endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003214
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003215static PyTypeObject repeat_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003216 PyVarObject_HEAD_INIT(NULL, 0)
3217 "itertools.repeat", /* tp_name */
3218 sizeof(repeatobject), /* tp_basicsize */
3219 0, /* tp_itemsize */
3220 /* methods */
3221 (destructor)repeat_dealloc, /* tp_dealloc */
3222 0, /* tp_print */
3223 0, /* tp_getattr */
3224 0, /* tp_setattr */
3225 0, /* tp_reserved */
3226 (reprfunc)repeat_repr, /* tp_repr */
3227 0, /* tp_as_number */
3228 0, /* tp_as_sequence */
3229 0, /* tp_as_mapping */
3230 0, /* tp_hash */
3231 0, /* tp_call */
3232 0, /* tp_str */
3233 PyObject_GenericGetAttr, /* tp_getattro */
3234 0, /* tp_setattro */
3235 0, /* tp_as_buffer */
3236 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3237 Py_TPFLAGS_BASETYPE, /* tp_flags */
3238 repeat_doc, /* tp_doc */
3239 (traverseproc)repeat_traverse, /* tp_traverse */
3240 0, /* tp_clear */
3241 0, /* tp_richcompare */
3242 0, /* tp_weaklistoffset */
3243 PyObject_SelfIter, /* tp_iter */
3244 (iternextfunc)repeat_next, /* tp_iternext */
3245 repeat_methods, /* tp_methods */
3246 0, /* tp_members */
3247 0, /* tp_getset */
3248 0, /* tp_base */
3249 0, /* tp_dict */
3250 0, /* tp_descr_get */
3251 0, /* tp_descr_set */
3252 0, /* tp_dictoffset */
3253 0, /* tp_init */
3254 0, /* tp_alloc */
3255 repeat_new, /* tp_new */
3256 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003257};
3258
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003259/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00003260
3261#include "Python.h"
3262
3263typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003264 PyObject_HEAD
3265 Py_ssize_t tuplesize;
3266 Py_ssize_t numactive;
3267 PyObject *ittuple; /* tuple of iterators */
3268 PyObject *result;
3269 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003270} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003271
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003272static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003273
3274static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003275zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003276{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003277 ziplongestobject *lz;
3278 Py_ssize_t i;
3279 PyObject *ittuple; /* tuple of iterators */
3280 PyObject *result;
3281 PyObject *fillvalue = Py_None;
3282 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003283
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003284 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3285 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3286 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3287 PyErr_SetString(PyExc_TypeError,
3288 "zip_longest() got an unexpected keyword argument");
3289 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003290 }
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003291 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003292
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003293 /* args must be a tuple */
3294 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00003295
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003296 /* obtain iterators */
3297 ittuple = PyTuple_New(tuplesize);
3298 if (ittuple == NULL)
3299 return NULL;
3300 for (i=0; i < tuplesize; ++i) {
3301 PyObject *item = PyTuple_GET_ITEM(args, i);
3302 PyObject *it = PyObject_GetIter(item);
3303 if (it == NULL) {
3304 if (PyErr_ExceptionMatches(PyExc_TypeError))
3305 PyErr_Format(PyExc_TypeError,
3306 "zip_longest argument #%zd must support iteration",
3307 i+1);
3308 Py_DECREF(ittuple);
3309 return NULL;
3310 }
3311 PyTuple_SET_ITEM(ittuple, i, it);
3312 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003313
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003314 /* create a result holder */
3315 result = PyTuple_New(tuplesize);
3316 if (result == NULL) {
3317 Py_DECREF(ittuple);
3318 return NULL;
3319 }
3320 for (i=0 ; i < tuplesize ; i++) {
3321 Py_INCREF(Py_None);
3322 PyTuple_SET_ITEM(result, i, Py_None);
3323 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003324
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003325 /* create ziplongestobject structure */
3326 lz = (ziplongestobject *)type->tp_alloc(type, 0);
3327 if (lz == NULL) {
3328 Py_DECREF(ittuple);
3329 Py_DECREF(result);
3330 return NULL;
3331 }
3332 lz->ittuple = ittuple;
3333 lz->tuplesize = tuplesize;
3334 lz->numactive = tuplesize;
3335 lz->result = result;
3336 Py_INCREF(fillvalue);
3337 lz->fillvalue = fillvalue;
3338 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003339}
3340
3341static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003342zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003343{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003344 PyObject_GC_UnTrack(lz);
3345 Py_XDECREF(lz->ittuple);
3346 Py_XDECREF(lz->result);
3347 Py_XDECREF(lz->fillvalue);
3348 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003349}
3350
3351static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003352zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003353{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003354 Py_VISIT(lz->ittuple);
3355 Py_VISIT(lz->result);
3356 Py_VISIT(lz->fillvalue);
3357 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003358}
3359
3360static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003361zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003362{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003363 Py_ssize_t i;
3364 Py_ssize_t tuplesize = lz->tuplesize;
3365 PyObject *result = lz->result;
3366 PyObject *it;
3367 PyObject *item;
3368 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003369
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003370 if (tuplesize == 0)
3371 return NULL;
3372 if (lz->numactive == 0)
3373 return NULL;
3374 if (Py_REFCNT(result) == 1) {
3375 Py_INCREF(result);
3376 for (i=0 ; i < tuplesize ; i++) {
3377 it = PyTuple_GET_ITEM(lz->ittuple, i);
3378 if (it == NULL) {
3379 Py_INCREF(lz->fillvalue);
3380 item = lz->fillvalue;
3381 } else {
3382 item = PyIter_Next(it);
3383 if (item == NULL) {
3384 lz->numactive -= 1;
3385 if (lz->numactive == 0 || PyErr_Occurred()) {
3386 lz->numactive = 0;
3387 Py_DECREF(result);
Raymond Hettingera9311a32009-11-01 21:02:38 +00003388 return NULL;
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003389 } else {
3390 Py_INCREF(lz->fillvalue);
3391 item = lz->fillvalue;
3392 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3393 Py_DECREF(it);
3394 }
Raymond Hettingera9311a32009-11-01 21:02:38 +00003395 }
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003396 }
3397 olditem = PyTuple_GET_ITEM(result, i);
3398 PyTuple_SET_ITEM(result, i, item);
3399 Py_DECREF(olditem);
Raymond Hettingera9311a32009-11-01 21:02:38 +00003400 }
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003401 } else {
3402 result = PyTuple_New(tuplesize);
3403 if (result == NULL)
3404 return NULL;
3405 for (i=0 ; i < tuplesize ; i++) {
3406 it = PyTuple_GET_ITEM(lz->ittuple, i);
3407 if (it == NULL) {
3408 Py_INCREF(lz->fillvalue);
3409 item = lz->fillvalue;
3410 } else {
3411 item = PyIter_Next(it);
3412 if (item == NULL) {
3413 lz->numactive -= 1;
3414 if (lz->numactive == 0 || PyErr_Occurred()) {
3415 lz->numactive = 0;
3416 Py_DECREF(result);
3417 return NULL;
3418 } else {
3419 Py_INCREF(lz->fillvalue);
3420 item = lz->fillvalue;
3421 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3422 Py_DECREF(it);
3423 }
3424 }
3425 }
3426 PyTuple_SET_ITEM(result, i, item);
3427 }
3428 }
3429 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003430}
3431
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003432PyDoc_STRVAR(zip_longest_doc,
3433"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003434\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003435Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00003436the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003437method continues until the longest iterable in the argument sequence\n\
3438is exhausted and then it raises StopIteration. When the shorter iterables\n\
3439are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3440defaults to None or can be specified by a keyword argument.\n\
3441");
3442
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003443static PyTypeObject ziplongest_type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003444 PyVarObject_HEAD_INIT(NULL, 0)
3445 "itertools.zip_longest", /* tp_name */
3446 sizeof(ziplongestobject), /* tp_basicsize */
3447 0, /* tp_itemsize */
3448 /* methods */
3449 (destructor)zip_longest_dealloc, /* tp_dealloc */
3450 0, /* tp_print */
3451 0, /* tp_getattr */
3452 0, /* tp_setattr */
3453 0, /* tp_reserved */
3454 0, /* tp_repr */
3455 0, /* tp_as_number */
3456 0, /* tp_as_sequence */
3457 0, /* tp_as_mapping */
3458 0, /* tp_hash */
3459 0, /* tp_call */
3460 0, /* tp_str */
3461 PyObject_GenericGetAttr, /* tp_getattro */
3462 0, /* tp_setattro */
3463 0, /* tp_as_buffer */
3464 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3465 Py_TPFLAGS_BASETYPE, /* tp_flags */
3466 zip_longest_doc, /* tp_doc */
3467 (traverseproc)zip_longest_traverse, /* tp_traverse */
3468 0, /* tp_clear */
3469 0, /* tp_richcompare */
3470 0, /* tp_weaklistoffset */
3471 PyObject_SelfIter, /* tp_iter */
3472 (iternextfunc)zip_longest_next, /* tp_iternext */
3473 0, /* tp_methods */
3474 0, /* tp_members */
3475 0, /* tp_getset */
3476 0, /* tp_base */
3477 0, /* tp_dict */
3478 0, /* tp_descr_get */
3479 0, /* tp_descr_set */
3480 0, /* tp_dictoffset */
3481 0, /* tp_init */
3482 0, /* tp_alloc */
3483 zip_longest_new, /* tp_new */
3484 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003485};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003486
3487/* module level code ********************************************************/
3488
3489PyDoc_STRVAR(module_doc,
3490"Functional tools for creating and using iterators.\n\
3491\n\
3492Infinite iterators:\n\
3493count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003494cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003495repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003496\n\
3497Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003498chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
3499compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
3500dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
3501groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003502filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003503islice(seq, [start,] stop [, step]) --> elements from\n\
3504 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003505starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003506tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003507takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003508zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00003509\n\
3510Combinatoric generators:\n\
3511product(p, q, ... [repeat=1]) --> cartesian product\n\
3512permutations(p[, r])\n\
Raymond Hettingerd3a77c02009-11-19 01:23:41 +00003513combinations(p, r)\n\
3514combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003515");
3516
3517
Raymond Hettingerad983e72003-11-12 14:32:26 +00003518static PyMethodDef module_methods[] = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003519 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3520 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00003521};
3522
Martin v. Löwis1a214512008-06-11 05:26:20 +00003523
3524static struct PyModuleDef itertoolsmodule = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003525 PyModuleDef_HEAD_INIT,
3526 "itertools",
3527 module_doc,
3528 -1,
3529 module_methods,
3530 NULL,
3531 NULL,
3532 NULL,
3533 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00003534};
3535
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003536PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00003537PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003538{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003539 int i;
3540 PyObject *m;
3541 char *name;
3542 PyTypeObject *typelist[] = {
3543 &combinations_type,
3544 &cwr_type,
3545 &cycle_type,
3546 &dropwhile_type,
3547 &takewhile_type,
3548 &islice_type,
3549 &starmap_type,
3550 &chain_type,
3551 &compress_type,
3552 &filterfalse_type,
3553 &count_type,
3554 &ziplongest_type,
3555 &permutations_type,
3556 &product_type,
3557 &repeat_type,
3558 &groupby_type,
3559 NULL
3560 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00003561
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003562 Py_TYPE(&teedataobject_type) = &PyType_Type;
3563 m = PyModule_Create(&itertoolsmodule);
3564 if (m == NULL)
3565 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003566
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003567 for (i=0 ; typelist[i] != NULL ; i++) {
3568 if (PyType_Ready(typelist[i]) < 0)
3569 return NULL;
3570 name = strchr(typelist[i]->tp_name, '.');
3571 assert (name != NULL);
3572 Py_INCREF(typelist[i]);
3573 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
3574 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003575
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00003576 if (PyType_Ready(&teedataobject_type) < 0)
3577 return NULL;
3578 if (PyType_Ready(&tee_type) < 0)
3579 return NULL;
3580 if (PyType_Ready(&_grouper_type) < 0)
3581 return NULL;
3582 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003583}