blob: 52a38a55df3aedede10d6a0fcf4358e0af4bc1f6 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
2#include "Python.h"
Fred Drake08ebfec2004-10-17 19:36:57 +00003#include "structmember.h"
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00005/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00006 by Raymond D. Hettinger <python@rcn.com>
7 Copyright (c) 2003 Python Software Foundation.
8 All rights reserved.
9*/
10
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000011
12/* groupby object ***********************************************************/
13
14typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000015 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000021} groupbyobject;
22
23static PyTypeObject groupby_type;
24static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26static PyObject *
27groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000029 static char *kwargs[] = {"iterable", "key", NULL};
30 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000032
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000033 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
34 &it, &keyfunc))
35 return NULL;
36
37 gbo = (groupbyobject *)type->tp_alloc(type, 0);
38 if (gbo == NULL)
39 return NULL;
40 gbo->tgtkey = NULL;
41 gbo->currkey = NULL;
42 gbo->currvalue = NULL;
43 gbo->keyfunc = keyfunc;
44 Py_INCREF(keyfunc);
45 gbo->it = PyObject_GetIter(it);
46 if (gbo->it == NULL) {
47 Py_DECREF(gbo);
48 return NULL;
49 }
50 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000051}
52
53static void
54groupby_dealloc(groupbyobject *gbo)
55{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000056 PyObject_GC_UnTrack(gbo);
57 Py_XDECREF(gbo->it);
58 Py_XDECREF(gbo->keyfunc);
59 Py_XDECREF(gbo->tgtkey);
60 Py_XDECREF(gbo->currkey);
61 Py_XDECREF(gbo->currvalue);
62 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000063}
64
65static int
66groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
67{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000068 Py_VISIT(gbo->it);
69 Py_VISIT(gbo->keyfunc);
70 Py_VISIT(gbo->tgtkey);
71 Py_VISIT(gbo->currkey);
72 Py_VISIT(gbo->currvalue);
73 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000074}
75
76static PyObject *
77groupby_next(groupbyobject *gbo)
78{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000081 /* skip to next iteration group */
82 for (;;) {
83 if (gbo->currkey == NULL)
84 /* pass */;
85 else if (gbo->tgtkey == NULL)
86 break;
87 else {
88 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000089
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000090 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
91 gbo->currkey, Py_EQ);
92 if (rcmp == -1)
93 return NULL;
94 else if (rcmp == 0)
95 break;
96 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000097
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000098 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000101
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000102 if (gbo->keyfunc == Py_None) {
103 newkey = newvalue;
104 Py_INCREF(newvalue);
105 } else {
106 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
107 newvalue, NULL);
108 if (newkey == NULL) {
109 Py_DECREF(newvalue);
110 return NULL;
111 }
112 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000113
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000114 tmp = gbo->currkey;
115 gbo->currkey = newkey;
116 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000117
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000118 tmp = gbo->currvalue;
119 gbo->currvalue = newvalue;
120 Py_XDECREF(tmp);
121 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000123 Py_INCREF(gbo->currkey);
124 tmp = gbo->tgtkey;
125 gbo->tgtkey = gbo->currkey;
126 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000127
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000131
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000135}
136
137PyDoc_STRVAR(groupby_doc,
138"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
139(key, sub-iterator) grouped by each value of key(value).\n");
140
141static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000142 PyVarObject_HEAD_INIT(NULL, 0)
143 "itertools.groupby", /* tp_name */
144 sizeof(groupbyobject), /* tp_basicsize */
145 0, /* tp_itemsize */
146 /* methods */
147 (destructor)groupby_dealloc, /* tp_dealloc */
148 0, /* tp_print */
149 0, /* tp_getattr */
150 0, /* tp_setattr */
151 0, /* tp_reserved */
152 0, /* tp_repr */
153 0, /* tp_as_number */
154 0, /* tp_as_sequence */
155 0, /* tp_as_mapping */
156 0, /* tp_hash */
157 0, /* tp_call */
158 0, /* tp_str */
159 PyObject_GenericGetAttr, /* tp_getattro */
160 0, /* tp_setattro */
161 0, /* tp_as_buffer */
162 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
163 Py_TPFLAGS_BASETYPE, /* tp_flags */
164 groupby_doc, /* tp_doc */
165 (traverseproc)groupby_traverse, /* tp_traverse */
166 0, /* tp_clear */
167 0, /* tp_richcompare */
168 0, /* tp_weaklistoffset */
169 PyObject_SelfIter, /* tp_iter */
170 (iternextfunc)groupby_next, /* tp_iternext */
171 0, /* tp_methods */
172 0, /* tp_members */
173 0, /* tp_getset */
174 0, /* tp_base */
175 0, /* tp_dict */
176 0, /* tp_descr_get */
177 0, /* tp_descr_set */
178 0, /* tp_dictoffset */
179 0, /* tp_init */
180 0, /* tp_alloc */
181 groupby_new, /* tp_new */
182 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000183};
184
185
186/* _grouper object (internal) ************************************************/
187
188typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000189 PyObject_HEAD
190 PyObject *parent;
191 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000192} _grouperobject;
193
194static PyTypeObject _grouper_type;
195
196static PyObject *
197_grouper_create(groupbyobject *parent, PyObject *tgtkey)
198{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000199 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000200
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000201 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
202 if (igo == NULL)
203 return NULL;
204 igo->parent = (PyObject *)parent;
205 Py_INCREF(parent);
206 igo->tgtkey = tgtkey;
207 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000208
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000209 PyObject_GC_Track(igo);
210 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000211}
212
213static void
214_grouper_dealloc(_grouperobject *igo)
215{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000216 PyObject_GC_UnTrack(igo);
217 Py_DECREF(igo->parent);
218 Py_DECREF(igo->tgtkey);
219 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000220}
221
222static int
223_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
224{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000225 Py_VISIT(igo->parent);
226 Py_VISIT(igo->tgtkey);
227 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000228}
229
230static PyObject *
231_grouper_next(_grouperobject *igo)
232{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000233 groupbyobject *gbo = (groupbyobject *)igo->parent;
234 PyObject *newvalue, *newkey, *r;
235 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000236
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000237 if (gbo->currvalue == NULL) {
238 newvalue = PyIter_Next(gbo->it);
239 if (newvalue == NULL)
240 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000241
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000242 if (gbo->keyfunc == Py_None) {
243 newkey = newvalue;
244 Py_INCREF(newvalue);
245 } else {
246 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
247 newvalue, NULL);
248 if (newkey == NULL) {
249 Py_DECREF(newvalue);
250 return NULL;
251 }
252 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000253
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000254 assert(gbo->currkey == NULL);
255 gbo->currkey = newkey;
256 gbo->currvalue = newvalue;
257 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000259 assert(gbo->currkey != NULL);
260 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
261 if (rcmp <= 0)
262 /* got any error or current group is end */
263 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000264
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000265 r = gbo->currvalue;
266 gbo->currvalue = NULL;
267 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000268
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000269 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000270}
271
272static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000273 PyVarObject_HEAD_INIT(NULL, 0)
274 "itertools._grouper", /* tp_name */
275 sizeof(_grouperobject), /* tp_basicsize */
276 0, /* tp_itemsize */
277 /* methods */
278 (destructor)_grouper_dealloc, /* tp_dealloc */
279 0, /* tp_print */
280 0, /* tp_getattr */
281 0, /* tp_setattr */
282 0, /* tp_reserved */
283 0, /* tp_repr */
284 0, /* tp_as_number */
285 0, /* tp_as_sequence */
286 0, /* tp_as_mapping */
287 0, /* tp_hash */
288 0, /* tp_call */
289 0, /* tp_str */
290 PyObject_GenericGetAttr, /* tp_getattro */
291 0, /* tp_setattro */
292 0, /* tp_as_buffer */
293 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
294 0, /* tp_doc */
295 (traverseproc)_grouper_traverse,/* tp_traverse */
296 0, /* tp_clear */
297 0, /* tp_richcompare */
298 0, /* tp_weaklistoffset */
299 PyObject_SelfIter, /* tp_iter */
300 (iternextfunc)_grouper_next, /* tp_iternext */
301 0, /* tp_methods */
302 0, /* tp_members */
303 0, /* tp_getset */
304 0, /* tp_base */
305 0, /* tp_dict */
306 0, /* tp_descr_get */
307 0, /* tp_descr_set */
308 0, /* tp_dictoffset */
309 0, /* tp_init */
310 0, /* tp_alloc */
311 0, /* tp_new */
312 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000313};
314
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000315
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000316
Raymond Hettingerad983e72003-11-12 14:32:26 +0000317/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000318
Raymond Hettingerad983e72003-11-12 14:32:26 +0000319/* The teedataobject pre-allocates space for LINKCELLS number of objects.
320 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000321 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000322 the other structure members including PyHEAD overhead. The larger the
323 value, the less memory overhead per object and the less time spent
324 allocating/deallocating new links. The smaller the number, the less
325 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000326*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000327#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000328
329typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000330 PyObject_HEAD
331 PyObject *it;
332 int numread;
333 PyObject *nextlink;
334 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000335} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000336
337typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000338 PyObject_HEAD
339 teedataobject *dataobj;
340 int index;
341 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000342} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000343
Raymond Hettingerad983e72003-11-12 14:32:26 +0000344static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000345
346static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000347teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000348{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000349 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000350
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000351 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
352 if (tdo == NULL)
353 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000354
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000355 tdo->numread = 0;
356 tdo->nextlink = NULL;
357 Py_INCREF(it);
358 tdo->it = it;
359 PyObject_GC_Track(tdo);
360 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000361}
362
363static PyObject *
364teedataobject_jumplink(teedataobject *tdo)
365{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000366 if (tdo->nextlink == NULL)
367 tdo->nextlink = teedataobject_new(tdo->it);
368 Py_XINCREF(tdo->nextlink);
369 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000370}
371
372static PyObject *
373teedataobject_getitem(teedataobject *tdo, int i)
374{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000375 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000376
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000377 assert(i < LINKCELLS);
378 if (i < tdo->numread)
379 value = tdo->values[i];
380 else {
381 /* this is the lead iterator, so fetch more data */
382 assert(i == tdo->numread);
383 value = PyIter_Next(tdo->it);
384 if (value == NULL)
385 return NULL;
386 tdo->numread++;
387 tdo->values[i] = value;
388 }
389 Py_INCREF(value);
390 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000391}
392
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000393static int
394teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
395{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000396 int i;
397 Py_VISIT(tdo->it);
398 for (i = 0; i < tdo->numread; i++)
399 Py_VISIT(tdo->values[i]);
400 Py_VISIT(tdo->nextlink);
401 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000402}
403
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200404static void
405teedataobject_safe_decref(PyObject *obj)
406{
407 while (obj && Py_TYPE(obj) == &teedataobject_type &&
408 Py_REFCNT(obj) == 1) {
409 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
410 ((teedataobject *)obj)->nextlink = NULL;
411 Py_DECREF(obj);
412 obj = nextlink;
413 }
414 Py_XDECREF(obj);
415}
416
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000417static int
418teedataobject_clear(teedataobject *tdo)
419{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000420 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200421 PyObject *tmp;
422
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000423 Py_CLEAR(tdo->it);
424 for (i=0 ; i<tdo->numread ; i++)
425 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200426 tmp = tdo->nextlink;
427 tdo->nextlink = NULL;
428 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000429 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000430}
431
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000432static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000433teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000434{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000435 PyObject_GC_UnTrack(tdo);
436 teedataobject_clear(tdo);
437 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000438}
439
Raymond Hettingerad983e72003-11-12 14:32:26 +0000440PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000441
Raymond Hettingerad983e72003-11-12 14:32:26 +0000442static PyTypeObject teedataobject_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000443 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
444 "itertools.tee_dataobject", /* tp_name */
445 sizeof(teedataobject), /* tp_basicsize */
446 0, /* tp_itemsize */
447 /* methods */
448 (destructor)teedataobject_dealloc, /* tp_dealloc */
449 0, /* tp_print */
450 0, /* tp_getattr */
451 0, /* tp_setattr */
452 0, /* tp_reserved */
453 0, /* tp_repr */
454 0, /* tp_as_number */
455 0, /* tp_as_sequence */
456 0, /* tp_as_mapping */
457 0, /* tp_hash */
458 0, /* tp_call */
459 0, /* tp_str */
460 PyObject_GenericGetAttr, /* tp_getattro */
461 0, /* tp_setattro */
462 0, /* tp_as_buffer */
463 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
464 teedataobject_doc, /* tp_doc */
465 (traverseproc)teedataobject_traverse, /* tp_traverse */
466 (inquiry)teedataobject_clear, /* tp_clear */
467 0, /* tp_richcompare */
468 0, /* tp_weaklistoffset */
469 0, /* tp_iter */
470 0, /* tp_iternext */
471 0, /* tp_methods */
472 0, /* tp_members */
473 0, /* tp_getset */
474 0, /* tp_base */
475 0, /* tp_dict */
476 0, /* tp_descr_get */
477 0, /* tp_descr_set */
478 0, /* tp_dictoffset */
479 0, /* tp_init */
480 0, /* tp_alloc */
481 0, /* tp_new */
482 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000483};
484
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000485
486static PyTypeObject tee_type;
487
488static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000489tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000490{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000491 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000492
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000493 if (to->index >= LINKCELLS) {
494 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200495 if (link == NULL)
496 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000497 Py_DECREF(to->dataobj);
498 to->dataobj = (teedataobject *)link;
499 to->index = 0;
500 }
501 value = teedataobject_getitem(to->dataobj, to->index);
502 if (value == NULL)
503 return NULL;
504 to->index++;
505 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000506}
507
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000508static int
509tee_traverse(teeobject *to, visitproc visit, void *arg)
510{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000511 Py_VISIT((PyObject *)to->dataobj);
512 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000513}
514
Raymond Hettingerad983e72003-11-12 14:32:26 +0000515static PyObject *
516tee_copy(teeobject *to)
517{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000518 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000519
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000520 newto = PyObject_GC_New(teeobject, &tee_type);
521 if (newto == NULL)
522 return NULL;
523 Py_INCREF(to->dataobj);
524 newto->dataobj = to->dataobj;
525 newto->index = to->index;
526 newto->weakreflist = NULL;
527 PyObject_GC_Track(newto);
528 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000529}
530
531PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
532
533static PyObject *
534tee_fromiterable(PyObject *iterable)
535{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000536 teeobject *to;
537 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000538
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000539 it = PyObject_GetIter(iterable);
540 if (it == NULL)
541 return NULL;
542 if (PyObject_TypeCheck(it, &tee_type)) {
543 to = (teeobject *)tee_copy((teeobject *)it);
544 goto done;
545 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000546
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000547 to = PyObject_GC_New(teeobject, &tee_type);
548 if (to == NULL)
549 goto done;
550 to->dataobj = (teedataobject *)teedataobject_new(it);
551 if (!to->dataobj) {
552 PyObject_GC_Del(to);
553 to = NULL;
554 goto done;
555 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000556
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000557 to->index = 0;
558 to->weakreflist = NULL;
559 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000560done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000561 Py_XDECREF(it);
562 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000563}
564
565static PyObject *
566tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
567{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000568 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000569
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000570 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
571 return NULL;
572 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000573}
574
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000575static int
576tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000577{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000578 if (to->weakreflist != NULL)
579 PyObject_ClearWeakRefs((PyObject *) to);
580 Py_CLEAR(to->dataobj);
581 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000582}
583
584static void
585tee_dealloc(teeobject *to)
586{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000587 PyObject_GC_UnTrack(to);
588 tee_clear(to);
589 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000590}
591
Raymond Hettingerad983e72003-11-12 14:32:26 +0000592PyDoc_STRVAR(teeobject_doc,
593"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000594
Raymond Hettingerad983e72003-11-12 14:32:26 +0000595static PyMethodDef tee_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000596 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
597 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000598};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000599
600static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000601 PyVarObject_HEAD_INIT(NULL, 0)
602 "itertools.tee", /* tp_name */
603 sizeof(teeobject), /* tp_basicsize */
604 0, /* tp_itemsize */
605 /* methods */
606 (destructor)tee_dealloc, /* tp_dealloc */
607 0, /* tp_print */
608 0, /* tp_getattr */
609 0, /* tp_setattr */
610 0, /* tp_reserved */
611 0, /* tp_repr */
612 0, /* tp_as_number */
613 0, /* tp_as_sequence */
614 0, /* tp_as_mapping */
615 0, /* tp_hash */
616 0, /* tp_call */
617 0, /* tp_str */
618 0, /* tp_getattro */
619 0, /* tp_setattro */
620 0, /* tp_as_buffer */
621 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
622 teeobject_doc, /* tp_doc */
623 (traverseproc)tee_traverse, /* tp_traverse */
624 (inquiry)tee_clear, /* tp_clear */
625 0, /* tp_richcompare */
626 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
627 PyObject_SelfIter, /* tp_iter */
628 (iternextfunc)tee_next, /* tp_iternext */
629 tee_methods, /* tp_methods */
630 0, /* tp_members */
631 0, /* tp_getset */
632 0, /* tp_base */
633 0, /* tp_dict */
634 0, /* tp_descr_get */
635 0, /* tp_descr_set */
636 0, /* tp_dictoffset */
637 0, /* tp_init */
638 0, /* tp_alloc */
639 tee_new, /* tp_new */
640 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000641};
642
Raymond Hettingerad983e72003-11-12 14:32:26 +0000643static PyObject *
644tee(PyObject *self, PyObject *args)
645{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000646 Py_ssize_t i, n=2;
647 PyObject *it, *iterable, *copyable, *result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000648
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000649 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
650 return NULL;
651 if (n < 0) {
652 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
653 return NULL;
654 }
655 result = PyTuple_New(n);
656 if (result == NULL)
657 return NULL;
658 if (n == 0)
659 return result;
660 it = PyObject_GetIter(iterable);
661 if (it == NULL) {
662 Py_DECREF(result);
663 return NULL;
664 }
665 if (!PyObject_HasAttrString(it, "__copy__")) {
666 copyable = tee_fromiterable(it);
667 Py_DECREF(it);
668 if (copyable == NULL) {
669 Py_DECREF(result);
670 return NULL;
671 }
672 } else
673 copyable = it;
674 PyTuple_SET_ITEM(result, 0, copyable);
675 for (i=1 ; i<n ; i++) {
676 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
677 if (copyable == NULL) {
678 Py_DECREF(result);
679 return NULL;
680 }
681 PyTuple_SET_ITEM(result, i, copyable);
682 }
683 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000684}
685
686PyDoc_STRVAR(tee_doc,
687"tee(iterable, n=2) --> tuple of n independent iterators.");
688
689
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000690/* cycle object **********************************************************/
691
692typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000693 PyObject_HEAD
694 PyObject *it;
695 PyObject *saved;
696 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000697} cycleobject;
698
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000699static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000700
701static PyObject *
702cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
703{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000704 PyObject *it;
705 PyObject *iterable;
706 PyObject *saved;
707 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000708
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000709 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
710 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000711
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000712 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
713 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000714
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000715 /* Get iterator. */
716 it = PyObject_GetIter(iterable);
717 if (it == NULL)
718 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000719
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000720 saved = PyList_New(0);
721 if (saved == NULL) {
722 Py_DECREF(it);
723 return NULL;
724 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000725
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000726 /* create cycleobject structure */
727 lz = (cycleobject *)type->tp_alloc(type, 0);
728 if (lz == NULL) {
729 Py_DECREF(it);
730 Py_DECREF(saved);
731 return NULL;
732 }
733 lz->it = it;
734 lz->saved = saved;
735 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000736
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000737 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000738}
739
740static void
741cycle_dealloc(cycleobject *lz)
742{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000743 PyObject_GC_UnTrack(lz);
744 Py_XDECREF(lz->saved);
745 Py_XDECREF(lz->it);
746 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000747}
748
749static int
750cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
751{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000752 Py_VISIT(lz->it);
753 Py_VISIT(lz->saved);
754 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000755}
756
757static PyObject *
758cycle_next(cycleobject *lz)
759{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000760 PyObject *item;
761 PyObject *it;
762 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000763
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000764 while (1) {
765 item = PyIter_Next(lz->it);
766 if (item != NULL) {
767 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
768 Py_DECREF(item);
769 return NULL;
770 }
771 return item;
772 }
773 if (PyErr_Occurred()) {
774 if (PyErr_ExceptionMatches(PyExc_StopIteration))
775 PyErr_Clear();
776 else
777 return NULL;
778 }
779 if (PyList_Size(lz->saved) == 0)
780 return NULL;
781 it = PyObject_GetIter(lz->saved);
782 if (it == NULL)
783 return NULL;
784 tmp = lz->it;
785 lz->it = it;
786 lz->firstpass = 1;
787 Py_DECREF(tmp);
788 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000789}
790
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000791PyDoc_STRVAR(cycle_doc,
792"cycle(iterable) --> cycle object\n\
793\n\
794Return elements from the iterable until it is exhausted.\n\
795Then repeat the sequence indefinitely.");
796
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000797static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000798 PyVarObject_HEAD_INIT(NULL, 0)
799 "itertools.cycle", /* tp_name */
800 sizeof(cycleobject), /* tp_basicsize */
801 0, /* tp_itemsize */
802 /* methods */
803 (destructor)cycle_dealloc, /* tp_dealloc */
804 0, /* tp_print */
805 0, /* tp_getattr */
806 0, /* tp_setattr */
807 0, /* tp_reserved */
808 0, /* tp_repr */
809 0, /* tp_as_number */
810 0, /* tp_as_sequence */
811 0, /* tp_as_mapping */
812 0, /* tp_hash */
813 0, /* tp_call */
814 0, /* tp_str */
815 PyObject_GenericGetAttr, /* tp_getattro */
816 0, /* tp_setattro */
817 0, /* tp_as_buffer */
818 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
819 Py_TPFLAGS_BASETYPE, /* tp_flags */
820 cycle_doc, /* tp_doc */
821 (traverseproc)cycle_traverse, /* tp_traverse */
822 0, /* tp_clear */
823 0, /* tp_richcompare */
824 0, /* tp_weaklistoffset */
825 PyObject_SelfIter, /* tp_iter */
826 (iternextfunc)cycle_next, /* tp_iternext */
827 0, /* tp_methods */
828 0, /* tp_members */
829 0, /* tp_getset */
830 0, /* tp_base */
831 0, /* tp_dict */
832 0, /* tp_descr_get */
833 0, /* tp_descr_set */
834 0, /* tp_dictoffset */
835 0, /* tp_init */
836 0, /* tp_alloc */
837 cycle_new, /* tp_new */
838 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000839};
840
841
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000842/* dropwhile object **********************************************************/
843
844typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000845 PyObject_HEAD
846 PyObject *func;
847 PyObject *it;
848 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000849} dropwhileobject;
850
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000851static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000852
853static PyObject *
854dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
855{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000856 PyObject *func, *seq;
857 PyObject *it;
858 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000859
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000860 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
861 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000862
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000863 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
864 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000865
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000866 /* Get iterator. */
867 it = PyObject_GetIter(seq);
868 if (it == NULL)
869 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000870
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000871 /* create dropwhileobject structure */
872 lz = (dropwhileobject *)type->tp_alloc(type, 0);
873 if (lz == NULL) {
874 Py_DECREF(it);
875 return NULL;
876 }
877 Py_INCREF(func);
878 lz->func = func;
879 lz->it = it;
880 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000881
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000882 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000883}
884
885static void
886dropwhile_dealloc(dropwhileobject *lz)
887{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000888 PyObject_GC_UnTrack(lz);
889 Py_XDECREF(lz->func);
890 Py_XDECREF(lz->it);
891 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000892}
893
894static int
895dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
896{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000897 Py_VISIT(lz->it);
898 Py_VISIT(lz->func);
899 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000900}
901
902static PyObject *
903dropwhile_next(dropwhileobject *lz)
904{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000905 PyObject *item, *good;
906 PyObject *it = lz->it;
907 long ok;
908 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000909
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000910 iternext = *Py_TYPE(it)->tp_iternext;
911 for (;;) {
912 item = iternext(it);
913 if (item == NULL)
914 return NULL;
915 if (lz->start == 1)
916 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000917
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000918 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
919 if (good == NULL) {
920 Py_DECREF(item);
921 return NULL;
922 }
923 ok = PyObject_IsTrue(good);
924 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +0200925 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000926 lz->start = 1;
927 return item;
928 }
929 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +0200930 if (ok < 0)
931 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000932 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000933}
934
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000935PyDoc_STRVAR(dropwhile_doc,
936"dropwhile(predicate, iterable) --> dropwhile object\n\
937\n\
938Drop items from the iterable while predicate(item) is true.\n\
939Afterwards, return every element until the iterable is exhausted.");
940
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000941static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000942 PyVarObject_HEAD_INIT(NULL, 0)
943 "itertools.dropwhile", /* tp_name */
944 sizeof(dropwhileobject), /* tp_basicsize */
945 0, /* tp_itemsize */
946 /* methods */
947 (destructor)dropwhile_dealloc, /* tp_dealloc */
948 0, /* tp_print */
949 0, /* tp_getattr */
950 0, /* tp_setattr */
951 0, /* tp_reserved */
952 0, /* tp_repr */
953 0, /* tp_as_number */
954 0, /* tp_as_sequence */
955 0, /* tp_as_mapping */
956 0, /* tp_hash */
957 0, /* tp_call */
958 0, /* tp_str */
959 PyObject_GenericGetAttr, /* tp_getattro */
960 0, /* tp_setattro */
961 0, /* tp_as_buffer */
962 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
963 Py_TPFLAGS_BASETYPE, /* tp_flags */
964 dropwhile_doc, /* tp_doc */
965 (traverseproc)dropwhile_traverse, /* tp_traverse */
966 0, /* tp_clear */
967 0, /* tp_richcompare */
968 0, /* tp_weaklistoffset */
969 PyObject_SelfIter, /* tp_iter */
970 (iternextfunc)dropwhile_next, /* tp_iternext */
971 0, /* tp_methods */
972 0, /* tp_members */
973 0, /* tp_getset */
974 0, /* tp_base */
975 0, /* tp_dict */
976 0, /* tp_descr_get */
977 0, /* tp_descr_set */
978 0, /* tp_dictoffset */
979 0, /* tp_init */
980 0, /* tp_alloc */
981 dropwhile_new, /* tp_new */
982 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000983};
984
985
986/* takewhile object **********************************************************/
987
988typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000989 PyObject_HEAD
990 PyObject *func;
991 PyObject *it;
992 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000993} takewhileobject;
994
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000995static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000996
997static PyObject *
998takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
999{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001000 PyObject *func, *seq;
1001 PyObject *it;
1002 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001003
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001004 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1005 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001006
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001007 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1008 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001009
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001010 /* Get iterator. */
1011 it = PyObject_GetIter(seq);
1012 if (it == NULL)
1013 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001014
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001015 /* create takewhileobject structure */
1016 lz = (takewhileobject *)type->tp_alloc(type, 0);
1017 if (lz == NULL) {
1018 Py_DECREF(it);
1019 return NULL;
1020 }
1021 Py_INCREF(func);
1022 lz->func = func;
1023 lz->it = it;
1024 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001025
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001026 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001027}
1028
1029static void
1030takewhile_dealloc(takewhileobject *lz)
1031{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001032 PyObject_GC_UnTrack(lz);
1033 Py_XDECREF(lz->func);
1034 Py_XDECREF(lz->it);
1035 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001036}
1037
1038static int
1039takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1040{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001041 Py_VISIT(lz->it);
1042 Py_VISIT(lz->func);
1043 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001044}
1045
1046static PyObject *
1047takewhile_next(takewhileobject *lz)
1048{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001049 PyObject *item, *good;
1050 PyObject *it = lz->it;
1051 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001052
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001053 if (lz->stop == 1)
1054 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001055
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001056 item = (*Py_TYPE(it)->tp_iternext)(it);
1057 if (item == NULL)
1058 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001059
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001060 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1061 if (good == NULL) {
1062 Py_DECREF(item);
1063 return NULL;
1064 }
1065 ok = PyObject_IsTrue(good);
1066 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001067 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001068 return item;
1069 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001070 if (ok == 0)
1071 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001072 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001073}
1074
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001075PyDoc_STRVAR(takewhile_doc,
1076"takewhile(predicate, iterable) --> takewhile object\n\
1077\n\
1078Return successive entries from an iterable as long as the \n\
1079predicate evaluates to true for each entry.");
1080
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001081static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001082 PyVarObject_HEAD_INIT(NULL, 0)
1083 "itertools.takewhile", /* tp_name */
1084 sizeof(takewhileobject), /* tp_basicsize */
1085 0, /* tp_itemsize */
1086 /* methods */
1087 (destructor)takewhile_dealloc, /* tp_dealloc */
1088 0, /* tp_print */
1089 0, /* tp_getattr */
1090 0, /* tp_setattr */
1091 0, /* tp_reserved */
1092 0, /* tp_repr */
1093 0, /* tp_as_number */
1094 0, /* tp_as_sequence */
1095 0, /* tp_as_mapping */
1096 0, /* tp_hash */
1097 0, /* tp_call */
1098 0, /* tp_str */
1099 PyObject_GenericGetAttr, /* tp_getattro */
1100 0, /* tp_setattro */
1101 0, /* tp_as_buffer */
1102 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1103 Py_TPFLAGS_BASETYPE, /* tp_flags */
1104 takewhile_doc, /* tp_doc */
1105 (traverseproc)takewhile_traverse, /* tp_traverse */
1106 0, /* tp_clear */
1107 0, /* tp_richcompare */
1108 0, /* tp_weaklistoffset */
1109 PyObject_SelfIter, /* tp_iter */
1110 (iternextfunc)takewhile_next, /* tp_iternext */
1111 0, /* tp_methods */
1112 0, /* tp_members */
1113 0, /* tp_getset */
1114 0, /* tp_base */
1115 0, /* tp_dict */
1116 0, /* tp_descr_get */
1117 0, /* tp_descr_set */
1118 0, /* tp_dictoffset */
1119 0, /* tp_init */
1120 0, /* tp_alloc */
1121 takewhile_new, /* tp_new */
1122 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001123};
1124
1125
1126/* islice object ************************************************************/
1127
1128typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001129 PyObject_HEAD
1130 PyObject *it;
1131 Py_ssize_t next;
1132 Py_ssize_t stop;
1133 Py_ssize_t step;
1134 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001135} isliceobject;
1136
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001137static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001138
1139static PyObject *
1140islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1141{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001142 PyObject *seq;
1143 Py_ssize_t start=0, stop=-1, step=1;
1144 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1145 Py_ssize_t numargs;
1146 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001147
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001148 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1149 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001150
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001151 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1152 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001153
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001154 numargs = PyTuple_Size(args);
1155 if (numargs == 2) {
1156 if (a1 != Py_None) {
1157 stop = PyLong_AsSsize_t(a1);
1158 if (stop == -1) {
1159 if (PyErr_Occurred())
1160 PyErr_Clear();
1161 PyErr_SetString(PyExc_ValueError,
1162 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1163 return NULL;
1164 }
1165 }
1166 } else {
1167 if (a1 != Py_None)
1168 start = PyLong_AsSsize_t(a1);
1169 if (start == -1 && PyErr_Occurred())
1170 PyErr_Clear();
1171 if (a2 != Py_None) {
1172 stop = PyLong_AsSsize_t(a2);
1173 if (stop == -1) {
1174 if (PyErr_Occurred())
1175 PyErr_Clear();
1176 PyErr_SetString(PyExc_ValueError,
1177 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1178 return NULL;
1179 }
1180 }
1181 }
1182 if (start<0 || stop<-1) {
1183 PyErr_SetString(PyExc_ValueError,
1184 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1185 return NULL;
1186 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001187
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001188 if (a3 != NULL) {
1189 if (a3 != Py_None)
1190 step = PyLong_AsSsize_t(a3);
1191 if (step == -1 && PyErr_Occurred())
1192 PyErr_Clear();
1193 }
1194 if (step<1) {
1195 PyErr_SetString(PyExc_ValueError,
1196 "Step for islice() must be a positive integer or None.");
1197 return NULL;
1198 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001199
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001200 /* Get iterator. */
1201 it = PyObject_GetIter(seq);
1202 if (it == NULL)
1203 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001205 /* create isliceobject structure */
1206 lz = (isliceobject *)type->tp_alloc(type, 0);
1207 if (lz == NULL) {
1208 Py_DECREF(it);
1209 return NULL;
1210 }
1211 lz->it = it;
1212 lz->next = start;
1213 lz->stop = stop;
1214 lz->step = step;
1215 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001216
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001217 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001218}
1219
1220static void
1221islice_dealloc(isliceobject *lz)
1222{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001223 PyObject_GC_UnTrack(lz);
1224 Py_XDECREF(lz->it);
1225 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001226}
1227
1228static int
1229islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1230{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001231 Py_VISIT(lz->it);
1232 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001233}
1234
1235static PyObject *
1236islice_next(isliceobject *lz)
1237{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001238 PyObject *item;
1239 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001240 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001241 Py_ssize_t oldnext;
1242 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001243
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001244 iternext = *Py_TYPE(it)->tp_iternext;
1245 while (lz->cnt < lz->next) {
1246 item = iternext(it);
1247 if (item == NULL)
1248 return NULL;
1249 Py_DECREF(item);
1250 lz->cnt++;
1251 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001252 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001253 return NULL;
1254 item = iternext(it);
1255 if (item == NULL)
1256 return NULL;
1257 lz->cnt++;
1258 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001259 /* The (size_t) cast below avoids the danger of undefined
1260 behaviour from signed integer overflow. */
1261 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001262 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1263 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001264 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001265}
1266
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001267PyDoc_STRVAR(islice_doc,
1268"islice(iterable, [start,] stop [, step]) --> islice object\n\
1269\n\
1270Return an iterator whose next() method returns selected values from an\n\
1271iterable. If start is specified, will skip all preceding elements;\n\
1272otherwise, start defaults to zero. Step defaults to one. If\n\
1273specified as another value, step determines how many values are \n\
1274skipped between successive calls. Works like a slice() on a list\n\
1275but returns an iterator.");
1276
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001277static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001278 PyVarObject_HEAD_INIT(NULL, 0)
1279 "itertools.islice", /* tp_name */
1280 sizeof(isliceobject), /* tp_basicsize */
1281 0, /* tp_itemsize */
1282 /* methods */
1283 (destructor)islice_dealloc, /* tp_dealloc */
1284 0, /* tp_print */
1285 0, /* tp_getattr */
1286 0, /* tp_setattr */
1287 0, /* tp_reserved */
1288 0, /* tp_repr */
1289 0, /* tp_as_number */
1290 0, /* tp_as_sequence */
1291 0, /* tp_as_mapping */
1292 0, /* tp_hash */
1293 0, /* tp_call */
1294 0, /* tp_str */
1295 PyObject_GenericGetAttr, /* tp_getattro */
1296 0, /* tp_setattro */
1297 0, /* tp_as_buffer */
1298 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1299 Py_TPFLAGS_BASETYPE, /* tp_flags */
1300 islice_doc, /* tp_doc */
1301 (traverseproc)islice_traverse, /* tp_traverse */
1302 0, /* tp_clear */
1303 0, /* tp_richcompare */
1304 0, /* tp_weaklistoffset */
1305 PyObject_SelfIter, /* tp_iter */
1306 (iternextfunc)islice_next, /* tp_iternext */
1307 0, /* tp_methods */
1308 0, /* tp_members */
1309 0, /* tp_getset */
1310 0, /* tp_base */
1311 0, /* tp_dict */
1312 0, /* tp_descr_get */
1313 0, /* tp_descr_set */
1314 0, /* tp_dictoffset */
1315 0, /* tp_init */
1316 0, /* tp_alloc */
1317 islice_new, /* tp_new */
1318 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001319};
1320
1321
1322/* starmap object ************************************************************/
1323
1324typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001325 PyObject_HEAD
1326 PyObject *func;
1327 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001328} starmapobject;
1329
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001330static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001331
1332static PyObject *
1333starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1334{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001335 PyObject *func, *seq;
1336 PyObject *it;
1337 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001338
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001339 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1340 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001341
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001342 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1343 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001344
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001345 /* Get iterator. */
1346 it = PyObject_GetIter(seq);
1347 if (it == NULL)
1348 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001349
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001350 /* create starmapobject structure */
1351 lz = (starmapobject *)type->tp_alloc(type, 0);
1352 if (lz == NULL) {
1353 Py_DECREF(it);
1354 return NULL;
1355 }
1356 Py_INCREF(func);
1357 lz->func = func;
1358 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001359
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001360 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001361}
1362
1363static void
1364starmap_dealloc(starmapobject *lz)
1365{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001366 PyObject_GC_UnTrack(lz);
1367 Py_XDECREF(lz->func);
1368 Py_XDECREF(lz->it);
1369 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001370}
1371
1372static int
1373starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1374{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001375 Py_VISIT(lz->it);
1376 Py_VISIT(lz->func);
1377 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001378}
1379
1380static PyObject *
1381starmap_next(starmapobject *lz)
1382{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001383 PyObject *args;
1384 PyObject *result;
1385 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001386
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001387 args = (*Py_TYPE(it)->tp_iternext)(it);
1388 if (args == NULL)
1389 return NULL;
1390 if (!PyTuple_CheckExact(args)) {
1391 PyObject *newargs = PySequence_Tuple(args);
1392 Py_DECREF(args);
1393 if (newargs == NULL)
1394 return NULL;
1395 args = newargs;
1396 }
1397 result = PyObject_Call(lz->func, args, NULL);
1398 Py_DECREF(args);
1399 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001400}
1401
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001402PyDoc_STRVAR(starmap_doc,
1403"starmap(function, sequence) --> starmap object\n\
1404\n\
1405Return an iterator whose values are returned from the function evaluated\n\
1406with a argument tuple taken from the given sequence.");
1407
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001408static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001409 PyVarObject_HEAD_INIT(NULL, 0)
1410 "itertools.starmap", /* tp_name */
1411 sizeof(starmapobject), /* tp_basicsize */
1412 0, /* tp_itemsize */
1413 /* methods */
1414 (destructor)starmap_dealloc, /* tp_dealloc */
1415 0, /* tp_print */
1416 0, /* tp_getattr */
1417 0, /* tp_setattr */
1418 0, /* tp_reserved */
1419 0, /* tp_repr */
1420 0, /* tp_as_number */
1421 0, /* tp_as_sequence */
1422 0, /* tp_as_mapping */
1423 0, /* tp_hash */
1424 0, /* tp_call */
1425 0, /* tp_str */
1426 PyObject_GenericGetAttr, /* tp_getattro */
1427 0, /* tp_setattro */
1428 0, /* tp_as_buffer */
1429 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1430 Py_TPFLAGS_BASETYPE, /* tp_flags */
1431 starmap_doc, /* tp_doc */
1432 (traverseproc)starmap_traverse, /* tp_traverse */
1433 0, /* tp_clear */
1434 0, /* tp_richcompare */
1435 0, /* tp_weaklistoffset */
1436 PyObject_SelfIter, /* tp_iter */
1437 (iternextfunc)starmap_next, /* tp_iternext */
1438 0, /* tp_methods */
1439 0, /* tp_members */
1440 0, /* tp_getset */
1441 0, /* tp_base */
1442 0, /* tp_dict */
1443 0, /* tp_descr_get */
1444 0, /* tp_descr_set */
1445 0, /* tp_dictoffset */
1446 0, /* tp_init */
1447 0, /* tp_alloc */
1448 starmap_new, /* tp_new */
1449 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001450};
1451
1452
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001453/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001454
1455typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001456 PyObject_HEAD
1457 PyObject *source; /* Iterator over input iterables */
1458 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001459} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001460
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001461static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001462
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001463static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001464chain_new_internal(PyTypeObject *type, PyObject *source)
1465{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001466 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001467
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001468 lz = (chainobject *)type->tp_alloc(type, 0);
1469 if (lz == NULL) {
1470 Py_DECREF(source);
1471 return NULL;
1472 }
1473
1474 lz->source = source;
1475 lz->active = NULL;
1476 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001477}
1478
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001479static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001480chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001481{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001482 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001483
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001484 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1485 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001486
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001487 source = PyObject_GetIter(args);
1488 if (source == NULL)
1489 return NULL;
1490
1491 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001492}
1493
1494static PyObject *
1495chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1496{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001497 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001498
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001499 source = PyObject_GetIter(arg);
1500 if (source == NULL)
1501 return NULL;
1502
1503 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001504}
1505
1506static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001507chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001508{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001509 PyObject_GC_UnTrack(lz);
1510 Py_XDECREF(lz->active);
1511 Py_XDECREF(lz->source);
1512 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001513}
1514
Raymond Hettinger2012f172003-02-07 05:32:58 +00001515static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001516chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001517{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001518 Py_VISIT(lz->source);
1519 Py_VISIT(lz->active);
1520 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001521}
1522
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001523static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001524chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001525{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001526 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001527
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001528 if (lz->source == NULL)
1529 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001530
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001531 if (lz->active == NULL) {
1532 PyObject *iterable = PyIter_Next(lz->source);
1533 if (iterable == NULL) {
1534 Py_CLEAR(lz->source);
1535 return NULL; /* no more input sources */
1536 }
1537 lz->active = PyObject_GetIter(iterable);
1538 Py_DECREF(iterable);
1539 if (lz->active == NULL) {
1540 Py_CLEAR(lz->source);
1541 return NULL; /* input not iterable */
1542 }
1543 }
1544 item = PyIter_Next(lz->active);
1545 if (item != NULL)
1546 return item;
1547 if (PyErr_Occurred()) {
1548 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1549 PyErr_Clear();
1550 else
1551 return NULL; /* input raised an exception */
1552 }
1553 Py_CLEAR(lz->active);
1554 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001555}
1556
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001557PyDoc_STRVAR(chain_doc,
1558"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001559\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001560Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001561first iterable until it is exhausted, then elements from the next\n\
1562iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001563
Christian Heimesf16baeb2008-02-29 14:57:44 +00001564PyDoc_STRVAR(chain_from_iterable_doc,
1565"chain.from_iterable(iterable) --> chain object\n\
1566\n\
1567Alternate chain() contructor taking a single iterable argument\n\
1568that evaluates lazily.");
1569
1570static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001571 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1572 chain_from_iterable_doc},
1573 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001574};
1575
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001576static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001577 PyVarObject_HEAD_INIT(NULL, 0)
1578 "itertools.chain", /* tp_name */
1579 sizeof(chainobject), /* tp_basicsize */
1580 0, /* tp_itemsize */
1581 /* methods */
1582 (destructor)chain_dealloc, /* tp_dealloc */
1583 0, /* tp_print */
1584 0, /* tp_getattr */
1585 0, /* tp_setattr */
1586 0, /* tp_reserved */
1587 0, /* tp_repr */
1588 0, /* tp_as_number */
1589 0, /* tp_as_sequence */
1590 0, /* tp_as_mapping */
1591 0, /* tp_hash */
1592 0, /* tp_call */
1593 0, /* tp_str */
1594 PyObject_GenericGetAttr, /* tp_getattro */
1595 0, /* tp_setattro */
1596 0, /* tp_as_buffer */
1597 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1598 Py_TPFLAGS_BASETYPE, /* tp_flags */
1599 chain_doc, /* tp_doc */
1600 (traverseproc)chain_traverse, /* tp_traverse */
1601 0, /* tp_clear */
1602 0, /* tp_richcompare */
1603 0, /* tp_weaklistoffset */
1604 PyObject_SelfIter, /* tp_iter */
1605 (iternextfunc)chain_next, /* tp_iternext */
1606 chain_methods, /* tp_methods */
1607 0, /* tp_members */
1608 0, /* tp_getset */
1609 0, /* tp_base */
1610 0, /* tp_dict */
1611 0, /* tp_descr_get */
1612 0, /* tp_descr_set */
1613 0, /* tp_dictoffset */
1614 0, /* tp_init */
1615 0, /* tp_alloc */
1616 chain_new, /* tp_new */
1617 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001618};
1619
1620
Christian Heimesc3f30c42008-02-22 16:37:40 +00001621/* product object ************************************************************/
1622
1623typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001624 PyObject_HEAD
1625 PyObject *pools; /* tuple of pool tuples */
1626 Py_ssize_t *indices; /* one index per pool */
1627 PyObject *result; /* most recently returned result tuple */
1628 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001629} productobject;
1630
1631static PyTypeObject product_type;
1632
1633static PyObject *
1634product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1635{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001636 productobject *lz;
1637 Py_ssize_t nargs, npools, repeat=1;
1638 PyObject *pools = NULL;
1639 Py_ssize_t *indices = NULL;
1640 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001641
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001642 if (kwds != NULL) {
1643 char *kwlist[] = {"repeat", 0};
1644 PyObject *tmpargs = PyTuple_New(0);
1645 if (tmpargs == NULL)
1646 return NULL;
1647 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1648 Py_DECREF(tmpargs);
1649 return NULL;
1650 }
1651 Py_DECREF(tmpargs);
1652 if (repeat < 0) {
1653 PyErr_SetString(PyExc_ValueError,
1654 "repeat argument cannot be negative");
1655 return NULL;
1656 }
1657 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001658
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001659 assert(PyTuple_Check(args));
1660 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1661 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001662
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001663 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1664 if (indices == NULL) {
1665 PyErr_NoMemory();
1666 goto error;
1667 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001668
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001669 pools = PyTuple_New(npools);
1670 if (pools == NULL)
1671 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001672
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001673 for (i=0; i < nargs ; ++i) {
1674 PyObject *item = PyTuple_GET_ITEM(args, i);
1675 PyObject *pool = PySequence_Tuple(item);
1676 if (pool == NULL)
1677 goto error;
1678 PyTuple_SET_ITEM(pools, i, pool);
1679 indices[i] = 0;
1680 }
1681 for ( ; i < npools; ++i) {
1682 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1683 Py_INCREF(pool);
1684 PyTuple_SET_ITEM(pools, i, pool);
1685 indices[i] = 0;
1686 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001687
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001688 /* create productobject structure */
1689 lz = (productobject *)type->tp_alloc(type, 0);
1690 if (lz == NULL)
1691 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001692
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001693 lz->pools = pools;
1694 lz->indices = indices;
1695 lz->result = NULL;
1696 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001697
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001698 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001699
1700error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001701 if (indices != NULL)
1702 PyMem_Free(indices);
1703 Py_XDECREF(pools);
1704 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001705}
1706
1707static void
1708product_dealloc(productobject *lz)
1709{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001710 PyObject_GC_UnTrack(lz);
1711 Py_XDECREF(lz->pools);
1712 Py_XDECREF(lz->result);
1713 if (lz->indices != NULL)
1714 PyMem_Free(lz->indices);
1715 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001716}
1717
1718static int
1719product_traverse(productobject *lz, visitproc visit, void *arg)
1720{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001721 Py_VISIT(lz->pools);
1722 Py_VISIT(lz->result);
1723 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001724}
1725
1726static PyObject *
1727product_next(productobject *lz)
1728{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001729 PyObject *pool;
1730 PyObject *elem;
1731 PyObject *oldelem;
1732 PyObject *pools = lz->pools;
1733 PyObject *result = lz->result;
1734 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1735 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001736
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001737 if (lz->stopped)
1738 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001739
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001740 if (result == NULL) {
1741 /* On the first pass, return an initial tuple filled with the
1742 first element from each pool. */
1743 result = PyTuple_New(npools);
1744 if (result == NULL)
1745 goto empty;
1746 lz->result = result;
1747 for (i=0; i < npools; i++) {
1748 pool = PyTuple_GET_ITEM(pools, i);
1749 if (PyTuple_GET_SIZE(pool) == 0)
1750 goto empty;
1751 elem = PyTuple_GET_ITEM(pool, 0);
1752 Py_INCREF(elem);
1753 PyTuple_SET_ITEM(result, i, elem);
1754 }
1755 } else {
1756 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001757
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001758 /* Copy the previous result tuple or re-use it if available */
1759 if (Py_REFCNT(result) > 1) {
1760 PyObject *old_result = result;
1761 result = PyTuple_New(npools);
1762 if (result == NULL)
1763 goto empty;
1764 lz->result = result;
1765 for (i=0; i < npools; i++) {
1766 elem = PyTuple_GET_ITEM(old_result, i);
1767 Py_INCREF(elem);
1768 PyTuple_SET_ITEM(result, i, elem);
1769 }
1770 Py_DECREF(old_result);
1771 }
1772 /* Now, we've got the only copy so we can update it in-place */
1773 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001774
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001775 /* Update the pool indices right-to-left. Only advance to the
1776 next pool when the previous one rolls-over */
1777 for (i=npools-1 ; i >= 0 ; i--) {
1778 pool = PyTuple_GET_ITEM(pools, i);
1779 indices[i]++;
1780 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1781 /* Roll-over and advance to next pool */
1782 indices[i] = 0;
1783 elem = PyTuple_GET_ITEM(pool, 0);
1784 Py_INCREF(elem);
1785 oldelem = PyTuple_GET_ITEM(result, i);
1786 PyTuple_SET_ITEM(result, i, elem);
1787 Py_DECREF(oldelem);
1788 } else {
1789 /* No rollover. Just increment and stop here. */
1790 elem = PyTuple_GET_ITEM(pool, indices[i]);
1791 Py_INCREF(elem);
1792 oldelem = PyTuple_GET_ITEM(result, i);
1793 PyTuple_SET_ITEM(result, i, elem);
1794 Py_DECREF(oldelem);
1795 break;
1796 }
1797 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001798
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001799 /* If i is negative, then the indices have all rolled-over
1800 and we're done. */
1801 if (i < 0)
1802 goto empty;
1803 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001804
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001805 Py_INCREF(result);
1806 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001807
1808empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001809 lz->stopped = 1;
1810 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001811}
1812
1813PyDoc_STRVAR(product_doc,
1814"product(*iterables) --> product object\n\
1815\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001816Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001817For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1818The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1819cycle in a manner similar to an odometer (with the rightmost element changing\n\
1820on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00001821To compute the product of an iterable with itself, specify the number\n\
1822of repetitions with the optional repeat keyword argument. For example,\n\
1823product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001824product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1825product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1826
1827static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001828 PyVarObject_HEAD_INIT(NULL, 0)
1829 "itertools.product", /* tp_name */
1830 sizeof(productobject), /* tp_basicsize */
1831 0, /* tp_itemsize */
1832 /* methods */
1833 (destructor)product_dealloc, /* tp_dealloc */
1834 0, /* tp_print */
1835 0, /* tp_getattr */
1836 0, /* tp_setattr */
1837 0, /* tp_reserved */
1838 0, /* tp_repr */
1839 0, /* tp_as_number */
1840 0, /* tp_as_sequence */
1841 0, /* tp_as_mapping */
1842 0, /* tp_hash */
1843 0, /* tp_call */
1844 0, /* tp_str */
1845 PyObject_GenericGetAttr, /* tp_getattro */
1846 0, /* tp_setattro */
1847 0, /* tp_as_buffer */
1848 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1849 Py_TPFLAGS_BASETYPE, /* tp_flags */
1850 product_doc, /* tp_doc */
1851 (traverseproc)product_traverse, /* tp_traverse */
1852 0, /* tp_clear */
1853 0, /* tp_richcompare */
1854 0, /* tp_weaklistoffset */
1855 PyObject_SelfIter, /* tp_iter */
1856 (iternextfunc)product_next, /* tp_iternext */
1857 0, /* tp_methods */
1858 0, /* tp_members */
1859 0, /* tp_getset */
1860 0, /* tp_base */
1861 0, /* tp_dict */
1862 0, /* tp_descr_get */
1863 0, /* tp_descr_set */
1864 0, /* tp_dictoffset */
1865 0, /* tp_init */
1866 0, /* tp_alloc */
1867 product_new, /* tp_new */
1868 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001869};
1870
1871
Christian Heimes380f7f22008-02-28 11:19:05 +00001872/* combinations object ************************************************************/
1873
1874typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001875 PyObject_HEAD
1876 PyObject *pool; /* input converted to a tuple */
1877 Py_ssize_t *indices; /* one index per result element */
1878 PyObject *result; /* most recently returned result tuple */
1879 Py_ssize_t r; /* size of result tuple */
1880 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00001881} combinationsobject;
1882
1883static PyTypeObject combinations_type;
1884
1885static PyObject *
1886combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1887{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001888 combinationsobject *co;
1889 Py_ssize_t n;
1890 Py_ssize_t r;
1891 PyObject *pool = NULL;
1892 PyObject *iterable = NULL;
1893 Py_ssize_t *indices = NULL;
1894 Py_ssize_t i;
1895 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00001896
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001897 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
1898 &iterable, &r))
1899 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001900
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001901 pool = PySequence_Tuple(iterable);
1902 if (pool == NULL)
1903 goto error;
1904 n = PyTuple_GET_SIZE(pool);
1905 if (r < 0) {
1906 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
1907 goto error;
1908 }
Christian Heimes380f7f22008-02-28 11:19:05 +00001909
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001910 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
1911 if (indices == NULL) {
1912 PyErr_NoMemory();
1913 goto error;
1914 }
Christian Heimes380f7f22008-02-28 11:19:05 +00001915
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001916 for (i=0 ; i<r ; i++)
1917 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00001918
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001919 /* create combinationsobject structure */
1920 co = (combinationsobject *)type->tp_alloc(type, 0);
1921 if (co == NULL)
1922 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00001923
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001924 co->pool = pool;
1925 co->indices = indices;
1926 co->result = NULL;
1927 co->r = r;
1928 co->stopped = r > n ? 1 : 0;
1929
1930 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00001931
1932error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001933 if (indices != NULL)
1934 PyMem_Free(indices);
1935 Py_XDECREF(pool);
1936 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001937}
1938
1939static void
1940combinations_dealloc(combinationsobject *co)
1941{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001942 PyObject_GC_UnTrack(co);
1943 Py_XDECREF(co->pool);
1944 Py_XDECREF(co->result);
1945 if (co->indices != NULL)
1946 PyMem_Free(co->indices);
1947 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00001948}
1949
1950static int
1951combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
1952{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001953 Py_VISIT(co->pool);
1954 Py_VISIT(co->result);
1955 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00001956}
1957
1958static PyObject *
1959combinations_next(combinationsobject *co)
1960{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001961 PyObject *elem;
1962 PyObject *oldelem;
1963 PyObject *pool = co->pool;
1964 Py_ssize_t *indices = co->indices;
1965 PyObject *result = co->result;
1966 Py_ssize_t n = PyTuple_GET_SIZE(pool);
1967 Py_ssize_t r = co->r;
1968 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00001969
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001970 if (co->stopped)
1971 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00001972
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001973 if (result == NULL) {
1974 /* On the first pass, initialize result tuple using the indices */
1975 result = PyTuple_New(r);
1976 if (result == NULL)
1977 goto empty;
1978 co->result = result;
1979 for (i=0; i<r ; i++) {
1980 index = indices[i];
1981 elem = PyTuple_GET_ITEM(pool, index);
1982 Py_INCREF(elem);
1983 PyTuple_SET_ITEM(result, i, elem);
1984 }
1985 } else {
1986 /* Copy the previous result tuple or re-use it if available */
1987 if (Py_REFCNT(result) > 1) {
1988 PyObject *old_result = result;
1989 result = PyTuple_New(r);
1990 if (result == NULL)
1991 goto empty;
1992 co->result = result;
1993 for (i=0; i<r ; i++) {
1994 elem = PyTuple_GET_ITEM(old_result, i);
1995 Py_INCREF(elem);
1996 PyTuple_SET_ITEM(result, i, elem);
1997 }
1998 Py_DECREF(old_result);
1999 }
2000 /* Now, we've got the only copy so we can update it in-place
2001 * CPython's empty tuple is a singleton and cached in
2002 * PyTuple's freelist.
2003 */
2004 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002005
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002006 /* Scan indices right-to-left until finding one that is not
2007 at its maximum (i + n - r). */
2008 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2009 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002010
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002011 /* If i is negative, then the indices are all at
2012 their maximum value and we're done. */
2013 if (i < 0)
2014 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002015
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002016 /* Increment the current index which we know is not at its
2017 maximum. Then move back to the right setting each index
2018 to its lowest possible value (one higher than the index
2019 to its left -- this maintains the sort order invariant). */
2020 indices[i]++;
2021 for (j=i+1 ; j<r ; j++)
2022 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002023
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002024 /* Update the result tuple for the new indices
2025 starting with i, the leftmost index that changed */
2026 for ( ; i<r ; i++) {
2027 index = indices[i];
2028 elem = PyTuple_GET_ITEM(pool, index);
2029 Py_INCREF(elem);
2030 oldelem = PyTuple_GET_ITEM(result, i);
2031 PyTuple_SET_ITEM(result, i, elem);
2032 Py_DECREF(oldelem);
2033 }
2034 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002035
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002036 Py_INCREF(result);
2037 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002038
2039empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002040 co->stopped = 1;
2041 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002042}
2043
2044PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002045"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002046\n\
2047Return successive r-length combinations of elements in the iterable.\n\n\
2048combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2049
2050static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002051 PyVarObject_HEAD_INIT(NULL, 0)
2052 "itertools.combinations", /* tp_name */
2053 sizeof(combinationsobject), /* tp_basicsize */
2054 0, /* tp_itemsize */
2055 /* methods */
2056 (destructor)combinations_dealloc, /* tp_dealloc */
2057 0, /* tp_print */
2058 0, /* tp_getattr */
2059 0, /* tp_setattr */
2060 0, /* tp_reserved */
2061 0, /* tp_repr */
2062 0, /* tp_as_number */
2063 0, /* tp_as_sequence */
2064 0, /* tp_as_mapping */
2065 0, /* tp_hash */
2066 0, /* tp_call */
2067 0, /* tp_str */
2068 PyObject_GenericGetAttr, /* tp_getattro */
2069 0, /* tp_setattro */
2070 0, /* tp_as_buffer */
2071 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2072 Py_TPFLAGS_BASETYPE, /* tp_flags */
2073 combinations_doc, /* tp_doc */
2074 (traverseproc)combinations_traverse, /* tp_traverse */
2075 0, /* tp_clear */
2076 0, /* tp_richcompare */
2077 0, /* tp_weaklistoffset */
2078 PyObject_SelfIter, /* tp_iter */
2079 (iternextfunc)combinations_next, /* tp_iternext */
2080 0, /* tp_methods */
2081 0, /* tp_members */
2082 0, /* tp_getset */
2083 0, /* tp_base */
2084 0, /* tp_dict */
2085 0, /* tp_descr_get */
2086 0, /* tp_descr_set */
2087 0, /* tp_dictoffset */
2088 0, /* tp_init */
2089 0, /* tp_alloc */
2090 combinations_new, /* tp_new */
2091 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002092};
2093
2094
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002095/* combinations with replacement object *******************************************/
2096
2097/* Equivalent to:
2098
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002099 def combinations_with_replacement(iterable, r):
2100 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2101 # number items returned: (n+r-1)! / r! / (n-1)!
2102 pool = tuple(iterable)
2103 n = len(pool)
2104 indices = [0] * r
2105 yield tuple(pool[i] for i in indices)
2106 while 1:
2107 for i in reversed(range(r)):
2108 if indices[i] != n - 1:
2109 break
2110 else:
2111 return
2112 indices[i:] = [indices[i] + 1] * (r - i)
2113 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002114
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002115 def combinations_with_replacement2(iterable, r):
2116 'Alternate version that filters from product()'
2117 pool = tuple(iterable)
2118 n = len(pool)
2119 for indices in product(range(n), repeat=r):
2120 if sorted(indices) == list(indices):
2121 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002122*/
2123typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002124 PyObject_HEAD
2125 PyObject *pool; /* input converted to a tuple */
2126 Py_ssize_t *indices; /* one index per result element */
2127 PyObject *result; /* most recently returned result tuple */
2128 Py_ssize_t r; /* size of result tuple */
2129 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002130} cwrobject;
2131
2132static PyTypeObject cwr_type;
2133
2134static PyObject *
2135cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2136{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002137 cwrobject *co;
2138 Py_ssize_t n;
2139 Py_ssize_t r;
2140 PyObject *pool = NULL;
2141 PyObject *iterable = NULL;
2142 Py_ssize_t *indices = NULL;
2143 Py_ssize_t i;
2144 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002145
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002146 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2147 &iterable, &r))
2148 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002149
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002150 pool = PySequence_Tuple(iterable);
2151 if (pool == NULL)
2152 goto error;
2153 n = PyTuple_GET_SIZE(pool);
2154 if (r < 0) {
2155 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2156 goto error;
2157 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002158
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002159 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2160 if (indices == NULL) {
2161 PyErr_NoMemory();
2162 goto error;
2163 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002164
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002165 for (i=0 ; i<r ; i++)
2166 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002167
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002168 /* create cwrobject structure */
2169 co = (cwrobject *)type->tp_alloc(type, 0);
2170 if (co == NULL)
2171 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002172
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002173 co->pool = pool;
2174 co->indices = indices;
2175 co->result = NULL;
2176 co->r = r;
2177 co->stopped = !n && r;
2178
2179 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002180
2181error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002182 if (indices != NULL)
2183 PyMem_Free(indices);
2184 Py_XDECREF(pool);
2185 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002186}
2187
2188static void
2189cwr_dealloc(cwrobject *co)
2190{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002191 PyObject_GC_UnTrack(co);
2192 Py_XDECREF(co->pool);
2193 Py_XDECREF(co->result);
2194 if (co->indices != NULL)
2195 PyMem_Free(co->indices);
2196 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002197}
2198
2199static int
2200cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2201{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002202 Py_VISIT(co->pool);
2203 Py_VISIT(co->result);
2204 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002205}
2206
2207static PyObject *
2208cwr_next(cwrobject *co)
2209{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002210 PyObject *elem;
2211 PyObject *oldelem;
2212 PyObject *pool = co->pool;
2213 Py_ssize_t *indices = co->indices;
2214 PyObject *result = co->result;
2215 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2216 Py_ssize_t r = co->r;
2217 Py_ssize_t i, j, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002218
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002219 if (co->stopped)
2220 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002221
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002222 if (result == NULL) {
2223 /* On the first pass, initialize result tuple using the indices */
2224 result = PyTuple_New(r);
2225 if (result == NULL)
2226 goto empty;
2227 co->result = result;
2228 for (i=0; i<r ; i++) {
2229 index = indices[i];
2230 elem = PyTuple_GET_ITEM(pool, index);
2231 Py_INCREF(elem);
2232 PyTuple_SET_ITEM(result, i, elem);
2233 }
2234 } else {
2235 /* Copy the previous result tuple or re-use it if available */
2236 if (Py_REFCNT(result) > 1) {
2237 PyObject *old_result = result;
2238 result = PyTuple_New(r);
2239 if (result == NULL)
2240 goto empty;
2241 co->result = result;
2242 for (i=0; i<r ; i++) {
2243 elem = PyTuple_GET_ITEM(old_result, i);
2244 Py_INCREF(elem);
2245 PyTuple_SET_ITEM(result, i, elem);
2246 }
2247 Py_DECREF(old_result);
2248 }
2249 /* Now, we've got the only copy so we can update it in-place CPython's
2250 empty tuple is a singleton and cached in PyTuple's freelist. */
2251 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002252
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002253 /* Scan indices right-to-left until finding one that is not
2254 * at its maximum (n-1). */
2255 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2256 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002257
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002258 /* If i is negative, then the indices are all at
2259 their maximum value and we're done. */
2260 if (i < 0)
2261 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002262
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002263 /* Increment the current index which we know is not at its
2264 maximum. Then set all to the right to the same value. */
2265 indices[i]++;
2266 for (j=i+1 ; j<r ; j++)
2267 indices[j] = indices[j-1];
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002268
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002269 /* Update the result tuple for the new indices
2270 starting with i, the leftmost index that changed */
2271 for ( ; i<r ; i++) {
2272 index = indices[i];
2273 elem = PyTuple_GET_ITEM(pool, index);
2274 Py_INCREF(elem);
2275 oldelem = PyTuple_GET_ITEM(result, i);
2276 PyTuple_SET_ITEM(result, i, elem);
2277 Py_DECREF(oldelem);
2278 }
2279 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002280
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002281 Py_INCREF(result);
2282 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002283
2284empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002285 co->stopped = 1;
2286 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002287}
2288
2289PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002290"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002291\n\
2292Return successive r-length combinations of elements in the iterable\n\
2293allowing individual elements to have successive repeats.\n\
2294combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2295
2296static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002297 PyVarObject_HEAD_INIT(NULL, 0)
2298 "itertools.combinations_with_replacement", /* tp_name */
2299 sizeof(cwrobject), /* tp_basicsize */
2300 0, /* tp_itemsize */
2301 /* methods */
2302 (destructor)cwr_dealloc, /* tp_dealloc */
2303 0, /* tp_print */
2304 0, /* tp_getattr */
2305 0, /* tp_setattr */
2306 0, /* tp_reserved */
2307 0, /* tp_repr */
2308 0, /* tp_as_number */
2309 0, /* tp_as_sequence */
2310 0, /* tp_as_mapping */
2311 0, /* tp_hash */
2312 0, /* tp_call */
2313 0, /* tp_str */
2314 PyObject_GenericGetAttr, /* tp_getattro */
2315 0, /* tp_setattro */
2316 0, /* tp_as_buffer */
2317 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2318 Py_TPFLAGS_BASETYPE, /* tp_flags */
2319 cwr_doc, /* tp_doc */
2320 (traverseproc)cwr_traverse, /* tp_traverse */
2321 0, /* tp_clear */
2322 0, /* tp_richcompare */
2323 0, /* tp_weaklistoffset */
2324 PyObject_SelfIter, /* tp_iter */
2325 (iternextfunc)cwr_next, /* tp_iternext */
2326 0, /* tp_methods */
2327 0, /* tp_members */
2328 0, /* tp_getset */
2329 0, /* tp_base */
2330 0, /* tp_dict */
2331 0, /* tp_descr_get */
2332 0, /* tp_descr_set */
2333 0, /* tp_dictoffset */
2334 0, /* tp_init */
2335 0, /* tp_alloc */
2336 cwr_new, /* tp_new */
2337 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002338};
2339
2340
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002341/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002342
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002343def permutations(iterable, r=None):
2344 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2345 pool = tuple(iterable)
2346 n = len(pool)
2347 r = n if r is None else r
2348 indices = range(n)
2349 cycles = range(n-r+1, n+1)[::-1]
2350 yield tuple(pool[i] for i in indices[:r])
2351 while n:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002352 for i in reversed(range(r)):
2353 cycles[i] -= 1
2354 if cycles[i] == 0:
2355 indices[i:] = indices[i+1:] + indices[i:i+1]
2356 cycles[i] = n - i
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002357 else:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002358 j = cycles[i]
2359 indices[i], indices[-j] = indices[-j], indices[i]
2360 yield tuple(pool[i] for i in indices[:r])
2361 break
2362 else:
2363 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002364*/
2365
2366typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002367 PyObject_HEAD
2368 PyObject *pool; /* input converted to a tuple */
2369 Py_ssize_t *indices; /* one index per element in the pool */
2370 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2371 PyObject *result; /* most recently returned result tuple */
2372 Py_ssize_t r; /* size of result tuple */
2373 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002374} permutationsobject;
2375
2376static PyTypeObject permutations_type;
2377
2378static PyObject *
2379permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2380{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002381 permutationsobject *po;
2382 Py_ssize_t n;
2383 Py_ssize_t r;
2384 PyObject *robj = Py_None;
2385 PyObject *pool = NULL;
2386 PyObject *iterable = NULL;
2387 Py_ssize_t *indices = NULL;
2388 Py_ssize_t *cycles = NULL;
2389 Py_ssize_t i;
2390 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002391
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002392 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2393 &iterable, &robj))
2394 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002395
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002396 pool = PySequence_Tuple(iterable);
2397 if (pool == NULL)
2398 goto error;
2399 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002400
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002401 r = n;
2402 if (robj != Py_None) {
2403 if (!PyLong_Check(robj)) {
2404 PyErr_SetString(PyExc_TypeError, "Expected int as r");
2405 goto error;
2406 }
2407 r = PyLong_AsSsize_t(robj);
2408 if (r == -1 && PyErr_Occurred())
2409 goto error;
2410 }
2411 if (r < 0) {
2412 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2413 goto error;
2414 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002415
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002416 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2417 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2418 if (indices == NULL || cycles == NULL) {
2419 PyErr_NoMemory();
2420 goto error;
2421 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002422
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002423 for (i=0 ; i<n ; i++)
2424 indices[i] = i;
2425 for (i=0 ; i<r ; i++)
2426 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002427
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002428 /* create permutationsobject structure */
2429 po = (permutationsobject *)type->tp_alloc(type, 0);
2430 if (po == NULL)
2431 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002432
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002433 po->pool = pool;
2434 po->indices = indices;
2435 po->cycles = cycles;
2436 po->result = NULL;
2437 po->r = r;
2438 po->stopped = r > n ? 1 : 0;
2439
2440 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002441
2442error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002443 if (indices != NULL)
2444 PyMem_Free(indices);
2445 if (cycles != NULL)
2446 PyMem_Free(cycles);
2447 Py_XDECREF(pool);
2448 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002449}
2450
2451static void
2452permutations_dealloc(permutationsobject *po)
2453{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002454 PyObject_GC_UnTrack(po);
2455 Py_XDECREF(po->pool);
2456 Py_XDECREF(po->result);
2457 PyMem_Free(po->indices);
2458 PyMem_Free(po->cycles);
2459 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002460}
2461
2462static int
2463permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2464{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002465 Py_VISIT(po->pool);
2466 Py_VISIT(po->result);
2467 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002468}
2469
2470static PyObject *
2471permutations_next(permutationsobject *po)
2472{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002473 PyObject *elem;
2474 PyObject *oldelem;
2475 PyObject *pool = po->pool;
2476 Py_ssize_t *indices = po->indices;
2477 Py_ssize_t *cycles = po->cycles;
2478 PyObject *result = po->result;
2479 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2480 Py_ssize_t r = po->r;
2481 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002482
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002483 if (po->stopped)
2484 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002485
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002486 if (result == NULL) {
2487 /* On the first pass, initialize result tuple using the indices */
2488 result = PyTuple_New(r);
2489 if (result == NULL)
2490 goto empty;
2491 po->result = result;
2492 for (i=0; i<r ; i++) {
2493 index = indices[i];
2494 elem = PyTuple_GET_ITEM(pool, index);
2495 Py_INCREF(elem);
2496 PyTuple_SET_ITEM(result, i, elem);
2497 }
2498 } else {
2499 if (n == 0)
2500 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002501
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002502 /* Copy the previous result tuple or re-use it if available */
2503 if (Py_REFCNT(result) > 1) {
2504 PyObject *old_result = result;
2505 result = PyTuple_New(r);
2506 if (result == NULL)
2507 goto empty;
2508 po->result = result;
2509 for (i=0; i<r ; i++) {
2510 elem = PyTuple_GET_ITEM(old_result, i);
2511 Py_INCREF(elem);
2512 PyTuple_SET_ITEM(result, i, elem);
2513 }
2514 Py_DECREF(old_result);
2515 }
2516 /* Now, we've got the only copy so we can update it in-place */
2517 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002518
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002519 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2520 for (i=r-1 ; i>=0 ; i--) {
2521 cycles[i] -= 1;
2522 if (cycles[i] == 0) {
2523 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2524 index = indices[i];
2525 for (j=i ; j<n-1 ; j++)
2526 indices[j] = indices[j+1];
2527 indices[n-1] = index;
2528 cycles[i] = n - i;
2529 } else {
2530 j = cycles[i];
2531 index = indices[i];
2532 indices[i] = indices[n-j];
2533 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002534
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002535 for (k=i; k<r ; k++) {
2536 /* start with i, the leftmost element that changed */
2537 /* yield tuple(pool[k] for k in indices[:r]) */
2538 index = indices[k];
2539 elem = PyTuple_GET_ITEM(pool, index);
2540 Py_INCREF(elem);
2541 oldelem = PyTuple_GET_ITEM(result, k);
2542 PyTuple_SET_ITEM(result, k, elem);
2543 Py_DECREF(oldelem);
2544 }
2545 break;
2546 }
2547 }
2548 /* If i is negative, then the cycles have all
2549 rolled-over and we're done. */
2550 if (i < 0)
2551 goto empty;
2552 }
2553 Py_INCREF(result);
2554 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002555
2556empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002557 po->stopped = 1;
2558 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002559}
2560
2561PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00002562"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002563\n\
2564Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00002565permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002566
2567static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002568 PyVarObject_HEAD_INIT(NULL, 0)
2569 "itertools.permutations", /* tp_name */
2570 sizeof(permutationsobject), /* tp_basicsize */
2571 0, /* tp_itemsize */
2572 /* methods */
2573 (destructor)permutations_dealloc, /* tp_dealloc */
2574 0, /* tp_print */
2575 0, /* tp_getattr */
2576 0, /* tp_setattr */
2577 0, /* tp_reserved */
2578 0, /* tp_repr */
2579 0, /* tp_as_number */
2580 0, /* tp_as_sequence */
2581 0, /* tp_as_mapping */
2582 0, /* tp_hash */
2583 0, /* tp_call */
2584 0, /* tp_str */
2585 PyObject_GenericGetAttr, /* tp_getattro */
2586 0, /* tp_setattro */
2587 0, /* tp_as_buffer */
2588 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2589 Py_TPFLAGS_BASETYPE, /* tp_flags */
2590 permutations_doc, /* tp_doc */
2591 (traverseproc)permutations_traverse, /* tp_traverse */
2592 0, /* tp_clear */
2593 0, /* tp_richcompare */
2594 0, /* tp_weaklistoffset */
2595 PyObject_SelfIter, /* tp_iter */
2596 (iternextfunc)permutations_next, /* tp_iternext */
2597 0, /* tp_methods */
2598 0, /* tp_members */
2599 0, /* tp_getset */
2600 0, /* tp_base */
2601 0, /* tp_dict */
2602 0, /* tp_descr_get */
2603 0, /* tp_descr_set */
2604 0, /* tp_dictoffset */
2605 0, /* tp_init */
2606 0, /* tp_alloc */
2607 permutations_new, /* tp_new */
2608 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002609};
2610
Raymond Hettinger482ba772010-12-01 22:48:00 +00002611/* accumulate object ************************************************************/
2612
2613typedef struct {
2614 PyObject_HEAD
2615 PyObject *total;
2616 PyObject *it;
2617} accumulateobject;
2618
2619static PyTypeObject accumulate_type;
2620
2621static PyObject *
2622accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2623{
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002624 static char *kwargs[] = {"iterable", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00002625 PyObject *iterable;
2626 PyObject *it;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002627 accumulateobject *lz;
2628
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002629 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O:accumulate", kwargs, &iterable))
2630 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002631
2632 /* Get iterator. */
2633 it = PyObject_GetIter(iterable);
2634 if (it == NULL)
2635 return NULL;
2636
Raymond Hettinger482ba772010-12-01 22:48:00 +00002637 /* create accumulateobject structure */
2638 lz = (accumulateobject *)type->tp_alloc(type, 0);
2639 if (lz == NULL) {
2640 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002641 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002642 }
2643
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002644 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002645 lz->it = it;
2646 return (PyObject *)lz;
2647}
2648
2649static void
2650accumulate_dealloc(accumulateobject *lz)
2651{
2652 PyObject_GC_UnTrack(lz);
2653 Py_XDECREF(lz->total);
2654 Py_XDECREF(lz->it);
2655 Py_TYPE(lz)->tp_free(lz);
2656}
2657
2658static int
2659accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
2660{
2661 Py_VISIT(lz->it);
2662 Py_VISIT(lz->total);
2663 return 0;
2664}
2665
2666static PyObject *
2667accumulate_next(accumulateobject *lz)
2668{
2669 PyObject *val, *oldtotal, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02002670
Raymond Hettinger482ba772010-12-01 22:48:00 +00002671 val = PyIter_Next(lz->it);
2672 if (val == NULL)
2673 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02002674
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002675 if (lz->total == NULL) {
2676 Py_INCREF(val);
2677 lz->total = val;
2678 return lz->total;
2679 }
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02002680
Raymond Hettinger482ba772010-12-01 22:48:00 +00002681 newtotal = PyNumber_Add(lz->total, val);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002682 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00002683 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002684 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00002685
2686 oldtotal = lz->total;
2687 lz->total = newtotal;
2688 Py_DECREF(oldtotal);
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02002689
Raymond Hettinger482ba772010-12-01 22:48:00 +00002690 Py_INCREF(newtotal);
2691 return newtotal;
2692}
2693
2694PyDoc_STRVAR(accumulate_doc,
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00002695"accumulate(iterable) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00002696\n\
2697Return series of accumulated sums.");
2698
2699static PyTypeObject accumulate_type = {
2700 PyVarObject_HEAD_INIT(NULL, 0)
2701 "itertools.accumulate", /* tp_name */
2702 sizeof(accumulateobject), /* tp_basicsize */
2703 0, /* tp_itemsize */
2704 /* methods */
2705 (destructor)accumulate_dealloc, /* tp_dealloc */
2706 0, /* tp_print */
2707 0, /* tp_getattr */
2708 0, /* tp_setattr */
2709 0, /* tp_reserved */
2710 0, /* tp_repr */
2711 0, /* tp_as_number */
2712 0, /* tp_as_sequence */
2713 0, /* tp_as_mapping */
2714 0, /* tp_hash */
2715 0, /* tp_call */
2716 0, /* tp_str */
2717 PyObject_GenericGetAttr, /* tp_getattro */
2718 0, /* tp_setattro */
2719 0, /* tp_as_buffer */
2720 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2721 Py_TPFLAGS_BASETYPE, /* tp_flags */
2722 accumulate_doc, /* tp_doc */
2723 (traverseproc)accumulate_traverse, /* tp_traverse */
2724 0, /* tp_clear */
2725 0, /* tp_richcompare */
2726 0, /* tp_weaklistoffset */
2727 PyObject_SelfIter, /* tp_iter */
2728 (iternextfunc)accumulate_next, /* tp_iternext */
2729 0, /* tp_methods */
2730 0, /* tp_members */
2731 0, /* tp_getset */
2732 0, /* tp_base */
2733 0, /* tp_dict */
2734 0, /* tp_descr_get */
2735 0, /* tp_descr_set */
2736 0, /* tp_dictoffset */
2737 0, /* tp_init */
2738 0, /* tp_alloc */
2739 accumulate_new, /* tp_new */
2740 PyObject_GC_Del, /* tp_free */
2741};
2742
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002743
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002744/* compress object ************************************************************/
2745
2746/* Equivalent to:
2747
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002748 def compress(data, selectors):
2749 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2750 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002751*/
2752
2753typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002754 PyObject_HEAD
2755 PyObject *data;
2756 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002757} compressobject;
2758
2759static PyTypeObject compress_type;
2760
2761static PyObject *
2762compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2763{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002764 PyObject *seq1, *seq2;
2765 PyObject *data=NULL, *selectors=NULL;
2766 compressobject *lz;
2767 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002768
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002769 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2770 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002771
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002772 data = PyObject_GetIter(seq1);
2773 if (data == NULL)
2774 goto fail;
2775 selectors = PyObject_GetIter(seq2);
2776 if (selectors == NULL)
2777 goto fail;
2778
2779 /* create compressobject structure */
2780 lz = (compressobject *)type->tp_alloc(type, 0);
2781 if (lz == NULL)
2782 goto fail;
2783 lz->data = data;
2784 lz->selectors = selectors;
2785 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002786
2787fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002788 Py_XDECREF(data);
2789 Py_XDECREF(selectors);
2790 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002791}
2792
2793static void
2794compress_dealloc(compressobject *lz)
2795{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002796 PyObject_GC_UnTrack(lz);
2797 Py_XDECREF(lz->data);
2798 Py_XDECREF(lz->selectors);
2799 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002800}
2801
2802static int
2803compress_traverse(compressobject *lz, visitproc visit, void *arg)
2804{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002805 Py_VISIT(lz->data);
2806 Py_VISIT(lz->selectors);
2807 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002808}
2809
2810static PyObject *
2811compress_next(compressobject *lz)
2812{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002813 PyObject *data = lz->data, *selectors = lz->selectors;
2814 PyObject *datum, *selector;
2815 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2816 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2817 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002818
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002819 while (1) {
2820 /* Steps: get datum, get selector, evaluate selector.
2821 Order is important (to match the pure python version
2822 in terms of which input gets a chance to raise an
2823 exception first).
2824 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002825
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002826 datum = datanext(data);
2827 if (datum == NULL)
2828 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002829
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002830 selector = selectornext(selectors);
2831 if (selector == NULL) {
2832 Py_DECREF(datum);
2833 return NULL;
2834 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002835
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002836 ok = PyObject_IsTrue(selector);
2837 Py_DECREF(selector);
2838 if (ok == 1)
2839 return datum;
2840 Py_DECREF(datum);
2841 if (ok == -1)
2842 return NULL;
2843 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002844}
2845
2846PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00002847"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002848\n\
2849Return data elements corresponding to true selector elements.\n\
2850Forms a shorter iterator from selected data elements using the\n\
2851selectors to choose the data elements.");
2852
2853static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002854 PyVarObject_HEAD_INIT(NULL, 0)
2855 "itertools.compress", /* tp_name */
2856 sizeof(compressobject), /* tp_basicsize */
2857 0, /* tp_itemsize */
2858 /* methods */
2859 (destructor)compress_dealloc, /* tp_dealloc */
2860 0, /* tp_print */
2861 0, /* tp_getattr */
2862 0, /* tp_setattr */
2863 0, /* tp_reserved */
2864 0, /* tp_repr */
2865 0, /* tp_as_number */
2866 0, /* tp_as_sequence */
2867 0, /* tp_as_mapping */
2868 0, /* tp_hash */
2869 0, /* tp_call */
2870 0, /* tp_str */
2871 PyObject_GenericGetAttr, /* tp_getattro */
2872 0, /* tp_setattro */
2873 0, /* tp_as_buffer */
2874 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2875 Py_TPFLAGS_BASETYPE, /* tp_flags */
2876 compress_doc, /* tp_doc */
2877 (traverseproc)compress_traverse, /* tp_traverse */
2878 0, /* tp_clear */
2879 0, /* tp_richcompare */
2880 0, /* tp_weaklistoffset */
2881 PyObject_SelfIter, /* tp_iter */
2882 (iternextfunc)compress_next, /* tp_iternext */
2883 0, /* tp_methods */
2884 0, /* tp_members */
2885 0, /* tp_getset */
2886 0, /* tp_base */
2887 0, /* tp_dict */
2888 0, /* tp_descr_get */
2889 0, /* tp_descr_set */
2890 0, /* tp_dictoffset */
2891 0, /* tp_init */
2892 0, /* tp_alloc */
2893 compress_new, /* tp_new */
2894 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002895};
2896
2897
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002898/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00002899
2900typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002901 PyObject_HEAD
2902 PyObject *func;
2903 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002904} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002905
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002906static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002907
2908static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002909filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002910{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002911 PyObject *func, *seq;
2912 PyObject *it;
2913 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002914
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002915 if (type == &filterfalse_type &&
2916 !_PyArg_NoKeywords("filterfalse()", kwds))
2917 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002918
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002919 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
2920 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002921
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002922 /* Get iterator. */
2923 it = PyObject_GetIter(seq);
2924 if (it == NULL)
2925 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002926
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002927 /* create filterfalseobject structure */
2928 lz = (filterfalseobject *)type->tp_alloc(type, 0);
2929 if (lz == NULL) {
2930 Py_DECREF(it);
2931 return NULL;
2932 }
2933 Py_INCREF(func);
2934 lz->func = func;
2935 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002936
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002937 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002938}
2939
2940static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002941filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002942{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002943 PyObject_GC_UnTrack(lz);
2944 Py_XDECREF(lz->func);
2945 Py_XDECREF(lz->it);
2946 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002947}
2948
2949static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002950filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002951{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002952 Py_VISIT(lz->it);
2953 Py_VISIT(lz->func);
2954 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002955}
2956
2957static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002958filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002959{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002960 PyObject *item;
2961 PyObject *it = lz->it;
2962 long ok;
2963 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002964
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002965 iternext = *Py_TYPE(it)->tp_iternext;
2966 for (;;) {
2967 item = iternext(it);
2968 if (item == NULL)
2969 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002970
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002971 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2972 ok = PyObject_IsTrue(item);
2973 } else {
2974 PyObject *good;
2975 good = PyObject_CallFunctionObjArgs(lz->func,
2976 item, NULL);
2977 if (good == NULL) {
2978 Py_DECREF(item);
2979 return NULL;
2980 }
2981 ok = PyObject_IsTrue(good);
2982 Py_DECREF(good);
2983 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02002984 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002985 return item;
2986 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02002987 if (ok < 0)
2988 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002989 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00002990}
2991
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002992PyDoc_STRVAR(filterfalse_doc,
2993"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002994\n\
2995Return those items of sequence for which function(item) is false.\n\
2996If function is None, return the items that are false.");
2997
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002998static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002999 PyVarObject_HEAD_INIT(NULL, 0)
3000 "itertools.filterfalse", /* tp_name */
3001 sizeof(filterfalseobject), /* tp_basicsize */
3002 0, /* tp_itemsize */
3003 /* methods */
3004 (destructor)filterfalse_dealloc, /* tp_dealloc */
3005 0, /* tp_print */
3006 0, /* tp_getattr */
3007 0, /* tp_setattr */
3008 0, /* tp_reserved */
3009 0, /* tp_repr */
3010 0, /* tp_as_number */
3011 0, /* tp_as_sequence */
3012 0, /* tp_as_mapping */
3013 0, /* tp_hash */
3014 0, /* tp_call */
3015 0, /* tp_str */
3016 PyObject_GenericGetAttr, /* tp_getattro */
3017 0, /* tp_setattro */
3018 0, /* tp_as_buffer */
3019 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3020 Py_TPFLAGS_BASETYPE, /* tp_flags */
3021 filterfalse_doc, /* tp_doc */
3022 (traverseproc)filterfalse_traverse, /* tp_traverse */
3023 0, /* tp_clear */
3024 0, /* tp_richcompare */
3025 0, /* tp_weaklistoffset */
3026 PyObject_SelfIter, /* tp_iter */
3027 (iternextfunc)filterfalse_next, /* tp_iternext */
3028 0, /* tp_methods */
3029 0, /* tp_members */
3030 0, /* tp_getset */
3031 0, /* tp_base */
3032 0, /* tp_dict */
3033 0, /* tp_descr_get */
3034 0, /* tp_descr_set */
3035 0, /* tp_dictoffset */
3036 0, /* tp_init */
3037 0, /* tp_alloc */
3038 filterfalse_new, /* tp_new */
3039 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003040};
3041
3042
3043/* count object ************************************************************/
3044
3045typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003046 PyObject_HEAD
3047 Py_ssize_t cnt;
3048 PyObject *long_cnt;
3049 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003050} countobject;
3051
Raymond Hettinger30729212009-02-12 06:28:27 +00003052/* Counting logic and invariants:
3053
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003054fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003055
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003056 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3057 Advances with: cnt += 1
3058 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003059
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003060slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003061
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003062 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3063 All counting is done with python objects (no overflows or underflows).
3064 Advances with: long_cnt += long_step
3065 Step may be zero -- effectively a slow version of repeat(cnt).
3066 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003067*/
3068
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003069static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003070
3071static PyObject *
3072count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3073{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003074 countobject *lz;
3075 int slow_mode = 0;
3076 Py_ssize_t cnt = 0;
3077 PyObject *long_cnt = NULL;
3078 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003079 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003080 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003081
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003082 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3083 kwlist, &long_cnt, &long_step))
3084 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003085
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003086 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3087 (long_step != NULL && !PyNumber_Check(long_step))) {
3088 PyErr_SetString(PyExc_TypeError, "a number is required");
3089 return NULL;
3090 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003091
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003092 if (long_cnt != NULL) {
3093 cnt = PyLong_AsSsize_t(long_cnt);
3094 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3095 PyErr_Clear();
3096 slow_mode = 1;
3097 }
3098 Py_INCREF(long_cnt);
3099 } else {
3100 cnt = 0;
3101 long_cnt = PyLong_FromLong(0);
3102 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003103
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003104 /* If not specified, step defaults to 1 */
3105 if (long_step == NULL) {
3106 long_step = PyLong_FromLong(1);
3107 if (long_step == NULL) {
3108 Py_DECREF(long_cnt);
3109 return NULL;
3110 }
3111 } else
3112 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003113
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003114 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003115
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003116 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003117 step = PyLong_AsLong(long_step);
3118 if (step != 1) {
3119 slow_mode = 1;
3120 if (step == -1 && PyErr_Occurred())
3121 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003122 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003123
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003124 if (slow_mode)
3125 cnt = PY_SSIZE_T_MAX;
3126 else
3127 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003128
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003129 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3130 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3131 assert(slow_mode ||
3132 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003133
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003134 /* create countobject structure */
3135 lz = (countobject *)type->tp_alloc(type, 0);
3136 if (lz == NULL) {
3137 Py_XDECREF(long_cnt);
3138 return NULL;
3139 }
3140 lz->cnt = cnt;
3141 lz->long_cnt = long_cnt;
3142 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003143
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003144 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003145}
3146
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003147static void
3148count_dealloc(countobject *lz)
3149{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003150 PyObject_GC_UnTrack(lz);
3151 Py_XDECREF(lz->long_cnt);
3152 Py_XDECREF(lz->long_step);
3153 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003154}
3155
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003156static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003157count_traverse(countobject *lz, visitproc visit, void *arg)
3158{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003159 Py_VISIT(lz->long_cnt);
3160 Py_VISIT(lz->long_step);
3161 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003162}
3163
3164static PyObject *
3165count_nextlong(countobject *lz)
3166{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003167 PyObject *long_cnt;
3168 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003169
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003170 long_cnt = lz->long_cnt;
3171 if (long_cnt == NULL) {
3172 /* Switch to slow_mode */
3173 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3174 if (long_cnt == NULL)
3175 return NULL;
3176 }
3177 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00003178
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003179 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3180 if (stepped_up == NULL)
3181 return NULL;
3182 lz->long_cnt = stepped_up;
3183 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003184}
3185
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003186static PyObject *
3187count_next(countobject *lz)
3188{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003189 if (lz->cnt == PY_SSIZE_T_MAX)
3190 return count_nextlong(lz);
3191 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003192}
3193
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003194static PyObject *
3195count_repr(countobject *lz)
3196{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003197 if (lz->cnt != PY_SSIZE_T_MAX)
3198 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003199
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003200 if (PyLong_Check(lz->long_step)) {
3201 long step = PyLong_AsLong(lz->long_step);
3202 if (step == -1 && PyErr_Occurred()) {
3203 PyErr_Clear();
3204 }
3205 if (step == 1) {
3206 /* Don't display step when it is an integer equal to 1 */
3207 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
3208 }
3209 }
3210 return PyUnicode_FromFormat("count(%R, %R)",
3211 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003212}
3213
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003214static PyObject *
3215count_reduce(countobject *lz)
3216{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003217 if (lz->cnt == PY_SSIZE_T_MAX)
3218 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3219 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003220}
3221
3222PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3223
3224static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003225 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3226 count_reduce_doc},
3227 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003228};
3229
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003230PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00003231 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003232\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003233Return a count object whose .__next__() method returns consecutive values.\n\
3234Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00003235 def count(firstval=0, step=1):\n\
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003236 x = firstval\n\
3237 while 1:\n\
3238 yield x\n\
3239 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003240
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003241static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003242 PyVarObject_HEAD_INIT(NULL, 0)
3243 "itertools.count", /* tp_name */
3244 sizeof(countobject), /* tp_basicsize */
3245 0, /* tp_itemsize */
3246 /* methods */
3247 (destructor)count_dealloc, /* tp_dealloc */
3248 0, /* tp_print */
3249 0, /* tp_getattr */
3250 0, /* tp_setattr */
3251 0, /* tp_reserved */
3252 (reprfunc)count_repr, /* tp_repr */
3253 0, /* tp_as_number */
3254 0, /* tp_as_sequence */
3255 0, /* tp_as_mapping */
3256 0, /* tp_hash */
3257 0, /* tp_call */
3258 0, /* tp_str */
3259 PyObject_GenericGetAttr, /* tp_getattro */
3260 0, /* tp_setattro */
3261 0, /* tp_as_buffer */
3262 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3263 Py_TPFLAGS_BASETYPE, /* tp_flags */
3264 count_doc, /* tp_doc */
3265 (traverseproc)count_traverse, /* tp_traverse */
3266 0, /* tp_clear */
3267 0, /* tp_richcompare */
3268 0, /* tp_weaklistoffset */
3269 PyObject_SelfIter, /* tp_iter */
3270 (iternextfunc)count_next, /* tp_iternext */
3271 count_methods, /* tp_methods */
3272 0, /* tp_members */
3273 0, /* tp_getset */
3274 0, /* tp_base */
3275 0, /* tp_dict */
3276 0, /* tp_descr_get */
3277 0, /* tp_descr_set */
3278 0, /* tp_dictoffset */
3279 0, /* tp_init */
3280 0, /* tp_alloc */
3281 count_new, /* tp_new */
3282 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003283};
3284
3285
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003286/* repeat object ************************************************************/
3287
3288typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003289 PyObject_HEAD
3290 PyObject *element;
3291 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003292} repeatobject;
3293
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003294static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003295
3296static PyObject *
3297repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3298{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003299 repeatobject *ro;
3300 PyObject *element;
3301 Py_ssize_t cnt = -1;
3302 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003303
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003304 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3305 &element, &cnt))
3306 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003307
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003308 if (PyTuple_Size(args) == 2 && cnt < 0)
3309 cnt = 0;
3310
3311 ro = (repeatobject *)type->tp_alloc(type, 0);
3312 if (ro == NULL)
3313 return NULL;
3314 Py_INCREF(element);
3315 ro->element = element;
3316 ro->cnt = cnt;
3317 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003318}
3319
3320static void
3321repeat_dealloc(repeatobject *ro)
3322{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003323 PyObject_GC_UnTrack(ro);
3324 Py_XDECREF(ro->element);
3325 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003326}
3327
3328static int
3329repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3330{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003331 Py_VISIT(ro->element);
3332 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003333}
3334
3335static PyObject *
3336repeat_next(repeatobject *ro)
3337{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003338 if (ro->cnt == 0)
3339 return NULL;
3340 if (ro->cnt > 0)
3341 ro->cnt--;
3342 Py_INCREF(ro->element);
3343 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003344}
3345
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003346static PyObject *
3347repeat_repr(repeatobject *ro)
3348{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003349 if (ro->cnt == -1)
3350 return PyUnicode_FromFormat("repeat(%R)", ro->element);
3351 else
3352 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
3353}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003354
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003355static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003356repeat_len(repeatobject *ro)
3357{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003358 if (ro->cnt == -1) {
3359 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3360 return NULL;
3361 }
3362 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003363}
3364
Armin Rigof5b3e362006-02-11 21:32:43 +00003365PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003366
3367static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003368 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3369 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003370};
3371
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003372PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00003373"repeat(object [,times]) -> create an iterator which returns the object\n\
3374for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003375endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003376
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003377static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003378 PyVarObject_HEAD_INIT(NULL, 0)
3379 "itertools.repeat", /* tp_name */
3380 sizeof(repeatobject), /* tp_basicsize */
3381 0, /* tp_itemsize */
3382 /* methods */
3383 (destructor)repeat_dealloc, /* tp_dealloc */
3384 0, /* tp_print */
3385 0, /* tp_getattr */
3386 0, /* tp_setattr */
3387 0, /* tp_reserved */
3388 (reprfunc)repeat_repr, /* tp_repr */
3389 0, /* tp_as_number */
3390 0, /* tp_as_sequence */
3391 0, /* tp_as_mapping */
3392 0, /* tp_hash */
3393 0, /* tp_call */
3394 0, /* tp_str */
3395 PyObject_GenericGetAttr, /* tp_getattro */
3396 0, /* tp_setattro */
3397 0, /* tp_as_buffer */
3398 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3399 Py_TPFLAGS_BASETYPE, /* tp_flags */
3400 repeat_doc, /* tp_doc */
3401 (traverseproc)repeat_traverse, /* tp_traverse */
3402 0, /* tp_clear */
3403 0, /* tp_richcompare */
3404 0, /* tp_weaklistoffset */
3405 PyObject_SelfIter, /* tp_iter */
3406 (iternextfunc)repeat_next, /* tp_iternext */
3407 repeat_methods, /* tp_methods */
3408 0, /* tp_members */
3409 0, /* tp_getset */
3410 0, /* tp_base */
3411 0, /* tp_dict */
3412 0, /* tp_descr_get */
3413 0, /* tp_descr_set */
3414 0, /* tp_dictoffset */
3415 0, /* tp_init */
3416 0, /* tp_alloc */
3417 repeat_new, /* tp_new */
3418 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003419};
3420
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003421/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00003422
3423#include "Python.h"
3424
3425typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003426 PyObject_HEAD
3427 Py_ssize_t tuplesize;
3428 Py_ssize_t numactive;
3429 PyObject *ittuple; /* tuple of iterators */
3430 PyObject *result;
3431 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003432} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003433
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003434static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003435
3436static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003437zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003438{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003439 ziplongestobject *lz;
3440 Py_ssize_t i;
3441 PyObject *ittuple; /* tuple of iterators */
3442 PyObject *result;
3443 PyObject *fillvalue = Py_None;
3444 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003445
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003446 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3447 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3448 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3449 PyErr_SetString(PyExc_TypeError,
3450 "zip_longest() got an unexpected keyword argument");
3451 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003452 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003453 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003454
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003455 /* args must be a tuple */
3456 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00003457
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003458 /* obtain iterators */
3459 ittuple = PyTuple_New(tuplesize);
3460 if (ittuple == NULL)
3461 return NULL;
3462 for (i=0; i < tuplesize; ++i) {
3463 PyObject *item = PyTuple_GET_ITEM(args, i);
3464 PyObject *it = PyObject_GetIter(item);
3465 if (it == NULL) {
3466 if (PyErr_ExceptionMatches(PyExc_TypeError))
3467 PyErr_Format(PyExc_TypeError,
3468 "zip_longest argument #%zd must support iteration",
3469 i+1);
3470 Py_DECREF(ittuple);
3471 return NULL;
3472 }
3473 PyTuple_SET_ITEM(ittuple, i, it);
3474 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003475
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003476 /* create a result holder */
3477 result = PyTuple_New(tuplesize);
3478 if (result == NULL) {
3479 Py_DECREF(ittuple);
3480 return NULL;
3481 }
3482 for (i=0 ; i < tuplesize ; i++) {
3483 Py_INCREF(Py_None);
3484 PyTuple_SET_ITEM(result, i, Py_None);
3485 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00003486
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003487 /* create ziplongestobject structure */
3488 lz = (ziplongestobject *)type->tp_alloc(type, 0);
3489 if (lz == NULL) {
3490 Py_DECREF(ittuple);
3491 Py_DECREF(result);
3492 return NULL;
3493 }
3494 lz->ittuple = ittuple;
3495 lz->tuplesize = tuplesize;
3496 lz->numactive = tuplesize;
3497 lz->result = result;
3498 Py_INCREF(fillvalue);
3499 lz->fillvalue = fillvalue;
3500 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003501}
3502
3503static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003504zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003505{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003506 PyObject_GC_UnTrack(lz);
3507 Py_XDECREF(lz->ittuple);
3508 Py_XDECREF(lz->result);
3509 Py_XDECREF(lz->fillvalue);
3510 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003511}
3512
3513static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003514zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003515{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003516 Py_VISIT(lz->ittuple);
3517 Py_VISIT(lz->result);
3518 Py_VISIT(lz->fillvalue);
3519 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003520}
3521
3522static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003523zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003524{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003525 Py_ssize_t i;
3526 Py_ssize_t tuplesize = lz->tuplesize;
3527 PyObject *result = lz->result;
3528 PyObject *it;
3529 PyObject *item;
3530 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003531
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003532 if (tuplesize == 0)
3533 return NULL;
3534 if (lz->numactive == 0)
3535 return NULL;
3536 if (Py_REFCNT(result) == 1) {
3537 Py_INCREF(result);
3538 for (i=0 ; i < tuplesize ; i++) {
3539 it = PyTuple_GET_ITEM(lz->ittuple, i);
3540 if (it == NULL) {
3541 Py_INCREF(lz->fillvalue);
3542 item = lz->fillvalue;
3543 } else {
3544 item = PyIter_Next(it);
3545 if (item == NULL) {
3546 lz->numactive -= 1;
3547 if (lz->numactive == 0 || PyErr_Occurred()) {
3548 lz->numactive = 0;
3549 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00003550 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003551 } else {
3552 Py_INCREF(lz->fillvalue);
3553 item = lz->fillvalue;
3554 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3555 Py_DECREF(it);
3556 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00003557 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003558 }
3559 olditem = PyTuple_GET_ITEM(result, i);
3560 PyTuple_SET_ITEM(result, i, item);
3561 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00003562 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003563 } else {
3564 result = PyTuple_New(tuplesize);
3565 if (result == NULL)
3566 return NULL;
3567 for (i=0 ; i < tuplesize ; i++) {
3568 it = PyTuple_GET_ITEM(lz->ittuple, i);
3569 if (it == NULL) {
3570 Py_INCREF(lz->fillvalue);
3571 item = lz->fillvalue;
3572 } else {
3573 item = PyIter_Next(it);
3574 if (item == NULL) {
3575 lz->numactive -= 1;
3576 if (lz->numactive == 0 || PyErr_Occurred()) {
3577 lz->numactive = 0;
3578 Py_DECREF(result);
3579 return NULL;
3580 } else {
3581 Py_INCREF(lz->fillvalue);
3582 item = lz->fillvalue;
3583 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3584 Py_DECREF(it);
3585 }
3586 }
3587 }
3588 PyTuple_SET_ITEM(result, i, item);
3589 }
3590 }
3591 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003592}
3593
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003594PyDoc_STRVAR(zip_longest_doc,
3595"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003596\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003597Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00003598the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003599method continues until the longest iterable in the argument sequence\n\
3600is exhausted and then it raises StopIteration. When the shorter iterables\n\
3601are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3602defaults to None or can be specified by a keyword argument.\n\
3603");
3604
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003605static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003606 PyVarObject_HEAD_INIT(NULL, 0)
3607 "itertools.zip_longest", /* tp_name */
3608 sizeof(ziplongestobject), /* tp_basicsize */
3609 0, /* tp_itemsize */
3610 /* methods */
3611 (destructor)zip_longest_dealloc, /* tp_dealloc */
3612 0, /* tp_print */
3613 0, /* tp_getattr */
3614 0, /* tp_setattr */
3615 0, /* tp_reserved */
3616 0, /* tp_repr */
3617 0, /* tp_as_number */
3618 0, /* tp_as_sequence */
3619 0, /* tp_as_mapping */
3620 0, /* tp_hash */
3621 0, /* tp_call */
3622 0, /* tp_str */
3623 PyObject_GenericGetAttr, /* tp_getattro */
3624 0, /* tp_setattro */
3625 0, /* tp_as_buffer */
3626 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3627 Py_TPFLAGS_BASETYPE, /* tp_flags */
3628 zip_longest_doc, /* tp_doc */
3629 (traverseproc)zip_longest_traverse, /* tp_traverse */
3630 0, /* tp_clear */
3631 0, /* tp_richcompare */
3632 0, /* tp_weaklistoffset */
3633 PyObject_SelfIter, /* tp_iter */
3634 (iternextfunc)zip_longest_next, /* tp_iternext */
3635 0, /* tp_methods */
3636 0, /* tp_members */
3637 0, /* tp_getset */
3638 0, /* tp_base */
3639 0, /* tp_dict */
3640 0, /* tp_descr_get */
3641 0, /* tp_descr_set */
3642 0, /* tp_dictoffset */
3643 0, /* tp_init */
3644 0, /* tp_alloc */
3645 zip_longest_new, /* tp_new */
3646 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003647};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003648
3649/* module level code ********************************************************/
3650
3651PyDoc_STRVAR(module_doc,
3652"Functional tools for creating and using iterators.\n\
3653\n\
3654Infinite iterators:\n\
3655count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003656cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003657repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003658\n\
3659Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003660accumulate(p, start=0) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003661chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
3662compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
3663dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
3664groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003665filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003666islice(seq, [start,] stop [, step]) --> elements from\n\
3667 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003668starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003669tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003670takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003671zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00003672\n\
3673Combinatoric generators:\n\
3674product(p, q, ... [repeat=1]) --> cartesian product\n\
3675permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00003676combinations(p, r)\n\
3677combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003678");
3679
3680
Raymond Hettingerad983e72003-11-12 14:32:26 +00003681static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003682 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3683 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00003684};
3685
Martin v. Löwis1a214512008-06-11 05:26:20 +00003686
3687static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003688 PyModuleDef_HEAD_INIT,
3689 "itertools",
3690 module_doc,
3691 -1,
3692 module_methods,
3693 NULL,
3694 NULL,
3695 NULL,
3696 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00003697};
3698
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003699PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00003700PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003701{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003702 int i;
3703 PyObject *m;
3704 char *name;
3705 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00003706 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003707 &combinations_type,
3708 &cwr_type,
3709 &cycle_type,
3710 &dropwhile_type,
3711 &takewhile_type,
3712 &islice_type,
3713 &starmap_type,
3714 &chain_type,
3715 &compress_type,
3716 &filterfalse_type,
3717 &count_type,
3718 &ziplongest_type,
3719 &permutations_type,
3720 &product_type,
3721 &repeat_type,
3722 &groupby_type,
3723 NULL
3724 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00003725
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003726 Py_TYPE(&teedataobject_type) = &PyType_Type;
3727 m = PyModule_Create(&itertoolsmodule);
3728 if (m == NULL)
3729 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003730
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003731 for (i=0 ; typelist[i] != NULL ; i++) {
3732 if (PyType_Ready(typelist[i]) < 0)
3733 return NULL;
3734 name = strchr(typelist[i]->tp_name, '.');
3735 assert (name != NULL);
3736 Py_INCREF(typelist[i]);
3737 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
3738 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003739
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003740 if (PyType_Ready(&teedataobject_type) < 0)
3741 return NULL;
3742 if (PyType_Ready(&tee_type) < 0)
3743 return NULL;
3744 if (PyType_Ready(&_grouper_type) < 0)
3745 return NULL;
3746 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003747}