blob: a7cd3f4a644403f082646303cbd1bafcc8f4cec4 [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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +000098 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000101
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +0000114 tmp = gbo->currkey;
115 gbo->currkey = newkey;
116 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000117
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000118 tmp = gbo->currvalue;
119 gbo->currvalue = newvalue;
120 Py_XDECREF(tmp);
121 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +0000128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000131
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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_compare */
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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +0000199 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000200
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +0000216 PyObject_GC_UnTrack(igo);
217 Py_DECREF(igo->parent);
218 Py_DECREF(igo->tgtkey);
219 PyObject_GC_Del(igo);
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000220}
221
222static int
223_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
224{
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +0000233 groupbyobject *gbo = (groupbyobject *)igo->parent;
234 PyObject *newvalue, *newkey, *r;
235 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000236
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +0000254 assert(gbo->currkey == NULL);
255 gbo->currkey = newkey;
256 gbo->currvalue = newvalue;
257 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +0000265 r = gbo->currvalue;
266 gbo->currvalue = NULL;
267 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000268
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000269 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000270}
271
272static PyTypeObject _grouper_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +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_compare */
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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +0000349 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000350
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000351 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
352 if (tdo == NULL)
353 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000354
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +0000375 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000376
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Wouters19bf33b2006-03-27 21:02:13 +0000393static int
394teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
395{
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Wouters19bf33b2006-03-27 21:02:13 +0000402}
403
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +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 Wouters19bf33b2006-03-27 21:02:13 +0000417static int
418teedataobject_clear(teedataobject *tdo)
419{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000420 int i;
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200421 PyObject *tmp;
422
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000423 Py_CLEAR(tdo->it);
424 for (i=0 ; i<tdo->numread ; i++)
425 Py_CLEAR(tdo->values[i]);
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200426 tmp = tdo->nextlink;
427 tdo->nextlink = NULL;
428 teedataobject_safe_decref(tmp);
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000429 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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_compare */
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 Pitrouc83ea132010-05-09 14:46:46 +0000491 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000492
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000493 if (to->index >= LINKCELLS) {
494 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200495 if (link == NULL)
496 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Wouters19bf33b2006-03-27 21:02:13 +0000508static int
509tee_traverse(teeobject *to, visitproc visit, void *arg)
510{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000511 Py_VISIT((PyObject *)to->dataobj);
512 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000513}
514
Raymond Hettingerad983e72003-11-12 14:32:26 +0000515static PyObject *
516tee_copy(teeobject *to)
517{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000518 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000519
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +0000536 teeobject *to;
537 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000538
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 }
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000556
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000557 to->index = 0;
558 to->weakreflist = NULL;
559 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000560done:
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +0000568 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000569
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Wouters19bf33b2006-03-27 21:02:13 +0000575static int
576tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000577{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000578 if (to->weakreflist != NULL)
579 PyObject_ClearWeakRefs((PyObject *) to);
580 Py_CLEAR(to->dataobj);
581 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000582}
583
584static void
585tee_dealloc(teeobject *to)
586{
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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_compare */
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 Pitrouc83ea132010-05-09 14:46:46 +0000646 Py_ssize_t i, n=2;
647 PyObject *it, *iterable, *copyable, *result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000648
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +0000704 PyObject *it;
705 PyObject *iterable;
706 PyObject *saved;
707 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000708
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000709 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
710 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000711
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000712 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
713 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000714
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000715 /* Get iterator. */
716 it = PyObject_GetIter(iterable);
717 if (it == NULL)
718 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000719
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +0000737 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000738}
739
740static void
741cycle_dealloc(cycleobject *lz)
742{
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +0000760 PyObject *item;
761 PyObject *it;
762 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000763
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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_compare */
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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +0000856 PyObject *func, *seq;
857 PyObject *it;
858 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000859
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000860 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
861 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000862
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000863 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
864 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000865
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000866 /* Get iterator. */
867 it = PyObject_GetIter(seq);
868 if (it == NULL)
869 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000870
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +0000882 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000883}
884
885static void
886dropwhile_dealloc(dropwhileobject *lz)
887{
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc5bef752012-08-15 23:16:51 +0200925 if (ok == 0) {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000926 lz->start = 1;
927 return item;
928 }
929 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +0200930 if (ok < 0)
931 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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_compare */
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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +00001000 PyObject *func, *seq;
1001 PyObject *it;
1002 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001003
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001004 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1005 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001006
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001007 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1008 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001009
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001010 /* Get iterator. */
1011 it = PyObject_GetIter(seq);
1012 if (it == NULL)
1013 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001014
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +00001026 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001027}
1028
1029static void
1030takewhile_dealloc(takewhileobject *lz)
1031{
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +00001049 PyObject *item, *good;
1050 PyObject *it = lz->it;
1051 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001052
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001053 if (lz->stop == 1)
1054 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001055
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001056 item = (*Py_TYPE(it)->tp_iternext)(it);
1057 if (item == NULL)
1058 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001059
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc5bef752012-08-15 23:16:51 +02001067 if (ok > 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001068 return item;
1069 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02001070 if (ok == 0)
1071 lz->stop = 1;
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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_compare */
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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +00001148 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1149 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001150
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001151 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1152 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001153
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001154 numargs = PyTuple_Size(args);
1155 if (numargs == 2) {
1156 if (a1 != Py_None) {
1157 stop = PyInt_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 <= maxint.");
1163 return NULL;
1164 }
1165 }
1166 } else {
1167 if (a1 != Py_None)
1168 start = PyInt_AsSsize_t(a1);
1169 if (start == -1 && PyErr_Occurred())
1170 PyErr_Clear();
1171 if (a2 != Py_None) {
1172 stop = PyInt_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 <= maxint.");
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 <= maxint.");
1185 return NULL;
1186 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001187
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001188 if (a3 != NULL) {
1189 if (a3 != Py_None)
1190 step = PyInt_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 Pitrouc83ea132010-05-09 14:46:46 +00001200 /* Get iterator. */
1201 it = PyObject_GetIter(seq);
1202 if (it == NULL)
1203 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +00001217 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001218}
1219
1220static void
1221islice_dealloc(isliceobject *lz)
1222{
Antoine Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +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 Pitrouc83ea132010-05-09 14:46:46 +00001238 PyObject *item;
1239 PyObject *it = lz->it;
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001240 Py_ssize_t stop = lz->stop;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001241 Py_ssize_t oldnext;
1242 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001243
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001244 if (it == NULL)
1245 return NULL;
1246
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001247 iternext = *Py_TYPE(it)->tp_iternext;
1248 while (lz->cnt < lz->next) {
1249 item = iternext(it);
1250 if (item == NULL)
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001251 goto empty;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001252 Py_DECREF(item);
1253 lz->cnt++;
1254 }
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001255 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001256 goto empty;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001257 item = iternext(it);
1258 if (item == NULL)
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001259 goto empty;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001260 lz->cnt++;
1261 oldnext = lz->next;
Mark Dickinsona96b0d12011-09-24 09:01:16 +01001262 /* The (size_t) cast below avoids the danger of undefined
1263 behaviour from signed integer overflow. */
1264 lz->next += (size_t)lz->step;
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001265 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1266 lz->next = stop;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001267 return item;
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001268
1269empty:
1270 Py_CLEAR(lz->it);
1271 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001272}
1273
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001274PyDoc_STRVAR(islice_doc,
1275"islice(iterable, [start,] stop [, step]) --> islice object\n\
1276\n\
1277Return an iterator whose next() method returns selected values from an\n\
1278iterable. If start is specified, will skip all preceding elements;\n\
1279otherwise, start defaults to zero. Step defaults to one. If\n\
1280specified as another value, step determines how many values are \n\
1281skipped between successive calls. Works like a slice() on a list\n\
1282but returns an iterator.");
1283
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001284static PyTypeObject islice_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001285 PyVarObject_HEAD_INIT(NULL, 0)
1286 "itertools.islice", /* tp_name */
1287 sizeof(isliceobject), /* tp_basicsize */
1288 0, /* tp_itemsize */
1289 /* methods */
1290 (destructor)islice_dealloc, /* tp_dealloc */
1291 0, /* tp_print */
1292 0, /* tp_getattr */
1293 0, /* tp_setattr */
1294 0, /* tp_compare */
1295 0, /* tp_repr */
1296 0, /* tp_as_number */
1297 0, /* tp_as_sequence */
1298 0, /* tp_as_mapping */
1299 0, /* tp_hash */
1300 0, /* tp_call */
1301 0, /* tp_str */
1302 PyObject_GenericGetAttr, /* tp_getattro */
1303 0, /* tp_setattro */
1304 0, /* tp_as_buffer */
1305 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1306 Py_TPFLAGS_BASETYPE, /* tp_flags */
1307 islice_doc, /* tp_doc */
1308 (traverseproc)islice_traverse, /* tp_traverse */
1309 0, /* tp_clear */
1310 0, /* tp_richcompare */
1311 0, /* tp_weaklistoffset */
1312 PyObject_SelfIter, /* tp_iter */
1313 (iternextfunc)islice_next, /* tp_iternext */
1314 0, /* tp_methods */
1315 0, /* tp_members */
1316 0, /* tp_getset */
1317 0, /* tp_base */
1318 0, /* tp_dict */
1319 0, /* tp_descr_get */
1320 0, /* tp_descr_set */
1321 0, /* tp_dictoffset */
1322 0, /* tp_init */
1323 0, /* tp_alloc */
1324 islice_new, /* tp_new */
1325 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001326};
1327
1328
1329/* starmap object ************************************************************/
1330
1331typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001332 PyObject_HEAD
1333 PyObject *func;
1334 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001335} starmapobject;
1336
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001337static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001338
1339static PyObject *
1340starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1341{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001342 PyObject *func, *seq;
1343 PyObject *it;
1344 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001345
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001346 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1347 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001348
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001349 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1350 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001351
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001352 /* Get iterator. */
1353 it = PyObject_GetIter(seq);
1354 if (it == NULL)
1355 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001356
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001357 /* create starmapobject structure */
1358 lz = (starmapobject *)type->tp_alloc(type, 0);
1359 if (lz == NULL) {
1360 Py_DECREF(it);
1361 return NULL;
1362 }
1363 Py_INCREF(func);
1364 lz->func = func;
1365 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001366
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001367 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001368}
1369
1370static void
1371starmap_dealloc(starmapobject *lz)
1372{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001373 PyObject_GC_UnTrack(lz);
1374 Py_XDECREF(lz->func);
1375 Py_XDECREF(lz->it);
1376 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001377}
1378
1379static int
1380starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1381{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001382 Py_VISIT(lz->it);
1383 Py_VISIT(lz->func);
1384 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001385}
1386
1387static PyObject *
1388starmap_next(starmapobject *lz)
1389{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001390 PyObject *args;
1391 PyObject *result;
1392 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001393
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001394 args = (*Py_TYPE(it)->tp_iternext)(it);
1395 if (args == NULL)
1396 return NULL;
1397 if (!PyTuple_CheckExact(args)) {
1398 PyObject *newargs = PySequence_Tuple(args);
1399 Py_DECREF(args);
1400 if (newargs == NULL)
1401 return NULL;
1402 args = newargs;
1403 }
1404 result = PyObject_Call(lz->func, args, NULL);
1405 Py_DECREF(args);
1406 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001407}
1408
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001409PyDoc_STRVAR(starmap_doc,
1410"starmap(function, sequence) --> starmap object\n\
1411\n\
1412Return an iterator whose values are returned from the function evaluated\n\
1413with a argument tuple taken from the given sequence.");
1414
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001415static PyTypeObject starmap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001416 PyVarObject_HEAD_INIT(NULL, 0)
1417 "itertools.starmap", /* tp_name */
1418 sizeof(starmapobject), /* tp_basicsize */
1419 0, /* tp_itemsize */
1420 /* methods */
1421 (destructor)starmap_dealloc, /* tp_dealloc */
1422 0, /* tp_print */
1423 0, /* tp_getattr */
1424 0, /* tp_setattr */
1425 0, /* tp_compare */
1426 0, /* tp_repr */
1427 0, /* tp_as_number */
1428 0, /* tp_as_sequence */
1429 0, /* tp_as_mapping */
1430 0, /* tp_hash */
1431 0, /* tp_call */
1432 0, /* tp_str */
1433 PyObject_GenericGetAttr, /* tp_getattro */
1434 0, /* tp_setattro */
1435 0, /* tp_as_buffer */
1436 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1437 Py_TPFLAGS_BASETYPE, /* tp_flags */
1438 starmap_doc, /* tp_doc */
1439 (traverseproc)starmap_traverse, /* tp_traverse */
1440 0, /* tp_clear */
1441 0, /* tp_richcompare */
1442 0, /* tp_weaklistoffset */
1443 PyObject_SelfIter, /* tp_iter */
1444 (iternextfunc)starmap_next, /* tp_iternext */
1445 0, /* tp_methods */
1446 0, /* tp_members */
1447 0, /* tp_getset */
1448 0, /* tp_base */
1449 0, /* tp_dict */
1450 0, /* tp_descr_get */
1451 0, /* tp_descr_set */
1452 0, /* tp_dictoffset */
1453 0, /* tp_init */
1454 0, /* tp_alloc */
1455 starmap_new, /* tp_new */
1456 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457};
1458
1459
1460/* imap object ************************************************************/
1461
1462typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001463 PyObject_HEAD
1464 PyObject *iters;
1465 PyObject *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001466} imapobject;
1467
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001468static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469
1470static PyObject *
1471imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1472{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001473 PyObject *it, *iters, *func;
1474 imapobject *lz;
1475 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001476
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001477 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
1478 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001479
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001480 numargs = PyTuple_Size(args);
1481 if (numargs < 2) {
1482 PyErr_SetString(PyExc_TypeError,
1483 "imap() must have at least two arguments.");
1484 return NULL;
1485 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001486
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001487 iters = PyTuple_New(numargs-1);
1488 if (iters == NULL)
1489 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001490
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001491 for (i=1 ; i<numargs ; i++) {
1492 /* Get iterator. */
1493 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1494 if (it == NULL) {
1495 Py_DECREF(iters);
1496 return NULL;
1497 }
1498 PyTuple_SET_ITEM(iters, i-1, it);
1499 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001500
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001501 /* create imapobject structure */
1502 lz = (imapobject *)type->tp_alloc(type, 0);
1503 if (lz == NULL) {
1504 Py_DECREF(iters);
1505 return NULL;
1506 }
1507 lz->iters = iters;
1508 func = PyTuple_GET_ITEM(args, 0);
1509 Py_INCREF(func);
1510 lz->func = func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001511
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001512 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001513}
1514
1515static void
1516imap_dealloc(imapobject *lz)
1517{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001518 PyObject_GC_UnTrack(lz);
1519 Py_XDECREF(lz->iters);
1520 Py_XDECREF(lz->func);
1521 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001522}
1523
1524static int
1525imap_traverse(imapobject *lz, visitproc visit, void *arg)
1526{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001527 Py_VISIT(lz->iters);
1528 Py_VISIT(lz->func);
1529 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001530}
1531
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001532/*
Raymond Hettinger2012f172003-02-07 05:32:58 +00001533imap() is an iterator version of __builtins__.map() except that it does
1534not have the None fill-in feature. That was intentionally left out for
1535the following reasons:
1536
1537 1) Itertools are designed to be easily combined and chained together.
1538 Having all tools stop with the shortest input is a unifying principle
1539 that makes it easier to combine finite iterators (supplying data) with
1540 infinite iterators like count() and repeat() (for supplying sequential
1541 or constant arguments to a function).
1542
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001543 2) In typical use cases for combining itertools, having one finite data
Raymond Hettinger2012f172003-02-07 05:32:58 +00001544 supplier run out before another is likely to be an error condition which
1545 should not pass silently by automatically supplying None.
1546
1547 3) The use cases for automatic None fill-in are rare -- not many functions
1548 do something useful when a parameter suddenly switches type and becomes
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001549 None.
Raymond Hettinger2012f172003-02-07 05:32:58 +00001550
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001551 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001552 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001553
1554 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1555*/
1556
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001557static PyObject *
1558imap_next(imapobject *lz)
1559{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001560 PyObject *val;
1561 PyObject *argtuple;
1562 PyObject *result;
1563 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001564
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001565 numargs = PyTuple_Size(lz->iters);
1566 argtuple = PyTuple_New(numargs);
1567 if (argtuple == NULL)
1568 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001569
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001570 for (i=0 ; i<numargs ; i++) {
1571 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1572 if (val == NULL) {
1573 Py_DECREF(argtuple);
1574 return NULL;
1575 }
1576 PyTuple_SET_ITEM(argtuple, i, val);
1577 }
1578 if (lz->func == Py_None)
1579 return argtuple;
1580 result = PyObject_Call(lz->func, argtuple, NULL);
1581 Py_DECREF(argtuple);
1582 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001583}
1584
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001585PyDoc_STRVAR(imap_doc,
1586"imap(func, *iterables) --> imap object\n\
1587\n\
1588Make an iterator that computes the function using arguments from\n\
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001589each of the iterables. Like map() except that it returns\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001590an iterator instead of a list and that it stops when the shortest\n\
1591iterable is exhausted instead of filling in None for shorter\n\
1592iterables.");
1593
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001594static PyTypeObject imap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001595 PyVarObject_HEAD_INIT(NULL, 0)
1596 "itertools.imap", /* tp_name */
1597 sizeof(imapobject), /* tp_basicsize */
1598 0, /* tp_itemsize */
1599 /* methods */
1600 (destructor)imap_dealloc, /* tp_dealloc */
1601 0, /* tp_print */
1602 0, /* tp_getattr */
1603 0, /* tp_setattr */
1604 0, /* tp_compare */
1605 0, /* tp_repr */
1606 0, /* tp_as_number */
1607 0, /* tp_as_sequence */
1608 0, /* tp_as_mapping */
1609 0, /* tp_hash */
1610 0, /* tp_call */
1611 0, /* tp_str */
1612 PyObject_GenericGetAttr, /* tp_getattro */
1613 0, /* tp_setattro */
1614 0, /* tp_as_buffer */
1615 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1616 Py_TPFLAGS_BASETYPE, /* tp_flags */
1617 imap_doc, /* tp_doc */
1618 (traverseproc)imap_traverse, /* tp_traverse */
1619 0, /* tp_clear */
1620 0, /* tp_richcompare */
1621 0, /* tp_weaklistoffset */
1622 PyObject_SelfIter, /* tp_iter */
1623 (iternextfunc)imap_next, /* tp_iternext */
1624 0, /* tp_methods */
1625 0, /* tp_members */
1626 0, /* tp_getset */
1627 0, /* tp_base */
1628 0, /* tp_dict */
1629 0, /* tp_descr_get */
1630 0, /* tp_descr_set */
1631 0, /* tp_dictoffset */
1632 0, /* tp_init */
1633 0, /* tp_alloc */
1634 imap_new, /* tp_new */
1635 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001636};
1637
1638
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001639/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001640
1641typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001642 PyObject_HEAD
1643 PyObject *source; /* Iterator over input iterables */
1644 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001645} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001646
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001647static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001648
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001649static PyObject *
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001650chain_new_internal(PyTypeObject *type, PyObject *source)
1651{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001652 chainobject *lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001653
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001654 lz = (chainobject *)type->tp_alloc(type, 0);
1655 if (lz == NULL) {
1656 Py_DECREF(source);
1657 return NULL;
1658 }
1659
1660 lz->source = source;
1661 lz->active = NULL;
1662 return (PyObject *)lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001663}
1664
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001665static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001666chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001667{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001668 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001669
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001670 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1671 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001672
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001673 source = PyObject_GetIter(args);
1674 if (source == NULL)
1675 return NULL;
1676
1677 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001678}
1679
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001680static PyObject *
1681chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1682{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001683 PyObject *source;
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001684
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001685 source = PyObject_GetIter(arg);
1686 if (source == NULL)
1687 return NULL;
1688
1689 return chain_new_internal(type, source);
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001690}
1691
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001692static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001693chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001694{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001695 PyObject_GC_UnTrack(lz);
1696 Py_XDECREF(lz->active);
1697 Py_XDECREF(lz->source);
1698 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001699}
1700
Raymond Hettinger2012f172003-02-07 05:32:58 +00001701static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001702chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001703{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001704 Py_VISIT(lz->source);
1705 Py_VISIT(lz->active);
1706 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001707}
1708
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001709static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001710chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001711{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001712 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001713
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001714 if (lz->source == NULL)
1715 return NULL; /* already stopped */
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001716
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001717 if (lz->active == NULL) {
1718 PyObject *iterable = PyIter_Next(lz->source);
1719 if (iterable == NULL) {
1720 Py_CLEAR(lz->source);
1721 return NULL; /* no more input sources */
1722 }
1723 lz->active = PyObject_GetIter(iterable);
1724 Py_DECREF(iterable);
1725 if (lz->active == NULL) {
1726 Py_CLEAR(lz->source);
1727 return NULL; /* input not iterable */
1728 }
1729 }
1730 item = PyIter_Next(lz->active);
1731 if (item != NULL)
1732 return item;
1733 if (PyErr_Occurred()) {
1734 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1735 PyErr_Clear();
1736 else
1737 return NULL; /* input raised an exception */
1738 }
1739 Py_CLEAR(lz->active);
1740 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001741}
1742
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001743PyDoc_STRVAR(chain_doc,
1744"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001745\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001746Return a chain object whose .next() method returns elements from the\n\
1747first iterable until it is exhausted, then elements from the next\n\
1748iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001749
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001750PyDoc_STRVAR(chain_from_iterable_doc,
1751"chain.from_iterable(iterable) --> chain object\n\
1752\n\
1753Alternate chain() contructor taking a single iterable argument\n\
1754that evaluates lazily.");
1755
1756static PyMethodDef chain_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001757 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1758 chain_from_iterable_doc},
1759 {NULL, NULL} /* sentinel */
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001760};
1761
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001762static PyTypeObject chain_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001763 PyVarObject_HEAD_INIT(NULL, 0)
1764 "itertools.chain", /* tp_name */
1765 sizeof(chainobject), /* tp_basicsize */
1766 0, /* tp_itemsize */
1767 /* methods */
1768 (destructor)chain_dealloc, /* tp_dealloc */
1769 0, /* tp_print */
1770 0, /* tp_getattr */
1771 0, /* tp_setattr */
1772 0, /* tp_compare */
1773 0, /* tp_repr */
1774 0, /* tp_as_number */
1775 0, /* tp_as_sequence */
1776 0, /* tp_as_mapping */
1777 0, /* tp_hash */
1778 0, /* tp_call */
1779 0, /* tp_str */
1780 PyObject_GenericGetAttr, /* tp_getattro */
1781 0, /* tp_setattro */
1782 0, /* tp_as_buffer */
1783 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1784 Py_TPFLAGS_BASETYPE, /* tp_flags */
1785 chain_doc, /* tp_doc */
1786 (traverseproc)chain_traverse, /* tp_traverse */
1787 0, /* tp_clear */
1788 0, /* tp_richcompare */
1789 0, /* tp_weaklistoffset */
1790 PyObject_SelfIter, /* tp_iter */
1791 (iternextfunc)chain_next, /* tp_iternext */
1792 chain_methods, /* tp_methods */
1793 0, /* tp_members */
1794 0, /* tp_getset */
1795 0, /* tp_base */
1796 0, /* tp_dict */
1797 0, /* tp_descr_get */
1798 0, /* tp_descr_set */
1799 0, /* tp_dictoffset */
1800 0, /* tp_init */
1801 0, /* tp_alloc */
1802 chain_new, /* tp_new */
1803 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001804};
1805
1806
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001807/* product object ************************************************************/
1808
1809typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001810 PyObject_HEAD
1811 PyObject *pools; /* tuple of pool tuples */
1812 Py_ssize_t *indices; /* one index per pool */
1813 PyObject *result; /* most recently returned result tuple */
1814 int stopped; /* set to 1 when the product iterator is exhausted */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001815} productobject;
1816
1817static PyTypeObject product_type;
1818
1819static PyObject *
1820product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1821{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001822 productobject *lz;
1823 Py_ssize_t nargs, npools, repeat=1;
1824 PyObject *pools = NULL;
1825 Py_ssize_t *indices = NULL;
1826 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001827
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001828 if (kwds != NULL) {
1829 char *kwlist[] = {"repeat", 0};
1830 PyObject *tmpargs = PyTuple_New(0);
1831 if (tmpargs == NULL)
1832 return NULL;
1833 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1834 Py_DECREF(tmpargs);
1835 return NULL;
1836 }
1837 Py_DECREF(tmpargs);
1838 if (repeat < 0) {
1839 PyErr_SetString(PyExc_ValueError,
1840 "repeat argument cannot be negative");
1841 return NULL;
1842 }
1843 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001844
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001845 assert(PyTuple_Check(args));
1846 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1847 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001848
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001849 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1850 if (indices == NULL) {
1851 PyErr_NoMemory();
1852 goto error;
1853 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001854
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001855 pools = PyTuple_New(npools);
1856 if (pools == NULL)
1857 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001858
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001859 for (i=0; i < nargs ; ++i) {
1860 PyObject *item = PyTuple_GET_ITEM(args, i);
1861 PyObject *pool = PySequence_Tuple(item);
1862 if (pool == NULL)
1863 goto error;
1864 PyTuple_SET_ITEM(pools, i, pool);
1865 indices[i] = 0;
1866 }
1867 for ( ; i < npools; ++i) {
1868 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1869 Py_INCREF(pool);
1870 PyTuple_SET_ITEM(pools, i, pool);
1871 indices[i] = 0;
1872 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001873
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001874 /* create productobject structure */
1875 lz = (productobject *)type->tp_alloc(type, 0);
1876 if (lz == NULL)
1877 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001878
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001879 lz->pools = pools;
1880 lz->indices = indices;
1881 lz->result = NULL;
1882 lz->stopped = 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001883
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001884 return (PyObject *)lz;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001885
1886error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001887 if (indices != NULL)
1888 PyMem_Free(indices);
1889 Py_XDECREF(pools);
1890 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001891}
1892
1893static void
1894product_dealloc(productobject *lz)
1895{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001896 PyObject_GC_UnTrack(lz);
1897 Py_XDECREF(lz->pools);
1898 Py_XDECREF(lz->result);
1899 if (lz->indices != NULL)
1900 PyMem_Free(lz->indices);
1901 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001902}
1903
1904static int
1905product_traverse(productobject *lz, visitproc visit, void *arg)
1906{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001907 Py_VISIT(lz->pools);
1908 Py_VISIT(lz->result);
1909 return 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001910}
1911
1912static PyObject *
1913product_next(productobject *lz)
1914{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001915 PyObject *pool;
1916 PyObject *elem;
1917 PyObject *oldelem;
1918 PyObject *pools = lz->pools;
1919 PyObject *result = lz->result;
1920 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1921 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001922
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001923 if (lz->stopped)
1924 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001925
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001926 if (result == NULL) {
1927 /* On the first pass, return an initial tuple filled with the
1928 first element from each pool. */
1929 result = PyTuple_New(npools);
1930 if (result == NULL)
1931 goto empty;
1932 lz->result = result;
1933 for (i=0; i < npools; i++) {
1934 pool = PyTuple_GET_ITEM(pools, i);
1935 if (PyTuple_GET_SIZE(pool) == 0)
1936 goto empty;
1937 elem = PyTuple_GET_ITEM(pool, 0);
1938 Py_INCREF(elem);
1939 PyTuple_SET_ITEM(result, i, elem);
1940 }
1941 } else {
1942 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001943
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001944 /* Copy the previous result tuple or re-use it if available */
1945 if (Py_REFCNT(result) > 1) {
1946 PyObject *old_result = result;
1947 result = PyTuple_New(npools);
1948 if (result == NULL)
1949 goto empty;
1950 lz->result = result;
1951 for (i=0; i < npools; i++) {
1952 elem = PyTuple_GET_ITEM(old_result, i);
1953 Py_INCREF(elem);
1954 PyTuple_SET_ITEM(result, i, elem);
1955 }
1956 Py_DECREF(old_result);
1957 }
1958 /* Now, we've got the only copy so we can update it in-place */
1959 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001960
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001961 /* Update the pool indices right-to-left. Only advance to the
1962 next pool when the previous one rolls-over */
1963 for (i=npools-1 ; i >= 0 ; i--) {
1964 pool = PyTuple_GET_ITEM(pools, i);
1965 indices[i]++;
1966 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1967 /* Roll-over and advance to next pool */
1968 indices[i] = 0;
1969 elem = PyTuple_GET_ITEM(pool, 0);
1970 Py_INCREF(elem);
1971 oldelem = PyTuple_GET_ITEM(result, i);
1972 PyTuple_SET_ITEM(result, i, elem);
1973 Py_DECREF(oldelem);
1974 } else {
1975 /* No rollover. Just increment and stop here. */
1976 elem = PyTuple_GET_ITEM(pool, indices[i]);
1977 Py_INCREF(elem);
1978 oldelem = PyTuple_GET_ITEM(result, i);
1979 PyTuple_SET_ITEM(result, i, elem);
1980 Py_DECREF(oldelem);
1981 break;
1982 }
1983 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00001984
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001985 /* If i is negative, then the indices have all rolled-over
1986 and we're done. */
1987 if (i < 0)
1988 goto empty;
1989 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001990
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001991 Py_INCREF(result);
1992 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001993
1994empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001995 lz->stopped = 1;
1996 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001997}
1998
1999PyDoc_STRVAR(product_doc,
2000"product(*iterables) --> product object\n\
2001\n\
Raymond Hettinger73d79632008-02-23 02:20:41 +00002002Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002003For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2004The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2005cycle in a manner similar to an odometer (with the rightmost element changing\n\
2006on every iteration).\n\n\
Raymond Hettinger94034ea2009-02-09 18:39:41 +00002007To compute the product of an iterable with itself, specify the number\n\
2008of repetitions with the optional repeat keyword argument. For example,\n\
2009product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002010product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2011product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2012
2013static PyTypeObject product_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002014 PyVarObject_HEAD_INIT(NULL, 0)
2015 "itertools.product", /* tp_name */
2016 sizeof(productobject), /* tp_basicsize */
2017 0, /* tp_itemsize */
2018 /* methods */
2019 (destructor)product_dealloc, /* tp_dealloc */
2020 0, /* tp_print */
2021 0, /* tp_getattr */
2022 0, /* tp_setattr */
2023 0, /* tp_compare */
2024 0, /* tp_repr */
2025 0, /* tp_as_number */
2026 0, /* tp_as_sequence */
2027 0, /* tp_as_mapping */
2028 0, /* tp_hash */
2029 0, /* tp_call */
2030 0, /* tp_str */
2031 PyObject_GenericGetAttr, /* tp_getattro */
2032 0, /* tp_setattro */
2033 0, /* tp_as_buffer */
2034 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2035 Py_TPFLAGS_BASETYPE, /* tp_flags */
2036 product_doc, /* tp_doc */
2037 (traverseproc)product_traverse, /* tp_traverse */
2038 0, /* tp_clear */
2039 0, /* tp_richcompare */
2040 0, /* tp_weaklistoffset */
2041 PyObject_SelfIter, /* tp_iter */
2042 (iternextfunc)product_next, /* tp_iternext */
2043 0, /* tp_methods */
2044 0, /* tp_members */
2045 0, /* tp_getset */
2046 0, /* tp_base */
2047 0, /* tp_dict */
2048 0, /* tp_descr_get */
2049 0, /* tp_descr_set */
2050 0, /* tp_dictoffset */
2051 0, /* tp_init */
2052 0, /* tp_alloc */
2053 product_new, /* tp_new */
2054 PyObject_GC_Del, /* tp_free */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002055};
2056
2057
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002058/* combinations object ************************************************************/
2059
2060typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002061 PyObject_HEAD
2062 PyObject *pool; /* input converted to a tuple */
2063 Py_ssize_t *indices; /* one index per result element */
2064 PyObject *result; /* most recently returned result tuple */
2065 Py_ssize_t r; /* size of result tuple */
2066 int stopped; /* set to 1 when the combinations iterator is exhausted */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002067} combinationsobject;
2068
2069static PyTypeObject combinations_type;
2070
2071static PyObject *
2072combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2073{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002074 combinationsobject *co;
2075 Py_ssize_t n;
2076 Py_ssize_t r;
2077 PyObject *pool = NULL;
2078 PyObject *iterable = NULL;
2079 Py_ssize_t *indices = NULL;
2080 Py_ssize_t i;
2081 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002082
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002083 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2084 &iterable, &r))
2085 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002086
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002087 pool = PySequence_Tuple(iterable);
2088 if (pool == NULL)
2089 goto error;
2090 n = PyTuple_GET_SIZE(pool);
2091 if (r < 0) {
2092 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2093 goto error;
2094 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002095
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002096 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2097 if (indices == NULL) {
2098 PyErr_NoMemory();
2099 goto error;
2100 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002101
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002102 for (i=0 ; i<r ; i++)
2103 indices[i] = i;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002104
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002105 /* create combinationsobject structure */
2106 co = (combinationsobject *)type->tp_alloc(type, 0);
2107 if (co == NULL)
2108 goto error;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002109
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002110 co->pool = pool;
2111 co->indices = indices;
2112 co->result = NULL;
2113 co->r = r;
2114 co->stopped = r > n ? 1 : 0;
2115
2116 return (PyObject *)co;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002117
2118error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002119 if (indices != NULL)
2120 PyMem_Free(indices);
2121 Py_XDECREF(pool);
2122 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002123}
2124
2125static void
2126combinations_dealloc(combinationsobject *co)
2127{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002128 PyObject_GC_UnTrack(co);
2129 Py_XDECREF(co->pool);
2130 Py_XDECREF(co->result);
2131 if (co->indices != NULL)
2132 PyMem_Free(co->indices);
2133 Py_TYPE(co)->tp_free(co);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002134}
2135
2136static int
2137combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2138{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002139 Py_VISIT(co->pool);
2140 Py_VISIT(co->result);
2141 return 0;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002142}
2143
2144static PyObject *
2145combinations_next(combinationsobject *co)
2146{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002147 PyObject *elem;
2148 PyObject *oldelem;
2149 PyObject *pool = co->pool;
2150 Py_ssize_t *indices = co->indices;
2151 PyObject *result = co->result;
2152 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2153 Py_ssize_t r = co->r;
2154 Py_ssize_t i, j, index;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002155
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002156 if (co->stopped)
2157 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002158
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002159 if (result == NULL) {
2160 /* On the first pass, initialize result tuple using the indices */
2161 result = PyTuple_New(r);
2162 if (result == NULL)
2163 goto empty;
2164 co->result = result;
2165 for (i=0; i<r ; i++) {
2166 index = indices[i];
2167 elem = PyTuple_GET_ITEM(pool, index);
2168 Py_INCREF(elem);
2169 PyTuple_SET_ITEM(result, i, elem);
2170 }
2171 } else {
2172 /* Copy the previous result tuple or re-use it if available */
2173 if (Py_REFCNT(result) > 1) {
2174 PyObject *old_result = result;
2175 result = PyTuple_New(r);
2176 if (result == NULL)
2177 goto empty;
2178 co->result = result;
2179 for (i=0; i<r ; i++) {
2180 elem = PyTuple_GET_ITEM(old_result, i);
2181 Py_INCREF(elem);
2182 PyTuple_SET_ITEM(result, i, elem);
2183 }
2184 Py_DECREF(old_result);
2185 }
2186 /* Now, we've got the only copy so we can update it in-place
2187 * CPython's empty tuple is a singleton and cached in
2188 * PyTuple's freelist.
2189 */
2190 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002191
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002192 /* Scan indices right-to-left until finding one that is not
2193 at its maximum (i + n - r). */
2194 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2195 ;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002196
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002197 /* If i is negative, then the indices are all at
2198 their maximum value and we're done. */
2199 if (i < 0)
2200 goto empty;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002201
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002202 /* Increment the current index which we know is not at its
2203 maximum. Then move back to the right setting each index
2204 to its lowest possible value (one higher than the index
2205 to its left -- this maintains the sort order invariant). */
2206 indices[i]++;
2207 for (j=i+1 ; j<r ; j++)
2208 indices[j] = indices[j-1] + 1;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002209
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002210 /* Update the result tuple for the new indices
2211 starting with i, the leftmost index that changed */
2212 for ( ; i<r ; i++) {
2213 index = indices[i];
2214 elem = PyTuple_GET_ITEM(pool, index);
2215 Py_INCREF(elem);
2216 oldelem = PyTuple_GET_ITEM(result, i);
2217 PyTuple_SET_ITEM(result, i, elem);
2218 Py_DECREF(oldelem);
2219 }
2220 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002221
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002222 Py_INCREF(result);
2223 return result;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002224
2225empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002226 co->stopped = 1;
2227 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002228}
2229
2230PyDoc_STRVAR(combinations_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002231"combinations(iterable, r) --> combinations object\n\
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002232\n\
2233Return successive r-length combinations of elements in the iterable.\n\n\
2234combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2235
2236static PyTypeObject combinations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002237 PyVarObject_HEAD_INIT(NULL, 0)
2238 "itertools.combinations", /* tp_name */
2239 sizeof(combinationsobject), /* tp_basicsize */
2240 0, /* tp_itemsize */
2241 /* methods */
2242 (destructor)combinations_dealloc, /* tp_dealloc */
2243 0, /* tp_print */
2244 0, /* tp_getattr */
2245 0, /* tp_setattr */
2246 0, /* tp_compare */
2247 0, /* tp_repr */
2248 0, /* tp_as_number */
2249 0, /* tp_as_sequence */
2250 0, /* tp_as_mapping */
2251 0, /* tp_hash */
2252 0, /* tp_call */
2253 0, /* tp_str */
2254 PyObject_GenericGetAttr, /* tp_getattro */
2255 0, /* tp_setattro */
2256 0, /* tp_as_buffer */
2257 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2258 Py_TPFLAGS_BASETYPE, /* tp_flags */
2259 combinations_doc, /* tp_doc */
2260 (traverseproc)combinations_traverse, /* tp_traverse */
2261 0, /* tp_clear */
2262 0, /* tp_richcompare */
2263 0, /* tp_weaklistoffset */
2264 PyObject_SelfIter, /* tp_iter */
2265 (iternextfunc)combinations_next, /* tp_iternext */
2266 0, /* tp_methods */
2267 0, /* tp_members */
2268 0, /* tp_getset */
2269 0, /* tp_base */
2270 0, /* tp_dict */
2271 0, /* tp_descr_get */
2272 0, /* tp_descr_set */
2273 0, /* tp_dictoffset */
2274 0, /* tp_init */
2275 0, /* tp_alloc */
2276 combinations_new, /* tp_new */
2277 PyObject_GC_Del, /* tp_free */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002278};
2279
2280
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002281/* combinations with replacement object *******************************************/
2282
2283/* Equivalent to:
2284
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002285 def combinations_with_replacement(iterable, r):
2286 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2287 # number items returned: (n+r-1)! / r! / (n-1)!
2288 pool = tuple(iterable)
2289 n = len(pool)
2290 indices = [0] * r
2291 yield tuple(pool[i] for i in indices)
2292 while 1:
2293 for i in reversed(range(r)):
2294 if indices[i] != n - 1:
2295 break
2296 else:
2297 return
2298 indices[i:] = [indices[i] + 1] * (r - i)
2299 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002300
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002301 def combinations_with_replacement2(iterable, r):
2302 'Alternate version that filters from product()'
2303 pool = tuple(iterable)
2304 n = len(pool)
2305 for indices in product(range(n), repeat=r):
2306 if sorted(indices) == list(indices):
2307 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002308*/
2309typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002310 PyObject_HEAD
2311 PyObject *pool; /* input converted to a tuple */
2312 Py_ssize_t *indices; /* one index per result element */
2313 PyObject *result; /* most recently returned result tuple */
2314 Py_ssize_t r; /* size of result tuple */
2315 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002316} cwrobject;
2317
2318static PyTypeObject cwr_type;
2319
2320static PyObject *
2321cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2322{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002323 cwrobject *co;
2324 Py_ssize_t n;
2325 Py_ssize_t r;
2326 PyObject *pool = NULL;
2327 PyObject *iterable = NULL;
2328 Py_ssize_t *indices = NULL;
2329 Py_ssize_t i;
2330 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002331
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002332 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2333 &iterable, &r))
2334 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002335
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002336 pool = PySequence_Tuple(iterable);
2337 if (pool == NULL)
2338 goto error;
2339 n = PyTuple_GET_SIZE(pool);
2340 if (r < 0) {
2341 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2342 goto error;
2343 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002344
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002345 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2346 if (indices == NULL) {
2347 PyErr_NoMemory();
2348 goto error;
2349 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002350
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002351 for (i=0 ; i<r ; i++)
2352 indices[i] = 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002353
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002354 /* create cwrobject structure */
2355 co = (cwrobject *)type->tp_alloc(type, 0);
2356 if (co == NULL)
2357 goto error;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002358
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002359 co->pool = pool;
2360 co->indices = indices;
2361 co->result = NULL;
2362 co->r = r;
2363 co->stopped = !n && r;
2364
2365 return (PyObject *)co;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002366
2367error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002368 if (indices != NULL)
2369 PyMem_Free(indices);
2370 Py_XDECREF(pool);
2371 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002372}
2373
2374static void
2375cwr_dealloc(cwrobject *co)
2376{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002377 PyObject_GC_UnTrack(co);
2378 Py_XDECREF(co->pool);
2379 Py_XDECREF(co->result);
2380 if (co->indices != NULL)
2381 PyMem_Free(co->indices);
2382 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002383}
2384
2385static int
2386cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2387{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002388 Py_VISIT(co->pool);
2389 Py_VISIT(co->result);
2390 return 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002391}
2392
2393static PyObject *
2394cwr_next(cwrobject *co)
2395{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002396 PyObject *elem;
2397 PyObject *oldelem;
2398 PyObject *pool = co->pool;
2399 Py_ssize_t *indices = co->indices;
2400 PyObject *result = co->result;
2401 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2402 Py_ssize_t r = co->r;
2403 Py_ssize_t i, j, index;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002404
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002405 if (co->stopped)
2406 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002407
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002408 if (result == NULL) {
2409 /* On the first pass, initialize result tuple using the indices */
2410 result = PyTuple_New(r);
2411 if (result == NULL)
2412 goto empty;
2413 co->result = result;
2414 for (i=0; i<r ; i++) {
2415 index = indices[i];
2416 elem = PyTuple_GET_ITEM(pool, index);
2417 Py_INCREF(elem);
2418 PyTuple_SET_ITEM(result, i, elem);
2419 }
2420 } else {
2421 /* Copy the previous result tuple or re-use it if available */
2422 if (Py_REFCNT(result) > 1) {
2423 PyObject *old_result = result;
2424 result = PyTuple_New(r);
2425 if (result == NULL)
2426 goto empty;
2427 co->result = result;
2428 for (i=0; i<r ; i++) {
2429 elem = PyTuple_GET_ITEM(old_result, i);
2430 Py_INCREF(elem);
2431 PyTuple_SET_ITEM(result, i, elem);
2432 }
2433 Py_DECREF(old_result);
2434 }
2435 /* Now, we've got the only copy so we can update it in-place CPython's
2436 empty tuple is a singleton and cached in PyTuple's freelist. */
2437 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002438
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002439 /* Scan indices right-to-left until finding one that is not
2440 * at its maximum (n-1). */
2441 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2442 ;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002443
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002444 /* If i is negative, then the indices are all at
2445 their maximum value and we're done. */
2446 if (i < 0)
2447 goto empty;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002448
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002449 /* Increment the current index which we know is not at its
2450 maximum. Then set all to the right to the same value. */
2451 indices[i]++;
2452 for (j=i+1 ; j<r ; j++)
2453 indices[j] = indices[j-1];
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002454
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002455 /* Update the result tuple for the new indices
2456 starting with i, the leftmost index that changed */
2457 for ( ; i<r ; i++) {
2458 index = indices[i];
2459 elem = PyTuple_GET_ITEM(pool, index);
2460 Py_INCREF(elem);
2461 oldelem = PyTuple_GET_ITEM(result, i);
2462 PyTuple_SET_ITEM(result, i, elem);
2463 Py_DECREF(oldelem);
2464 }
2465 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002466
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002467 Py_INCREF(result);
2468 return result;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002469
2470empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002471 co->stopped = 1;
2472 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002473}
2474
2475PyDoc_STRVAR(cwr_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002476"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002477\n\
2478Return successive r-length combinations of elements in the iterable\n\
2479allowing individual elements to have successive repeats.\n\
2480combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2481
2482static PyTypeObject cwr_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002483 PyVarObject_HEAD_INIT(NULL, 0)
2484 "itertools.combinations_with_replacement", /* tp_name */
2485 sizeof(cwrobject), /* tp_basicsize */
2486 0, /* tp_itemsize */
2487 /* methods */
2488 (destructor)cwr_dealloc, /* tp_dealloc */
2489 0, /* tp_print */
2490 0, /* tp_getattr */
2491 0, /* tp_setattr */
2492 0, /* tp_compare */
2493 0, /* tp_repr */
2494 0, /* tp_as_number */
2495 0, /* tp_as_sequence */
2496 0, /* tp_as_mapping */
2497 0, /* tp_hash */
2498 0, /* tp_call */
2499 0, /* tp_str */
2500 PyObject_GenericGetAttr, /* tp_getattro */
2501 0, /* tp_setattro */
2502 0, /* tp_as_buffer */
2503 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2504 Py_TPFLAGS_BASETYPE, /* tp_flags */
2505 cwr_doc, /* tp_doc */
2506 (traverseproc)cwr_traverse, /* tp_traverse */
2507 0, /* tp_clear */
2508 0, /* tp_richcompare */
2509 0, /* tp_weaklistoffset */
2510 PyObject_SelfIter, /* tp_iter */
2511 (iternextfunc)cwr_next, /* tp_iternext */
2512 0, /* tp_methods */
2513 0, /* tp_members */
2514 0, /* tp_getset */
2515 0, /* tp_base */
2516 0, /* tp_dict */
2517 0, /* tp_descr_get */
2518 0, /* tp_descr_set */
2519 0, /* tp_dictoffset */
2520 0, /* tp_init */
2521 0, /* tp_alloc */
2522 cwr_new, /* tp_new */
2523 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002524};
2525
2526
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002527/* permutations object ************************************************************
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002528
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002529def permutations(iterable, r=None):
2530 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2531 pool = tuple(iterable)
2532 n = len(pool)
2533 r = n if r is None else r
2534 indices = range(n)
2535 cycles = range(n-r+1, n+1)[::-1]
2536 yield tuple(pool[i] for i in indices[:r])
2537 while n:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002538 for i in reversed(range(r)):
2539 cycles[i] -= 1
2540 if cycles[i] == 0:
2541 indices[i:] = indices[i+1:] + indices[i:i+1]
2542 cycles[i] = n - i
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002543 else:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002544 j = cycles[i]
2545 indices[i], indices[-j] = indices[-j], indices[i]
2546 yield tuple(pool[i] for i in indices[:r])
2547 break
2548 else:
2549 return
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002550*/
2551
2552typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002553 PyObject_HEAD
2554 PyObject *pool; /* input converted to a tuple */
2555 Py_ssize_t *indices; /* one index per element in the pool */
2556 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2557 PyObject *result; /* most recently returned result tuple */
2558 Py_ssize_t r; /* size of result tuple */
2559 int stopped; /* set to 1 when the permutations iterator is exhausted */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002560} permutationsobject;
2561
2562static PyTypeObject permutations_type;
2563
2564static PyObject *
2565permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2566{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002567 permutationsobject *po;
2568 Py_ssize_t n;
2569 Py_ssize_t r;
2570 PyObject *robj = Py_None;
2571 PyObject *pool = NULL;
2572 PyObject *iterable = NULL;
2573 Py_ssize_t *indices = NULL;
2574 Py_ssize_t *cycles = NULL;
2575 Py_ssize_t i;
2576 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002577
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002578 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2579 &iterable, &robj))
2580 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002581
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002582 pool = PySequence_Tuple(iterable);
2583 if (pool == NULL)
2584 goto error;
2585 n = PyTuple_GET_SIZE(pool);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002586
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002587 r = n;
2588 if (robj != Py_None) {
2589 r = PyInt_AsSsize_t(robj);
2590 if (r == -1 && PyErr_Occurred())
2591 goto error;
2592 }
2593 if (r < 0) {
2594 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2595 goto error;
2596 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002597
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002598 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2599 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2600 if (indices == NULL || cycles == NULL) {
2601 PyErr_NoMemory();
2602 goto error;
2603 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002604
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002605 for (i=0 ; i<n ; i++)
2606 indices[i] = i;
2607 for (i=0 ; i<r ; i++)
2608 cycles[i] = n - i;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002609
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002610 /* create permutationsobject structure */
2611 po = (permutationsobject *)type->tp_alloc(type, 0);
2612 if (po == NULL)
2613 goto error;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002614
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002615 po->pool = pool;
2616 po->indices = indices;
2617 po->cycles = cycles;
2618 po->result = NULL;
2619 po->r = r;
2620 po->stopped = r > n ? 1 : 0;
2621
2622 return (PyObject *)po;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002623
2624error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002625 if (indices != NULL)
2626 PyMem_Free(indices);
2627 if (cycles != NULL)
2628 PyMem_Free(cycles);
2629 Py_XDECREF(pool);
2630 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002631}
2632
2633static void
2634permutations_dealloc(permutationsobject *po)
2635{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002636 PyObject_GC_UnTrack(po);
2637 Py_XDECREF(po->pool);
2638 Py_XDECREF(po->result);
2639 PyMem_Free(po->indices);
2640 PyMem_Free(po->cycles);
2641 Py_TYPE(po)->tp_free(po);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002642}
2643
2644static int
2645permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2646{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002647 Py_VISIT(po->pool);
2648 Py_VISIT(po->result);
2649 return 0;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002650}
2651
2652static PyObject *
2653permutations_next(permutationsobject *po)
2654{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002655 PyObject *elem;
2656 PyObject *oldelem;
2657 PyObject *pool = po->pool;
2658 Py_ssize_t *indices = po->indices;
2659 Py_ssize_t *cycles = po->cycles;
2660 PyObject *result = po->result;
2661 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2662 Py_ssize_t r = po->r;
2663 Py_ssize_t i, j, k, index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002664
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002665 if (po->stopped)
2666 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002667
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002668 if (result == NULL) {
2669 /* On the first pass, initialize result tuple using the indices */
2670 result = PyTuple_New(r);
2671 if (result == NULL)
2672 goto empty;
2673 po->result = result;
2674 for (i=0; i<r ; i++) {
2675 index = indices[i];
2676 elem = PyTuple_GET_ITEM(pool, index);
2677 Py_INCREF(elem);
2678 PyTuple_SET_ITEM(result, i, elem);
2679 }
2680 } else {
2681 if (n == 0)
2682 goto empty;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002683
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002684 /* Copy the previous result tuple or re-use it if available */
2685 if (Py_REFCNT(result) > 1) {
2686 PyObject *old_result = result;
2687 result = PyTuple_New(r);
2688 if (result == NULL)
2689 goto empty;
2690 po->result = result;
2691 for (i=0; i<r ; i++) {
2692 elem = PyTuple_GET_ITEM(old_result, i);
2693 Py_INCREF(elem);
2694 PyTuple_SET_ITEM(result, i, elem);
2695 }
2696 Py_DECREF(old_result);
2697 }
2698 /* Now, we've got the only copy so we can update it in-place */
2699 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002700
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002701 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2702 for (i=r-1 ; i>=0 ; i--) {
2703 cycles[i] -= 1;
2704 if (cycles[i] == 0) {
2705 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2706 index = indices[i];
2707 for (j=i ; j<n-1 ; j++)
2708 indices[j] = indices[j+1];
2709 indices[n-1] = index;
2710 cycles[i] = n - i;
2711 } else {
2712 j = cycles[i];
2713 index = indices[i];
2714 indices[i] = indices[n-j];
2715 indices[n-j] = index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002716
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002717 for (k=i; k<r ; k++) {
2718 /* start with i, the leftmost element that changed */
2719 /* yield tuple(pool[k] for k in indices[:r]) */
2720 index = indices[k];
2721 elem = PyTuple_GET_ITEM(pool, index);
2722 Py_INCREF(elem);
2723 oldelem = PyTuple_GET_ITEM(result, k);
2724 PyTuple_SET_ITEM(result, k, elem);
2725 Py_DECREF(oldelem);
2726 }
2727 break;
2728 }
2729 }
2730 /* If i is negative, then the cycles have all
2731 rolled-over and we're done. */
2732 if (i < 0)
2733 goto empty;
2734 }
2735 Py_INCREF(result);
2736 return result;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002737
2738empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002739 po->stopped = 1;
2740 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002741}
2742
2743PyDoc_STRVAR(permutations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +00002744"permutations(iterable[, r]) --> permutations object\n\
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002745\n\
2746Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl8550d692008-06-10 12:46:39 +00002747permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002748
2749static PyTypeObject permutations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002750 PyVarObject_HEAD_INIT(NULL, 0)
2751 "itertools.permutations", /* tp_name */
2752 sizeof(permutationsobject), /* tp_basicsize */
2753 0, /* tp_itemsize */
2754 /* methods */
2755 (destructor)permutations_dealloc, /* tp_dealloc */
2756 0, /* tp_print */
2757 0, /* tp_getattr */
2758 0, /* tp_setattr */
2759 0, /* tp_compare */
2760 0, /* tp_repr */
2761 0, /* tp_as_number */
2762 0, /* tp_as_sequence */
2763 0, /* tp_as_mapping */
2764 0, /* tp_hash */
2765 0, /* tp_call */
2766 0, /* tp_str */
2767 PyObject_GenericGetAttr, /* tp_getattro */
2768 0, /* tp_setattro */
2769 0, /* tp_as_buffer */
2770 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2771 Py_TPFLAGS_BASETYPE, /* tp_flags */
2772 permutations_doc, /* tp_doc */
2773 (traverseproc)permutations_traverse, /* tp_traverse */
2774 0, /* tp_clear */
2775 0, /* tp_richcompare */
2776 0, /* tp_weaklistoffset */
2777 PyObject_SelfIter, /* tp_iter */
2778 (iternextfunc)permutations_next, /* tp_iternext */
2779 0, /* tp_methods */
2780 0, /* tp_members */
2781 0, /* tp_getset */
2782 0, /* tp_base */
2783 0, /* tp_dict */
2784 0, /* tp_descr_get */
2785 0, /* tp_descr_set */
2786 0, /* tp_dictoffset */
2787 0, /* tp_init */
2788 0, /* tp_alloc */
2789 permutations_new, /* tp_new */
2790 PyObject_GC_Del, /* tp_free */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002791};
2792
2793
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002794/* compress object ************************************************************/
2795
2796/* Equivalent to:
2797
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002798 def compress(data, selectors):
2799 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2800 return (d for d, s in izip(data, selectors) if s)
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002801*/
2802
2803typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002804 PyObject_HEAD
2805 PyObject *data;
2806 PyObject *selectors;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002807} compressobject;
2808
2809static PyTypeObject compress_type;
2810
2811static PyObject *
2812compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2813{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002814 PyObject *seq1, *seq2;
2815 PyObject *data=NULL, *selectors=NULL;
2816 compressobject *lz;
2817 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002818
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002819 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2820 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002821
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002822 data = PyObject_GetIter(seq1);
2823 if (data == NULL)
2824 goto fail;
2825 selectors = PyObject_GetIter(seq2);
2826 if (selectors == NULL)
2827 goto fail;
2828
2829 /* create compressobject structure */
2830 lz = (compressobject *)type->tp_alloc(type, 0);
2831 if (lz == NULL)
2832 goto fail;
2833 lz->data = data;
2834 lz->selectors = selectors;
2835 return (PyObject *)lz;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002836
2837fail:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002838 Py_XDECREF(data);
2839 Py_XDECREF(selectors);
2840 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002841}
2842
2843static void
2844compress_dealloc(compressobject *lz)
2845{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002846 PyObject_GC_UnTrack(lz);
2847 Py_XDECREF(lz->data);
2848 Py_XDECREF(lz->selectors);
2849 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002850}
2851
2852static int
2853compress_traverse(compressobject *lz, visitproc visit, void *arg)
2854{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002855 Py_VISIT(lz->data);
2856 Py_VISIT(lz->selectors);
2857 return 0;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002858}
2859
2860static PyObject *
2861compress_next(compressobject *lz)
2862{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002863 PyObject *data = lz->data, *selectors = lz->selectors;
2864 PyObject *datum, *selector;
2865 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2866 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2867 int ok;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002868
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002869 while (1) {
2870 /* Steps: get datum, get selector, evaluate selector.
2871 Order is important (to match the pure python version
2872 in terms of which input gets a chance to raise an
2873 exception first).
2874 */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002875
Serhiy Storchakabb845652013-04-06 22:04:10 +03002876 datum = datanext(data);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03002877 if (datum == NULL)
Serhiy Storchakabb845652013-04-06 22:04:10 +03002878 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002879
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002880 selector = selectornext(selectors);
2881 if (selector == NULL) {
2882 Py_DECREF(datum);
2883 return NULL;
2884 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002885
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002886 ok = PyObject_IsTrue(selector);
2887 Py_DECREF(selector);
2888 if (ok == 1)
2889 return datum;
2890 Py_DECREF(datum);
2891 if (ok == -1)
2892 return NULL;
2893 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002894}
2895
2896PyDoc_STRVAR(compress_doc,
Raymond Hettinger2e2909f2009-02-19 02:15:14 +00002897"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002898\n\
2899Return data elements corresponding to true selector elements.\n\
2900Forms a shorter iterator from selected data elements using the\n\
2901selectors to choose the data elements.");
2902
2903static PyTypeObject compress_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002904 PyVarObject_HEAD_INIT(NULL, 0)
2905 "itertools.compress", /* tp_name */
2906 sizeof(compressobject), /* tp_basicsize */
2907 0, /* tp_itemsize */
2908 /* methods */
2909 (destructor)compress_dealloc, /* tp_dealloc */
2910 0, /* tp_print */
2911 0, /* tp_getattr */
2912 0, /* tp_setattr */
2913 0, /* tp_compare */
2914 0, /* tp_repr */
2915 0, /* tp_as_number */
2916 0, /* tp_as_sequence */
2917 0, /* tp_as_mapping */
2918 0, /* tp_hash */
2919 0, /* tp_call */
2920 0, /* tp_str */
2921 PyObject_GenericGetAttr, /* tp_getattro */
2922 0, /* tp_setattro */
2923 0, /* tp_as_buffer */
2924 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2925 Py_TPFLAGS_BASETYPE, /* tp_flags */
2926 compress_doc, /* tp_doc */
2927 (traverseproc)compress_traverse, /* tp_traverse */
2928 0, /* tp_clear */
2929 0, /* tp_richcompare */
2930 0, /* tp_weaklistoffset */
2931 PyObject_SelfIter, /* tp_iter */
2932 (iternextfunc)compress_next, /* tp_iternext */
2933 0, /* tp_methods */
2934 0, /* tp_members */
2935 0, /* tp_getset */
2936 0, /* tp_base */
2937 0, /* tp_dict */
2938 0, /* tp_descr_get */
2939 0, /* tp_descr_set */
2940 0, /* tp_dictoffset */
2941 0, /* tp_init */
2942 0, /* tp_alloc */
2943 compress_new, /* tp_new */
2944 PyObject_GC_Del, /* tp_free */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002945};
2946
2947
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002948/* ifilter object ************************************************************/
2949
2950typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002951 PyObject_HEAD
2952 PyObject *func;
2953 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002954} ifilterobject;
2955
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002956static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002957
2958static PyObject *
2959ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2960{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002961 PyObject *func, *seq;
2962 PyObject *it;
2963 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002964
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002965 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
2966 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002967
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002968 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
2969 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002970
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002971 /* Get iterator. */
2972 it = PyObject_GetIter(seq);
2973 if (it == NULL)
2974 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002975
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002976 /* create ifilterobject structure */
2977 lz = (ifilterobject *)type->tp_alloc(type, 0);
2978 if (lz == NULL) {
2979 Py_DECREF(it);
2980 return NULL;
2981 }
2982 Py_INCREF(func);
2983 lz->func = func;
2984 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002985
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002986 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002987}
2988
2989static void
2990ifilter_dealloc(ifilterobject *lz)
2991{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002992 PyObject_GC_UnTrack(lz);
2993 Py_XDECREF(lz->func);
2994 Py_XDECREF(lz->it);
2995 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002996}
2997
2998static int
2999ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
3000{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003001 Py_VISIT(lz->it);
3002 Py_VISIT(lz->func);
3003 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003004}
3005
3006static PyObject *
3007ifilter_next(ifilterobject *lz)
3008{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003009 PyObject *item;
3010 PyObject *it = lz->it;
3011 long ok;
3012 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003013
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003014 iternext = *Py_TYPE(it)->tp_iternext;
3015 for (;;) {
3016 item = iternext(it);
3017 if (item == NULL)
3018 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003019
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003020 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3021 ok = PyObject_IsTrue(item);
3022 } else {
3023 PyObject *good;
3024 good = PyObject_CallFunctionObjArgs(lz->func,
3025 item, NULL);
3026 if (good == NULL) {
3027 Py_DECREF(item);
3028 return NULL;
3029 }
3030 ok = PyObject_IsTrue(good);
3031 Py_DECREF(good);
3032 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003033 if (ok > 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003034 return item;
3035 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003036 if (ok < 0)
3037 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003038 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003039}
3040
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003041PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003042"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003043\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003044Return those items of sequence for which function(item) is true.\n\
3045If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003046
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003047static PyTypeObject ifilter_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003048 PyVarObject_HEAD_INIT(NULL, 0)
3049 "itertools.ifilter", /* tp_name */
3050 sizeof(ifilterobject), /* tp_basicsize */
3051 0, /* tp_itemsize */
3052 /* methods */
3053 (destructor)ifilter_dealloc, /* tp_dealloc */
3054 0, /* tp_print */
3055 0, /* tp_getattr */
3056 0, /* tp_setattr */
3057 0, /* tp_compare */
3058 0, /* tp_repr */
3059 0, /* tp_as_number */
3060 0, /* tp_as_sequence */
3061 0, /* tp_as_mapping */
3062 0, /* tp_hash */
3063 0, /* tp_call */
3064 0, /* tp_str */
3065 PyObject_GenericGetAttr, /* tp_getattro */
3066 0, /* tp_setattro */
3067 0, /* tp_as_buffer */
3068 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3069 Py_TPFLAGS_BASETYPE, /* tp_flags */
3070 ifilter_doc, /* tp_doc */
3071 (traverseproc)ifilter_traverse, /* tp_traverse */
3072 0, /* tp_clear */
3073 0, /* tp_richcompare */
3074 0, /* tp_weaklistoffset */
3075 PyObject_SelfIter, /* tp_iter */
3076 (iternextfunc)ifilter_next, /* tp_iternext */
3077 0, /* tp_methods */
3078 0, /* tp_members */
3079 0, /* tp_getset */
3080 0, /* tp_base */
3081 0, /* tp_dict */
3082 0, /* tp_descr_get */
3083 0, /* tp_descr_set */
3084 0, /* tp_dictoffset */
3085 0, /* tp_init */
3086 0, /* tp_alloc */
3087 ifilter_new, /* tp_new */
3088 PyObject_GC_Del, /* tp_free */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003089};
3090
3091
3092/* ifilterfalse object ************************************************************/
3093
3094typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003095 PyObject_HEAD
3096 PyObject *func;
3097 PyObject *it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003098} ifilterfalseobject;
3099
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003100static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003101
3102static PyObject *
3103ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3104{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003105 PyObject *func, *seq;
3106 PyObject *it;
3107 ifilterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003108
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003109 if (type == &ifilterfalse_type &&
3110 !_PyArg_NoKeywords("ifilterfalse()", kwds))
3111 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003112
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003113 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
3114 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003115
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003116 /* Get iterator. */
3117 it = PyObject_GetIter(seq);
3118 if (it == NULL)
3119 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003120
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003121 /* create ifilterfalseobject structure */
3122 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
3123 if (lz == NULL) {
3124 Py_DECREF(it);
3125 return NULL;
3126 }
3127 Py_INCREF(func);
3128 lz->func = func;
3129 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003130
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003131 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003132}
3133
3134static void
3135ifilterfalse_dealloc(ifilterfalseobject *lz)
3136{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003137 PyObject_GC_UnTrack(lz);
3138 Py_XDECREF(lz->func);
3139 Py_XDECREF(lz->it);
3140 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003141}
3142
3143static int
3144ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
3145{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003146 Py_VISIT(lz->it);
3147 Py_VISIT(lz->func);
3148 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003149}
3150
3151static PyObject *
3152ifilterfalse_next(ifilterfalseobject *lz)
3153{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003154 PyObject *item;
3155 PyObject *it = lz->it;
3156 long ok;
3157 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003158
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003159 iternext = *Py_TYPE(it)->tp_iternext;
3160 for (;;) {
3161 item = iternext(it);
3162 if (item == NULL)
3163 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003164
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003165 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3166 ok = PyObject_IsTrue(item);
3167 } else {
3168 PyObject *good;
3169 good = PyObject_CallFunctionObjArgs(lz->func,
3170 item, NULL);
3171 if (good == NULL) {
3172 Py_DECREF(item);
3173 return NULL;
3174 }
3175 ok = PyObject_IsTrue(good);
3176 Py_DECREF(good);
3177 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003178 if (ok == 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003179 return item;
3180 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003181 if (ok < 0)
3182 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003183 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003184}
3185
Raymond Hettinger60eca932003-02-09 06:40:58 +00003186PyDoc_STRVAR(ifilterfalse_doc,
3187"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
3188\n\
3189Return those items of sequence for which function(item) is false.\n\
3190If function is None, return the items that are false.");
3191
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003192static PyTypeObject ifilterfalse_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003193 PyVarObject_HEAD_INIT(NULL, 0)
3194 "itertools.ifilterfalse", /* tp_name */
3195 sizeof(ifilterfalseobject), /* tp_basicsize */
3196 0, /* tp_itemsize */
3197 /* methods */
3198 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
3199 0, /* tp_print */
3200 0, /* tp_getattr */
3201 0, /* tp_setattr */
3202 0, /* tp_compare */
3203 0, /* tp_repr */
3204 0, /* tp_as_number */
3205 0, /* tp_as_sequence */
3206 0, /* tp_as_mapping */
3207 0, /* tp_hash */
3208 0, /* tp_call */
3209 0, /* tp_str */
3210 PyObject_GenericGetAttr, /* tp_getattro */
3211 0, /* tp_setattro */
3212 0, /* tp_as_buffer */
3213 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3214 Py_TPFLAGS_BASETYPE, /* tp_flags */
3215 ifilterfalse_doc, /* tp_doc */
3216 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
3217 0, /* tp_clear */
3218 0, /* tp_richcompare */
3219 0, /* tp_weaklistoffset */
3220 PyObject_SelfIter, /* tp_iter */
3221 (iternextfunc)ifilterfalse_next, /* tp_iternext */
3222 0, /* tp_methods */
3223 0, /* tp_members */
3224 0, /* tp_getset */
3225 0, /* tp_base */
3226 0, /* tp_dict */
3227 0, /* tp_descr_get */
3228 0, /* tp_descr_set */
3229 0, /* tp_dictoffset */
3230 0, /* tp_init */
3231 0, /* tp_alloc */
3232 ifilterfalse_new, /* tp_new */
3233 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003234};
3235
3236
3237/* count object ************************************************************/
3238
3239typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003240 PyObject_HEAD
3241 Py_ssize_t cnt;
3242 PyObject *long_cnt;
3243 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003244} countobject;
3245
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003246/* Counting logic and invariants:
3247
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003248fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003249
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003250 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3251 Advances with: cnt += 1
3252 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003253
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003254slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003255
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003256 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3257 All counting is done with python objects (no overflows or underflows).
3258 Advances with: long_cnt += long_step
3259 Step may be zero -- effectively a slow version of repeat(cnt).
3260 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003261*/
3262
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003263static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003264
3265static PyObject *
3266count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3267{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003268 countobject *lz;
3269 int slow_mode = 0;
3270 Py_ssize_t cnt = 0;
3271 PyObject *long_cnt = NULL;
3272 PyObject *long_step = NULL;
3273 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003274
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003275 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3276 kwlist, &long_cnt, &long_step))
3277 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003278
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003279 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3280 (long_step != NULL && !PyNumber_Check(long_step))) {
3281 PyErr_SetString(PyExc_TypeError, "a number is required");
3282 return NULL;
3283 }
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003284
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003285 if (long_cnt != NULL) {
3286 cnt = PyInt_AsSsize_t(long_cnt);
3287 if ((cnt == -1 && PyErr_Occurred()) || !PyInt_Check(long_cnt)) {
3288 PyErr_Clear();
3289 slow_mode = 1;
3290 }
3291 Py_INCREF(long_cnt);
3292 } else {
3293 cnt = 0;
3294 long_cnt = PyInt_FromLong(0);
3295 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003296
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003297 /* If not specified, step defaults to 1 */
3298 if (long_step == NULL) {
3299 long_step = PyInt_FromLong(1);
3300 if (long_step == NULL) {
3301 Py_DECREF(long_cnt);
3302 return NULL;
3303 }
3304 } else
3305 Py_INCREF(long_step);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003306
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003307 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003308
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003309 /* Fast mode only works when the step is 1 */
3310 if (!PyInt_Check(long_step) ||
3311 PyInt_AS_LONG(long_step) != 1) {
3312 slow_mode = 1;
3313 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003314
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003315 if (slow_mode)
3316 cnt = PY_SSIZE_T_MAX;
3317 else
3318 Py_CLEAR(long_cnt);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003319
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003320 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3321 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3322 assert(slow_mode ||
3323 (PyInt_Check(long_step) && PyInt_AS_LONG(long_step) == 1));
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003324
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003325 /* create countobject structure */
3326 lz = (countobject *)type->tp_alloc(type, 0);
3327 if (lz == NULL) {
3328 Py_XDECREF(long_cnt);
3329 return NULL;
3330 }
3331 lz->cnt = cnt;
3332 lz->long_cnt = long_cnt;
3333 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003334
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003335 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003336}
3337
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003338static void
3339count_dealloc(countobject *lz)
3340{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003341 PyObject_GC_UnTrack(lz);
3342 Py_XDECREF(lz->long_cnt);
3343 Py_XDECREF(lz->long_step);
3344 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003345}
3346
Benjamin Peterson062a7c32009-02-16 21:07:52 +00003347static int
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003348count_traverse(countobject *lz, visitproc visit, void *arg)
3349{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003350 Py_VISIT(lz->long_cnt);
3351 Py_VISIT(lz->long_step);
3352 return 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003353}
3354
3355static PyObject *
3356count_nextlong(countobject *lz)
3357{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003358 PyObject *long_cnt;
3359 PyObject *stepped_up;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003360
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003361 long_cnt = lz->long_cnt;
3362 if (long_cnt == NULL) {
3363 /* Switch to slow_mode */
3364 long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
3365 if (long_cnt == NULL)
3366 return NULL;
3367 }
3368 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003369
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003370 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3371 if (stepped_up == NULL)
3372 return NULL;
3373 lz->long_cnt = stepped_up;
3374 return long_cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003375}
3376
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003377static PyObject *
3378count_next(countobject *lz)
3379{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003380 if (lz->cnt == PY_SSIZE_T_MAX)
3381 return count_nextlong(lz);
3382 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003383}
3384
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003385static PyObject *
3386count_repr(countobject *lz)
3387{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003388 PyObject *cnt_repr, *step_repr = NULL;
3389 PyObject *result = NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003390
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003391 if (lz->cnt != PY_SSIZE_T_MAX)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003392 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003393
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003394 cnt_repr = PyObject_Repr(lz->long_cnt);
3395 if (cnt_repr == NULL)
3396 return NULL;
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003397
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003398 if (PyInt_Check(lz->long_step) && PyInt_AS_LONG(lz->long_step) == 1) {
3399 /* Don't display step when it is an integer equal to 1 */
3400 result = PyString_FromFormat("count(%s)",
3401 PyString_AS_STRING(cnt_repr));
3402 } else {
3403 step_repr = PyObject_Repr(lz->long_step);
3404 if (step_repr != NULL)
3405 result = PyString_FromFormat("count(%s, %s)",
3406 PyString_AS_STRING(cnt_repr),
3407 PyString_AS_STRING(step_repr));
3408 }
3409 Py_DECREF(cnt_repr);
3410 Py_XDECREF(step_repr);
3411 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003412}
3413
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003414static PyObject *
3415count_reduce(countobject *lz)
3416{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003417 if (lz->cnt == PY_SSIZE_T_MAX)
3418 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3419 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003420}
3421
3422PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3423
3424static PyMethodDef count_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003425 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3426 count_reduce_doc},
3427 {NULL, NULL} /* sentinel */
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003428};
3429
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003430PyDoc_STRVAR(count_doc,
Ezio Melottifdc1e0d2010-06-11 02:21:25 +00003431 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003432\n\
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003433Return a count object whose .next() method returns consecutive values.\n\
3434Equivalent to:\n\n\
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003435 def count(firstval=0, step=1):\n\
Eli Benderskye6802db2013-09-03 06:41:58 -07003436 x = firstval\n\
3437 while 1:\n\
3438 yield x\n\
3439 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003440
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003441static PyTypeObject count_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003442 PyVarObject_HEAD_INIT(NULL, 0)
3443 "itertools.count", /* tp_name */
3444 sizeof(countobject), /* tp_basicsize */
3445 0, /* tp_itemsize */
3446 /* methods */
3447 (destructor)count_dealloc, /* tp_dealloc */
3448 0, /* tp_print */
3449 0, /* tp_getattr */
3450 0, /* tp_setattr */
3451 0, /* tp_compare */
3452 (reprfunc)count_repr, /* tp_repr */
3453 0, /* tp_as_number */
3454 0, /* tp_as_sequence */
3455 0, /* tp_as_mapping */
3456 0, /* tp_hash */
3457 0, /* tp_call */
3458 0, /* tp_str */
3459 PyObject_GenericGetAttr, /* tp_getattro */
3460 0, /* tp_setattro */
3461 0, /* tp_as_buffer */
3462 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3463 Py_TPFLAGS_BASETYPE, /* tp_flags */
3464 count_doc, /* tp_doc */
3465 (traverseproc)count_traverse, /* tp_traverse */
3466 0, /* tp_clear */
3467 0, /* tp_richcompare */
3468 0, /* tp_weaklistoffset */
3469 PyObject_SelfIter, /* tp_iter */
3470 (iternextfunc)count_next, /* tp_iternext */
3471 count_methods, /* tp_methods */
3472 0, /* tp_members */
3473 0, /* tp_getset */
3474 0, /* tp_base */
3475 0, /* tp_dict */
3476 0, /* tp_descr_get */
3477 0, /* tp_descr_set */
3478 0, /* tp_dictoffset */
3479 0, /* tp_init */
3480 0, /* tp_alloc */
3481 count_new, /* tp_new */
3482 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003483};
3484
3485
3486/* izip object ************************************************************/
3487
3488#include "Python.h"
3489
3490typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003491 PyObject_HEAD
3492 Py_ssize_t tuplesize;
3493 PyObject *ittuple; /* tuple of iterators */
3494 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003495} izipobject;
3496
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003497static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003498
3499static PyObject *
3500izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3501{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003502 izipobject *lz;
3503 Py_ssize_t i;
3504 PyObject *ittuple; /* tuple of iterators */
3505 PyObject *result;
3506 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003507
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003508 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
3509 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003510
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003511 /* args must be a tuple */
3512 assert(PyTuple_Check(args));
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003513
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003514 /* obtain iterators */
3515 ittuple = PyTuple_New(tuplesize);
3516 if (ittuple == NULL)
3517 return NULL;
3518 for (i=0; i < tuplesize; ++i) {
3519 PyObject *item = PyTuple_GET_ITEM(args, i);
3520 PyObject *it = PyObject_GetIter(item);
3521 if (it == NULL) {
3522 if (PyErr_ExceptionMatches(PyExc_TypeError))
3523 PyErr_Format(PyExc_TypeError,
3524 "izip argument #%zd must support iteration",
3525 i+1);
3526 Py_DECREF(ittuple);
3527 return NULL;
3528 }
3529 PyTuple_SET_ITEM(ittuple, i, it);
3530 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003531
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003532 /* create a result holder */
3533 result = PyTuple_New(tuplesize);
3534 if (result == NULL) {
3535 Py_DECREF(ittuple);
3536 return NULL;
3537 }
3538 for (i=0 ; i < tuplesize ; i++) {
3539 Py_INCREF(Py_None);
3540 PyTuple_SET_ITEM(result, i, Py_None);
3541 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003542
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003543 /* create izipobject structure */
3544 lz = (izipobject *)type->tp_alloc(type, 0);
3545 if (lz == NULL) {
3546 Py_DECREF(ittuple);
3547 Py_DECREF(result);
3548 return NULL;
3549 }
3550 lz->ittuple = ittuple;
3551 lz->tuplesize = tuplesize;
3552 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003553
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003554 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003555}
3556
3557static void
3558izip_dealloc(izipobject *lz)
3559{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003560 PyObject_GC_UnTrack(lz);
3561 Py_XDECREF(lz->ittuple);
3562 Py_XDECREF(lz->result);
3563 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003564}
3565
3566static int
3567izip_traverse(izipobject *lz, visitproc visit, void *arg)
3568{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003569 Py_VISIT(lz->ittuple);
3570 Py_VISIT(lz->result);
3571 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003572}
3573
3574static PyObject *
3575izip_next(izipobject *lz)
3576{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003577 Py_ssize_t i;
3578 Py_ssize_t tuplesize = lz->tuplesize;
3579 PyObject *result = lz->result;
3580 PyObject *it;
3581 PyObject *item;
3582 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003583
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003584 if (tuplesize == 0)
3585 return NULL;
3586 if (Py_REFCNT(result) == 1) {
3587 Py_INCREF(result);
3588 for (i=0 ; i < tuplesize ; i++) {
3589 it = PyTuple_GET_ITEM(lz->ittuple, i);
3590 item = (*Py_TYPE(it)->tp_iternext)(it);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003591 if (item == NULL) {
3592 Py_DECREF(result);
3593 return NULL;
3594 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003595 olditem = PyTuple_GET_ITEM(result, i);
3596 PyTuple_SET_ITEM(result, i, item);
3597 Py_DECREF(olditem);
3598 }
3599 } else {
3600 result = PyTuple_New(tuplesize);
3601 if (result == NULL)
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003602 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003603 for (i=0 ; i < tuplesize ; i++) {
3604 it = PyTuple_GET_ITEM(lz->ittuple, i);
3605 item = (*Py_TYPE(it)->tp_iternext)(it);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003606 if (item == NULL) {
3607 Py_DECREF(result);
3608 return NULL;
3609 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003610 PyTuple_SET_ITEM(result, i, item);
3611 }
3612 }
3613 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003614}
3615
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003616PyDoc_STRVAR(izip_doc,
3617"izip(iter1 [,iter2 [...]]) --> izip object\n\
3618\n\
3619Return a izip object whose .next() method returns a tuple where\n\
3620the i-th element comes from the i-th iterable argument. The .next()\n\
3621method continues until the shortest iterable in the argument sequence\n\
3622is exhausted and then it raises StopIteration. Works like the zip()\n\
3623function but consumes less memory by returning an iterator instead of\n\
3624a list.");
3625
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003626static PyTypeObject izip_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003627 PyVarObject_HEAD_INIT(NULL, 0)
3628 "itertools.izip", /* tp_name */
3629 sizeof(izipobject), /* tp_basicsize */
3630 0, /* tp_itemsize */
3631 /* methods */
3632 (destructor)izip_dealloc, /* tp_dealloc */
3633 0, /* tp_print */
3634 0, /* tp_getattr */
3635 0, /* tp_setattr */
3636 0, /* tp_compare */
3637 0, /* tp_repr */
3638 0, /* tp_as_number */
3639 0, /* tp_as_sequence */
3640 0, /* tp_as_mapping */
3641 0, /* tp_hash */
3642 0, /* tp_call */
3643 0, /* tp_str */
3644 PyObject_GenericGetAttr, /* tp_getattro */
3645 0, /* tp_setattro */
3646 0, /* tp_as_buffer */
3647 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3648 Py_TPFLAGS_BASETYPE, /* tp_flags */
3649 izip_doc, /* tp_doc */
3650 (traverseproc)izip_traverse, /* tp_traverse */
3651 0, /* tp_clear */
3652 0, /* tp_richcompare */
3653 0, /* tp_weaklistoffset */
3654 PyObject_SelfIter, /* tp_iter */
3655 (iternextfunc)izip_next, /* tp_iternext */
3656 0, /* tp_methods */
3657 0, /* tp_members */
3658 0, /* tp_getset */
3659 0, /* tp_base */
3660 0, /* tp_dict */
3661 0, /* tp_descr_get */
3662 0, /* tp_descr_set */
3663 0, /* tp_dictoffset */
3664 0, /* tp_init */
3665 0, /* tp_alloc */
3666 izip_new, /* tp_new */
3667 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003668};
3669
3670
3671/* repeat object ************************************************************/
3672
3673typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003674 PyObject_HEAD
3675 PyObject *element;
3676 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003677} repeatobject;
3678
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003679static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003680
3681static PyObject *
3682repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3683{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003684 repeatobject *ro;
3685 PyObject *element;
3686 Py_ssize_t cnt = -1;
3687 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003688
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003689 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3690 &element, &cnt))
3691 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003692
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003693 if (PyTuple_Size(args) == 2 && cnt < 0)
3694 cnt = 0;
3695
3696 ro = (repeatobject *)type->tp_alloc(type, 0);
3697 if (ro == NULL)
3698 return NULL;
3699 Py_INCREF(element);
3700 ro->element = element;
3701 ro->cnt = cnt;
3702 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003703}
3704
3705static void
3706repeat_dealloc(repeatobject *ro)
3707{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003708 PyObject_GC_UnTrack(ro);
3709 Py_XDECREF(ro->element);
3710 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003711}
3712
3713static int
3714repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3715{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003716 Py_VISIT(ro->element);
3717 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003718}
3719
3720static PyObject *
3721repeat_next(repeatobject *ro)
3722{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003723 if (ro->cnt == 0)
3724 return NULL;
3725 if (ro->cnt > 0)
3726 ro->cnt--;
3727 Py_INCREF(ro->element);
3728 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003729}
3730
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003731static PyObject *
3732repeat_repr(repeatobject *ro)
3733{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003734 PyObject *result, *objrepr;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003735
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003736 objrepr = PyObject_Repr(ro->element);
3737 if (objrepr == NULL)
3738 return NULL;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003739
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003740 if (ro->cnt == -1)
3741 result = PyString_FromFormat("repeat(%s)",
3742 PyString_AS_STRING(objrepr));
3743 else
3744 result = PyString_FromFormat("repeat(%s, %zd)",
3745 PyString_AS_STRING(objrepr), ro->cnt);
3746 Py_DECREF(objrepr);
3747 return result;
3748}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003749
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003750static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003751repeat_len(repeatobject *ro)
3752{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003753 if (ro->cnt == -1) {
3754 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3755 return NULL;
3756 }
3757 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003758}
3759
Armin Rigof5b3e362006-02-11 21:32:43 +00003760PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003761
3762static PyMethodDef repeat_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003763 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3764 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003765};
3766
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003767PyDoc_STRVAR(repeat_doc,
Raymond Hettinger182edae2009-02-19 02:38:25 +00003768"repeat(object [,times]) -> create an iterator which returns the object\n\
3769for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003770endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003771
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003772static PyTypeObject repeat_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003773 PyVarObject_HEAD_INIT(NULL, 0)
3774 "itertools.repeat", /* tp_name */
3775 sizeof(repeatobject), /* tp_basicsize */
3776 0, /* tp_itemsize */
3777 /* methods */
3778 (destructor)repeat_dealloc, /* tp_dealloc */
3779 0, /* tp_print */
3780 0, /* tp_getattr */
3781 0, /* tp_setattr */
3782 0, /* tp_compare */
3783 (reprfunc)repeat_repr, /* tp_repr */
3784 0, /* tp_as_number */
3785 0, /* tp_as_sequence */
3786 0, /* tp_as_mapping */
3787 0, /* tp_hash */
3788 0, /* tp_call */
3789 0, /* tp_str */
3790 PyObject_GenericGetAttr, /* tp_getattro */
3791 0, /* tp_setattro */
3792 0, /* tp_as_buffer */
3793 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3794 Py_TPFLAGS_BASETYPE, /* tp_flags */
3795 repeat_doc, /* tp_doc */
3796 (traverseproc)repeat_traverse, /* tp_traverse */
3797 0, /* tp_clear */
3798 0, /* tp_richcompare */
3799 0, /* tp_weaklistoffset */
3800 PyObject_SelfIter, /* tp_iter */
3801 (iternextfunc)repeat_next, /* tp_iternext */
3802 repeat_methods, /* tp_methods */
3803 0, /* tp_members */
3804 0, /* tp_getset */
3805 0, /* tp_base */
3806 0, /* tp_dict */
3807 0, /* tp_descr_get */
3808 0, /* tp_descr_set */
3809 0, /* tp_dictoffset */
3810 0, /* tp_init */
3811 0, /* tp_alloc */
3812 repeat_new, /* tp_new */
3813 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003814};
3815
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003816/* iziplongest object ************************************************************/
3817
3818#include "Python.h"
3819
3820typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003821 PyObject_HEAD
3822 Py_ssize_t tuplesize;
3823 Py_ssize_t numactive;
3824 PyObject *ittuple; /* tuple of iterators */
3825 PyObject *result;
3826 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003827} iziplongestobject;
3828
3829static PyTypeObject iziplongest_type;
3830
3831static PyObject *
3832izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3833{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003834 iziplongestobject *lz;
3835 Py_ssize_t i;
3836 PyObject *ittuple; /* tuple of iterators */
3837 PyObject *result;
3838 PyObject *fillvalue = Py_None;
3839 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003840
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003841 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3842 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3843 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3844 PyErr_SetString(PyExc_TypeError,
3845 "izip_longest() got an unexpected keyword argument");
3846 return NULL;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003847 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003848 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003849
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003850 /* args must be a tuple */
3851 assert(PyTuple_Check(args));
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003852
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003853 /* obtain iterators */
3854 ittuple = PyTuple_New(tuplesize);
3855 if (ittuple == NULL)
3856 return NULL;
3857 for (i=0; i < tuplesize; ++i) {
3858 PyObject *item = PyTuple_GET_ITEM(args, i);
3859 PyObject *it = PyObject_GetIter(item);
3860 if (it == NULL) {
3861 if (PyErr_ExceptionMatches(PyExc_TypeError))
3862 PyErr_Format(PyExc_TypeError,
3863 "izip_longest argument #%zd must support iteration",
3864 i+1);
3865 Py_DECREF(ittuple);
3866 return NULL;
3867 }
3868 PyTuple_SET_ITEM(ittuple, i, it);
3869 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003870
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003871 /* create a result holder */
3872 result = PyTuple_New(tuplesize);
3873 if (result == NULL) {
3874 Py_DECREF(ittuple);
3875 return NULL;
3876 }
3877 for (i=0 ; i < tuplesize ; i++) {
3878 Py_INCREF(Py_None);
3879 PyTuple_SET_ITEM(result, i, Py_None);
3880 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003881
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003882 /* create iziplongestobject structure */
3883 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3884 if (lz == NULL) {
3885 Py_DECREF(ittuple);
3886 Py_DECREF(result);
3887 return NULL;
3888 }
3889 lz->ittuple = ittuple;
3890 lz->tuplesize = tuplesize;
3891 lz->numactive = tuplesize;
3892 lz->result = result;
3893 Py_INCREF(fillvalue);
3894 lz->fillvalue = fillvalue;
3895 return (PyObject *)lz;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003896}
3897
3898static void
3899izip_longest_dealloc(iziplongestobject *lz)
3900{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003901 PyObject_GC_UnTrack(lz);
3902 Py_XDECREF(lz->ittuple);
3903 Py_XDECREF(lz->result);
3904 Py_XDECREF(lz->fillvalue);
3905 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003906}
3907
3908static int
3909izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3910{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003911 Py_VISIT(lz->ittuple);
3912 Py_VISIT(lz->result);
3913 Py_VISIT(lz->fillvalue);
3914 return 0;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003915}
3916
3917static PyObject *
3918izip_longest_next(iziplongestobject *lz)
3919{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003920 Py_ssize_t i;
3921 Py_ssize_t tuplesize = lz->tuplesize;
3922 PyObject *result = lz->result;
3923 PyObject *it;
3924 PyObject *item;
3925 PyObject *olditem;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003926
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003927 if (tuplesize == 0)
3928 return NULL;
3929 if (lz->numactive == 0)
3930 return NULL;
3931 if (Py_REFCNT(result) == 1) {
3932 Py_INCREF(result);
3933 for (i=0 ; i < tuplesize ; i++) {
3934 it = PyTuple_GET_ITEM(lz->ittuple, i);
3935 if (it == NULL) {
3936 Py_INCREF(lz->fillvalue);
3937 item = lz->fillvalue;
3938 } else {
3939 item = PyIter_Next(it);
3940 if (item == NULL) {
3941 lz->numactive -= 1;
3942 if (lz->numactive == 0 || PyErr_Occurred()) {
3943 lz->numactive = 0;
3944 Py_DECREF(result);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003945 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003946 } else {
3947 Py_INCREF(lz->fillvalue);
3948 item = lz->fillvalue;
3949 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3950 Py_DECREF(it);
3951 }
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003952 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003953 }
3954 olditem = PyTuple_GET_ITEM(result, i);
3955 PyTuple_SET_ITEM(result, i, item);
3956 Py_DECREF(olditem);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003957 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003958 } else {
3959 result = PyTuple_New(tuplesize);
3960 if (result == NULL)
3961 return NULL;
3962 for (i=0 ; i < tuplesize ; i++) {
3963 it = PyTuple_GET_ITEM(lz->ittuple, i);
3964 if (it == NULL) {
3965 Py_INCREF(lz->fillvalue);
3966 item = lz->fillvalue;
3967 } else {
3968 item = PyIter_Next(it);
3969 if (item == NULL) {
3970 lz->numactive -= 1;
3971 if (lz->numactive == 0 || PyErr_Occurred()) {
3972 lz->numactive = 0;
3973 Py_DECREF(result);
3974 return NULL;
3975 } else {
3976 Py_INCREF(lz->fillvalue);
3977 item = lz->fillvalue;
3978 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3979 Py_DECREF(it);
3980 }
3981 }
3982 }
3983 PyTuple_SET_ITEM(result, i, item);
3984 }
3985 }
3986 return result;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003987}
3988
3989PyDoc_STRVAR(izip_longest_doc,
3990"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3991\n\
3992Return an izip_longest object whose .next() method returns a tuple where\n\
3993the i-th element comes from the i-th iterable argument. The .next()\n\
3994method continues until the longest iterable in the argument sequence\n\
3995is exhausted and then it raises StopIteration. When the shorter iterables\n\
3996are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3997defaults to None or can be specified by a keyword argument.\n\
3998");
3999
4000static PyTypeObject iziplongest_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004001 PyVarObject_HEAD_INIT(NULL, 0)
4002 "itertools.izip_longest", /* tp_name */
4003 sizeof(iziplongestobject), /* tp_basicsize */
4004 0, /* tp_itemsize */
4005 /* methods */
4006 (destructor)izip_longest_dealloc, /* tp_dealloc */
4007 0, /* tp_print */
4008 0, /* tp_getattr */
4009 0, /* tp_setattr */
4010 0, /* tp_compare */
4011 0, /* tp_repr */
4012 0, /* tp_as_number */
4013 0, /* tp_as_sequence */
4014 0, /* tp_as_mapping */
4015 0, /* tp_hash */
4016 0, /* tp_call */
4017 0, /* tp_str */
4018 PyObject_GenericGetAttr, /* tp_getattro */
4019 0, /* tp_setattro */
4020 0, /* tp_as_buffer */
4021 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4022 Py_TPFLAGS_BASETYPE, /* tp_flags */
4023 izip_longest_doc, /* tp_doc */
4024 (traverseproc)izip_longest_traverse, /* tp_traverse */
4025 0, /* tp_clear */
4026 0, /* tp_richcompare */
4027 0, /* tp_weaklistoffset */
4028 PyObject_SelfIter, /* tp_iter */
4029 (iternextfunc)izip_longest_next, /* tp_iternext */
4030 0, /* tp_methods */
4031 0, /* tp_members */
4032 0, /* tp_getset */
4033 0, /* tp_base */
4034 0, /* tp_dict */
4035 0, /* tp_descr_get */
4036 0, /* tp_descr_set */
4037 0, /* tp_dictoffset */
4038 0, /* tp_init */
4039 0, /* tp_alloc */
4040 izip_longest_new, /* tp_new */
4041 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd36862c2007-02-21 05:20:38 +00004042};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004043
4044/* module level code ********************************************************/
4045
4046PyDoc_STRVAR(module_doc,
4047"Functional tools for creating and using iterators.\n\
4048\n\
4049Infinite iterators:\n\
4050count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004051cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004052repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004053\n\
4054Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004055chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettingere76816b2009-01-29 03:43:44 +00004056compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004057dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4058groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00004059ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
4060ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004061islice(seq, [start,] stop [, step]) --> elements from\n\
4062 seq[start:stop:step]\n\
4063imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
4064starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004065tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004066takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004067izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4068izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4069\n\
4070Combinatoric generators:\n\
4071product(p, q, ... [repeat=1]) --> cartesian product\n\
4072permutations(p[, r])\n\
Raymond Hettinger9eac1192009-11-19 01:22:04 +00004073combinations(p, r)\n\
4074combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004075");
4076
4077
Raymond Hettingerad983e72003-11-12 14:32:26 +00004078static PyMethodDef module_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004079 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4080 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004081};
4082
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004083PyMODINIT_FUNC
4084inititertools(void)
4085{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004086 int i;
4087 PyObject *m;
4088 char *name;
4089 PyTypeObject *typelist[] = {
4090 &combinations_type,
4091 &cwr_type,
4092 &cycle_type,
4093 &dropwhile_type,
4094 &takewhile_type,
4095 &islice_type,
4096 &starmap_type,
4097 &imap_type,
4098 &chain_type,
4099 &compress_type,
4100 &ifilter_type,
4101 &ifilterfalse_type,
4102 &count_type,
4103 &izip_type,
4104 &iziplongest_type,
4105 &permutations_type,
4106 &product_type,
4107 &repeat_type,
4108 &groupby_type,
4109 NULL
4110 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004111
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004112 Py_TYPE(&teedataobject_type) = &PyType_Type;
4113 m = Py_InitModule3("itertools", module_methods, module_doc);
4114 if (m == NULL)
4115 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004116
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004117 for (i=0 ; typelist[i] != NULL ; i++) {
4118 if (PyType_Ready(typelist[i]) < 0)
4119 return;
4120 name = strchr(typelist[i]->tp_name, '.');
4121 assert (name != NULL);
4122 Py_INCREF(typelist[i]);
4123 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4124 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004125
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004126 if (PyType_Ready(&teedataobject_type) < 0)
4127 return;
4128 if (PyType_Ready(&tee_type) < 0)
4129 return;
4130 if (PyType_Ready(&_grouper_type) < 0)
4131 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004132}