blob: b63975c1e749fb2257f409a73b2077865a5cc7b8 [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
Benjamin Petersondda91212015-02-01 21:34:07 -05001845 assert(PyTuple_CheckExact(args));
1846 if (repeat == 0) {
1847 nargs = 0;
1848 } else {
1849 nargs = PyTuple_GET_SIZE(args);
1850 if (repeat > PY_SSIZE_T_MAX/sizeof(Py_ssize_t) ||
1851 nargs > PY_SSIZE_T_MAX/(repeat * sizeof(Py_ssize_t))) {
1852 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
1853 return NULL;
1854 }
1855 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001856 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001857
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001858 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1859 if (indices == NULL) {
1860 PyErr_NoMemory();
1861 goto error;
1862 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001863
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001864 pools = PyTuple_New(npools);
1865 if (pools == NULL)
1866 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001867
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001868 for (i=0; i < nargs ; ++i) {
1869 PyObject *item = PyTuple_GET_ITEM(args, i);
1870 PyObject *pool = PySequence_Tuple(item);
1871 if (pool == NULL)
1872 goto error;
1873 PyTuple_SET_ITEM(pools, i, pool);
1874 indices[i] = 0;
1875 }
1876 for ( ; i < npools; ++i) {
1877 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1878 Py_INCREF(pool);
1879 PyTuple_SET_ITEM(pools, i, pool);
1880 indices[i] = 0;
1881 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001882
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001883 /* create productobject structure */
1884 lz = (productobject *)type->tp_alloc(type, 0);
1885 if (lz == NULL)
1886 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001887
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001888 lz->pools = pools;
1889 lz->indices = indices;
1890 lz->result = NULL;
1891 lz->stopped = 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001892
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001893 return (PyObject *)lz;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001894
1895error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001896 if (indices != NULL)
1897 PyMem_Free(indices);
1898 Py_XDECREF(pools);
1899 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001900}
1901
1902static void
1903product_dealloc(productobject *lz)
1904{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001905 PyObject_GC_UnTrack(lz);
1906 Py_XDECREF(lz->pools);
1907 Py_XDECREF(lz->result);
1908 if (lz->indices != NULL)
1909 PyMem_Free(lz->indices);
1910 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001911}
1912
1913static int
1914product_traverse(productobject *lz, visitproc visit, void *arg)
1915{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001916 Py_VISIT(lz->pools);
1917 Py_VISIT(lz->result);
1918 return 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001919}
1920
1921static PyObject *
1922product_next(productobject *lz)
1923{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001924 PyObject *pool;
1925 PyObject *elem;
1926 PyObject *oldelem;
1927 PyObject *pools = lz->pools;
1928 PyObject *result = lz->result;
1929 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1930 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001931
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001932 if (lz->stopped)
1933 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001934
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001935 if (result == NULL) {
1936 /* On the first pass, return an initial tuple filled with the
1937 first element from each pool. */
1938 result = PyTuple_New(npools);
1939 if (result == NULL)
1940 goto empty;
1941 lz->result = result;
1942 for (i=0; i < npools; i++) {
1943 pool = PyTuple_GET_ITEM(pools, i);
1944 if (PyTuple_GET_SIZE(pool) == 0)
1945 goto empty;
1946 elem = PyTuple_GET_ITEM(pool, 0);
1947 Py_INCREF(elem);
1948 PyTuple_SET_ITEM(result, i, elem);
1949 }
1950 } else {
1951 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001952
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001953 /* Copy the previous result tuple or re-use it if available */
1954 if (Py_REFCNT(result) > 1) {
1955 PyObject *old_result = result;
1956 result = PyTuple_New(npools);
1957 if (result == NULL)
1958 goto empty;
1959 lz->result = result;
1960 for (i=0; i < npools; i++) {
1961 elem = PyTuple_GET_ITEM(old_result, i);
1962 Py_INCREF(elem);
1963 PyTuple_SET_ITEM(result, i, elem);
1964 }
1965 Py_DECREF(old_result);
1966 }
1967 /* Now, we've got the only copy so we can update it in-place */
1968 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001969
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001970 /* Update the pool indices right-to-left. Only advance to the
1971 next pool when the previous one rolls-over */
1972 for (i=npools-1 ; i >= 0 ; i--) {
1973 pool = PyTuple_GET_ITEM(pools, i);
1974 indices[i]++;
1975 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1976 /* Roll-over and advance to next pool */
1977 indices[i] = 0;
1978 elem = PyTuple_GET_ITEM(pool, 0);
1979 Py_INCREF(elem);
1980 oldelem = PyTuple_GET_ITEM(result, i);
1981 PyTuple_SET_ITEM(result, i, elem);
1982 Py_DECREF(oldelem);
1983 } else {
1984 /* No rollover. Just increment and stop here. */
1985 elem = PyTuple_GET_ITEM(pool, indices[i]);
1986 Py_INCREF(elem);
1987 oldelem = PyTuple_GET_ITEM(result, i);
1988 PyTuple_SET_ITEM(result, i, elem);
1989 Py_DECREF(oldelem);
1990 break;
1991 }
1992 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00001993
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001994 /* If i is negative, then the indices have all rolled-over
1995 and we're done. */
1996 if (i < 0)
1997 goto empty;
1998 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001999
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002000 Py_INCREF(result);
2001 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002002
2003empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002004 lz->stopped = 1;
2005 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002006}
2007
2008PyDoc_STRVAR(product_doc,
2009"product(*iterables) --> product object\n\
2010\n\
Raymond Hettinger73d79632008-02-23 02:20:41 +00002011Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002012For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2013The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2014cycle in a manner similar to an odometer (with the rightmost element changing\n\
2015on every iteration).\n\n\
Raymond Hettinger94034ea2009-02-09 18:39:41 +00002016To compute the product of an iterable with itself, specify the number\n\
2017of repetitions with the optional repeat keyword argument. For example,\n\
2018product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002019product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2020product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2021
2022static PyTypeObject product_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002023 PyVarObject_HEAD_INIT(NULL, 0)
2024 "itertools.product", /* tp_name */
2025 sizeof(productobject), /* tp_basicsize */
2026 0, /* tp_itemsize */
2027 /* methods */
2028 (destructor)product_dealloc, /* tp_dealloc */
2029 0, /* tp_print */
2030 0, /* tp_getattr */
2031 0, /* tp_setattr */
2032 0, /* tp_compare */
2033 0, /* tp_repr */
2034 0, /* tp_as_number */
2035 0, /* tp_as_sequence */
2036 0, /* tp_as_mapping */
2037 0, /* tp_hash */
2038 0, /* tp_call */
2039 0, /* tp_str */
2040 PyObject_GenericGetAttr, /* tp_getattro */
2041 0, /* tp_setattro */
2042 0, /* tp_as_buffer */
2043 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2044 Py_TPFLAGS_BASETYPE, /* tp_flags */
2045 product_doc, /* tp_doc */
2046 (traverseproc)product_traverse, /* tp_traverse */
2047 0, /* tp_clear */
2048 0, /* tp_richcompare */
2049 0, /* tp_weaklistoffset */
2050 PyObject_SelfIter, /* tp_iter */
2051 (iternextfunc)product_next, /* tp_iternext */
2052 0, /* tp_methods */
2053 0, /* tp_members */
2054 0, /* tp_getset */
2055 0, /* tp_base */
2056 0, /* tp_dict */
2057 0, /* tp_descr_get */
2058 0, /* tp_descr_set */
2059 0, /* tp_dictoffset */
2060 0, /* tp_init */
2061 0, /* tp_alloc */
2062 product_new, /* tp_new */
2063 PyObject_GC_Del, /* tp_free */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002064};
2065
2066
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002067/* combinations object ************************************************************/
2068
2069typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002070 PyObject_HEAD
2071 PyObject *pool; /* input converted to a tuple */
2072 Py_ssize_t *indices; /* one index per result element */
2073 PyObject *result; /* most recently returned result tuple */
2074 Py_ssize_t r; /* size of result tuple */
2075 int stopped; /* set to 1 when the combinations iterator is exhausted */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002076} combinationsobject;
2077
2078static PyTypeObject combinations_type;
2079
2080static PyObject *
2081combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2082{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002083 combinationsobject *co;
2084 Py_ssize_t n;
2085 Py_ssize_t r;
2086 PyObject *pool = NULL;
2087 PyObject *iterable = NULL;
2088 Py_ssize_t *indices = NULL;
2089 Py_ssize_t i;
2090 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002091
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002092 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2093 &iterable, &r))
2094 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002095
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002096 pool = PySequence_Tuple(iterable);
2097 if (pool == NULL)
2098 goto error;
2099 n = PyTuple_GET_SIZE(pool);
2100 if (r < 0) {
2101 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2102 goto error;
2103 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002104
Benjamin Peterson021dec12015-02-01 20:59:00 -05002105 if (r > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)) {
2106 PyErr_SetString(PyExc_OverflowError, "r is too big");
2107 goto error;
2108 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002109 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2110 if (indices == NULL) {
2111 PyErr_NoMemory();
2112 goto error;
2113 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002114
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002115 for (i=0 ; i<r ; i++)
2116 indices[i] = i;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002117
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002118 /* create combinationsobject structure */
2119 co = (combinationsobject *)type->tp_alloc(type, 0);
2120 if (co == NULL)
2121 goto error;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002122
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002123 co->pool = pool;
2124 co->indices = indices;
2125 co->result = NULL;
2126 co->r = r;
2127 co->stopped = r > n ? 1 : 0;
2128
2129 return (PyObject *)co;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002130
2131error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002132 if (indices != NULL)
2133 PyMem_Free(indices);
2134 Py_XDECREF(pool);
2135 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002136}
2137
2138static void
2139combinations_dealloc(combinationsobject *co)
2140{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002141 PyObject_GC_UnTrack(co);
2142 Py_XDECREF(co->pool);
2143 Py_XDECREF(co->result);
2144 if (co->indices != NULL)
2145 PyMem_Free(co->indices);
2146 Py_TYPE(co)->tp_free(co);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002147}
2148
2149static int
2150combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2151{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002152 Py_VISIT(co->pool);
2153 Py_VISIT(co->result);
2154 return 0;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002155}
2156
2157static PyObject *
2158combinations_next(combinationsobject *co)
2159{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002160 PyObject *elem;
2161 PyObject *oldelem;
2162 PyObject *pool = co->pool;
2163 Py_ssize_t *indices = co->indices;
2164 PyObject *result = co->result;
2165 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2166 Py_ssize_t r = co->r;
2167 Py_ssize_t i, j, index;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002168
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002169 if (co->stopped)
2170 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002171
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002172 if (result == NULL) {
2173 /* On the first pass, initialize result tuple using the indices */
2174 result = PyTuple_New(r);
2175 if (result == NULL)
2176 goto empty;
2177 co->result = result;
2178 for (i=0; i<r ; i++) {
2179 index = indices[i];
2180 elem = PyTuple_GET_ITEM(pool, index);
2181 Py_INCREF(elem);
2182 PyTuple_SET_ITEM(result, i, elem);
2183 }
2184 } else {
2185 /* Copy the previous result tuple or re-use it if available */
2186 if (Py_REFCNT(result) > 1) {
2187 PyObject *old_result = result;
2188 result = PyTuple_New(r);
2189 if (result == NULL)
2190 goto empty;
2191 co->result = result;
2192 for (i=0; i<r ; i++) {
2193 elem = PyTuple_GET_ITEM(old_result, i);
2194 Py_INCREF(elem);
2195 PyTuple_SET_ITEM(result, i, elem);
2196 }
2197 Py_DECREF(old_result);
2198 }
2199 /* Now, we've got the only copy so we can update it in-place
2200 * CPython's empty tuple is a singleton and cached in
2201 * PyTuple's freelist.
2202 */
2203 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002204
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002205 /* Scan indices right-to-left until finding one that is not
2206 at its maximum (i + n - r). */
2207 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2208 ;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002209
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002210 /* If i is negative, then the indices are all at
2211 their maximum value and we're done. */
2212 if (i < 0)
2213 goto empty;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002214
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002215 /* Increment the current index which we know is not at its
2216 maximum. Then move back to the right setting each index
2217 to its lowest possible value (one higher than the index
2218 to its left -- this maintains the sort order invariant). */
2219 indices[i]++;
2220 for (j=i+1 ; j<r ; j++)
2221 indices[j] = indices[j-1] + 1;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002222
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002223 /* Update the result tuple for the new indices
2224 starting with i, the leftmost index that changed */
2225 for ( ; i<r ; i++) {
2226 index = indices[i];
2227 elem = PyTuple_GET_ITEM(pool, index);
2228 Py_INCREF(elem);
2229 oldelem = PyTuple_GET_ITEM(result, i);
2230 PyTuple_SET_ITEM(result, i, elem);
2231 Py_DECREF(oldelem);
2232 }
2233 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002234
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002235 Py_INCREF(result);
2236 return result;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002237
2238empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002239 co->stopped = 1;
2240 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002241}
2242
2243PyDoc_STRVAR(combinations_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002244"combinations(iterable, r) --> combinations object\n\
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002245\n\
2246Return successive r-length combinations of elements in the iterable.\n\n\
2247combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2248
2249static PyTypeObject combinations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002250 PyVarObject_HEAD_INIT(NULL, 0)
2251 "itertools.combinations", /* tp_name */
2252 sizeof(combinationsobject), /* tp_basicsize */
2253 0, /* tp_itemsize */
2254 /* methods */
2255 (destructor)combinations_dealloc, /* tp_dealloc */
2256 0, /* tp_print */
2257 0, /* tp_getattr */
2258 0, /* tp_setattr */
2259 0, /* tp_compare */
2260 0, /* tp_repr */
2261 0, /* tp_as_number */
2262 0, /* tp_as_sequence */
2263 0, /* tp_as_mapping */
2264 0, /* tp_hash */
2265 0, /* tp_call */
2266 0, /* tp_str */
2267 PyObject_GenericGetAttr, /* tp_getattro */
2268 0, /* tp_setattro */
2269 0, /* tp_as_buffer */
2270 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2271 Py_TPFLAGS_BASETYPE, /* tp_flags */
2272 combinations_doc, /* tp_doc */
2273 (traverseproc)combinations_traverse, /* tp_traverse */
2274 0, /* tp_clear */
2275 0, /* tp_richcompare */
2276 0, /* tp_weaklistoffset */
2277 PyObject_SelfIter, /* tp_iter */
2278 (iternextfunc)combinations_next, /* tp_iternext */
2279 0, /* tp_methods */
2280 0, /* tp_members */
2281 0, /* tp_getset */
2282 0, /* tp_base */
2283 0, /* tp_dict */
2284 0, /* tp_descr_get */
2285 0, /* tp_descr_set */
2286 0, /* tp_dictoffset */
2287 0, /* tp_init */
2288 0, /* tp_alloc */
2289 combinations_new, /* tp_new */
2290 PyObject_GC_Del, /* tp_free */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002291};
2292
2293
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002294/* combinations with replacement object *******************************************/
2295
2296/* Equivalent to:
2297
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002298 def combinations_with_replacement(iterable, r):
2299 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2300 # number items returned: (n+r-1)! / r! / (n-1)!
2301 pool = tuple(iterable)
2302 n = len(pool)
2303 indices = [0] * r
2304 yield tuple(pool[i] for i in indices)
2305 while 1:
2306 for i in reversed(range(r)):
2307 if indices[i] != n - 1:
2308 break
2309 else:
2310 return
2311 indices[i:] = [indices[i] + 1] * (r - i)
2312 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002313
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002314 def combinations_with_replacement2(iterable, r):
2315 'Alternate version that filters from product()'
2316 pool = tuple(iterable)
2317 n = len(pool)
2318 for indices in product(range(n), repeat=r):
2319 if sorted(indices) == list(indices):
2320 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002321*/
2322typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002323 PyObject_HEAD
2324 PyObject *pool; /* input converted to a tuple */
2325 Py_ssize_t *indices; /* one index per result element */
2326 PyObject *result; /* most recently returned result tuple */
2327 Py_ssize_t r; /* size of result tuple */
2328 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002329} cwrobject;
2330
2331static PyTypeObject cwr_type;
2332
2333static PyObject *
2334cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2335{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002336 cwrobject *co;
2337 Py_ssize_t n;
2338 Py_ssize_t r;
2339 PyObject *pool = NULL;
2340 PyObject *iterable = NULL;
2341 Py_ssize_t *indices = NULL;
2342 Py_ssize_t i;
2343 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002344
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002345 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2346 &iterable, &r))
2347 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002348
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002349 pool = PySequence_Tuple(iterable);
2350 if (pool == NULL)
2351 goto error;
2352 n = PyTuple_GET_SIZE(pool);
2353 if (r < 0) {
2354 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2355 goto error;
2356 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002357
Benjamin Peterson17845c12015-02-01 21:10:47 -05002358 if (r > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)) {
2359 PyErr_SetString(PyExc_OverflowError, "r is too big");
2360 goto error;
2361 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002362 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2363 if (indices == NULL) {
2364 PyErr_NoMemory();
2365 goto error;
2366 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002367
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002368 for (i=0 ; i<r ; i++)
2369 indices[i] = 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002370
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002371 /* create cwrobject structure */
2372 co = (cwrobject *)type->tp_alloc(type, 0);
2373 if (co == NULL)
2374 goto error;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002375
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002376 co->pool = pool;
2377 co->indices = indices;
2378 co->result = NULL;
2379 co->r = r;
2380 co->stopped = !n && r;
2381
2382 return (PyObject *)co;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002383
2384error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002385 if (indices != NULL)
2386 PyMem_Free(indices);
2387 Py_XDECREF(pool);
2388 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002389}
2390
2391static void
2392cwr_dealloc(cwrobject *co)
2393{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002394 PyObject_GC_UnTrack(co);
2395 Py_XDECREF(co->pool);
2396 Py_XDECREF(co->result);
2397 if (co->indices != NULL)
2398 PyMem_Free(co->indices);
2399 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002400}
2401
2402static int
2403cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2404{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002405 Py_VISIT(co->pool);
2406 Py_VISIT(co->result);
2407 return 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002408}
2409
2410static PyObject *
2411cwr_next(cwrobject *co)
2412{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002413 PyObject *elem;
2414 PyObject *oldelem;
2415 PyObject *pool = co->pool;
2416 Py_ssize_t *indices = co->indices;
2417 PyObject *result = co->result;
2418 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2419 Py_ssize_t r = co->r;
2420 Py_ssize_t i, j, index;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002421
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002422 if (co->stopped)
2423 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002424
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002425 if (result == NULL) {
2426 /* On the first pass, initialize result tuple using the indices */
2427 result = PyTuple_New(r);
2428 if (result == NULL)
2429 goto empty;
2430 co->result = result;
2431 for (i=0; i<r ; i++) {
2432 index = indices[i];
2433 elem = PyTuple_GET_ITEM(pool, index);
2434 Py_INCREF(elem);
2435 PyTuple_SET_ITEM(result, i, elem);
2436 }
2437 } else {
2438 /* Copy the previous result tuple or re-use it if available */
2439 if (Py_REFCNT(result) > 1) {
2440 PyObject *old_result = result;
2441 result = PyTuple_New(r);
2442 if (result == NULL)
2443 goto empty;
2444 co->result = result;
2445 for (i=0; i<r ; i++) {
2446 elem = PyTuple_GET_ITEM(old_result, i);
2447 Py_INCREF(elem);
2448 PyTuple_SET_ITEM(result, i, elem);
2449 }
2450 Py_DECREF(old_result);
2451 }
2452 /* Now, we've got the only copy so we can update it in-place CPython's
2453 empty tuple is a singleton and cached in PyTuple's freelist. */
2454 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002455
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002456 /* Scan indices right-to-left until finding one that is not
2457 * at its maximum (n-1). */
2458 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2459 ;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002460
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002461 /* If i is negative, then the indices are all at
2462 their maximum value and we're done. */
2463 if (i < 0)
2464 goto empty;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002465
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002466 /* Increment the current index which we know is not at its
2467 maximum. Then set all to the right to the same value. */
2468 indices[i]++;
2469 for (j=i+1 ; j<r ; j++)
2470 indices[j] = indices[j-1];
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002471
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002472 /* Update the result tuple for the new indices
2473 starting with i, the leftmost index that changed */
2474 for ( ; i<r ; i++) {
2475 index = indices[i];
2476 elem = PyTuple_GET_ITEM(pool, index);
2477 Py_INCREF(elem);
2478 oldelem = PyTuple_GET_ITEM(result, i);
2479 PyTuple_SET_ITEM(result, i, elem);
2480 Py_DECREF(oldelem);
2481 }
2482 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002483
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002484 Py_INCREF(result);
2485 return result;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002486
2487empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002488 co->stopped = 1;
2489 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002490}
2491
2492PyDoc_STRVAR(cwr_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002493"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002494\n\
2495Return successive r-length combinations of elements in the iterable\n\
2496allowing individual elements to have successive repeats.\n\
2497combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2498
2499static PyTypeObject cwr_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002500 PyVarObject_HEAD_INIT(NULL, 0)
2501 "itertools.combinations_with_replacement", /* tp_name */
2502 sizeof(cwrobject), /* tp_basicsize */
2503 0, /* tp_itemsize */
2504 /* methods */
2505 (destructor)cwr_dealloc, /* tp_dealloc */
2506 0, /* tp_print */
2507 0, /* tp_getattr */
2508 0, /* tp_setattr */
2509 0, /* tp_compare */
2510 0, /* tp_repr */
2511 0, /* tp_as_number */
2512 0, /* tp_as_sequence */
2513 0, /* tp_as_mapping */
2514 0, /* tp_hash */
2515 0, /* tp_call */
2516 0, /* tp_str */
2517 PyObject_GenericGetAttr, /* tp_getattro */
2518 0, /* tp_setattro */
2519 0, /* tp_as_buffer */
2520 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2521 Py_TPFLAGS_BASETYPE, /* tp_flags */
2522 cwr_doc, /* tp_doc */
2523 (traverseproc)cwr_traverse, /* tp_traverse */
2524 0, /* tp_clear */
2525 0, /* tp_richcompare */
2526 0, /* tp_weaklistoffset */
2527 PyObject_SelfIter, /* tp_iter */
2528 (iternextfunc)cwr_next, /* tp_iternext */
2529 0, /* tp_methods */
2530 0, /* tp_members */
2531 0, /* tp_getset */
2532 0, /* tp_base */
2533 0, /* tp_dict */
2534 0, /* tp_descr_get */
2535 0, /* tp_descr_set */
2536 0, /* tp_dictoffset */
2537 0, /* tp_init */
2538 0, /* tp_alloc */
2539 cwr_new, /* tp_new */
2540 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002541};
2542
2543
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002544/* permutations object ************************************************************
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002545
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002546def permutations(iterable, r=None):
2547 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2548 pool = tuple(iterable)
2549 n = len(pool)
2550 r = n if r is None else r
2551 indices = range(n)
2552 cycles = range(n-r+1, n+1)[::-1]
2553 yield tuple(pool[i] for i in indices[:r])
2554 while n:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002555 for i in reversed(range(r)):
2556 cycles[i] -= 1
2557 if cycles[i] == 0:
2558 indices[i:] = indices[i+1:] + indices[i:i+1]
2559 cycles[i] = n - i
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002560 else:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002561 j = cycles[i]
2562 indices[i], indices[-j] = indices[-j], indices[i]
2563 yield tuple(pool[i] for i in indices[:r])
2564 break
2565 else:
2566 return
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002567*/
2568
2569typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002570 PyObject_HEAD
2571 PyObject *pool; /* input converted to a tuple */
2572 Py_ssize_t *indices; /* one index per element in the pool */
2573 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2574 PyObject *result; /* most recently returned result tuple */
2575 Py_ssize_t r; /* size of result tuple */
2576 int stopped; /* set to 1 when the permutations iterator is exhausted */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002577} permutationsobject;
2578
2579static PyTypeObject permutations_type;
2580
2581static PyObject *
2582permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2583{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002584 permutationsobject *po;
2585 Py_ssize_t n;
2586 Py_ssize_t r;
2587 PyObject *robj = Py_None;
2588 PyObject *pool = NULL;
2589 PyObject *iterable = NULL;
2590 Py_ssize_t *indices = NULL;
2591 Py_ssize_t *cycles = NULL;
2592 Py_ssize_t i;
2593 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002594
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002595 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2596 &iterable, &robj))
2597 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002598
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002599 pool = PySequence_Tuple(iterable);
2600 if (pool == NULL)
2601 goto error;
2602 n = PyTuple_GET_SIZE(pool);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002603
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002604 r = n;
2605 if (robj != Py_None) {
2606 r = PyInt_AsSsize_t(robj);
2607 if (r == -1 && PyErr_Occurred())
2608 goto error;
2609 }
2610 if (r < 0) {
2611 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2612 goto error;
2613 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002614
Benjamin Petersondda91212015-02-01 21:34:07 -05002615 if (n > PY_SSIZE_T_MAX/sizeof(Py_ssize_t) ||
2616 r > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)) {
2617 PyErr_SetString(PyExc_OverflowError, "parameters too large");
2618 goto error;
2619 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002620 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2621 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2622 if (indices == NULL || cycles == NULL) {
2623 PyErr_NoMemory();
2624 goto error;
2625 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002626
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002627 for (i=0 ; i<n ; i++)
2628 indices[i] = i;
2629 for (i=0 ; i<r ; i++)
2630 cycles[i] = n - i;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002631
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002632 /* create permutationsobject structure */
2633 po = (permutationsobject *)type->tp_alloc(type, 0);
2634 if (po == NULL)
2635 goto error;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002636
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002637 po->pool = pool;
2638 po->indices = indices;
2639 po->cycles = cycles;
2640 po->result = NULL;
2641 po->r = r;
2642 po->stopped = r > n ? 1 : 0;
2643
2644 return (PyObject *)po;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002645
2646error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002647 if (indices != NULL)
2648 PyMem_Free(indices);
2649 if (cycles != NULL)
2650 PyMem_Free(cycles);
2651 Py_XDECREF(pool);
2652 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002653}
2654
2655static void
2656permutations_dealloc(permutationsobject *po)
2657{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002658 PyObject_GC_UnTrack(po);
2659 Py_XDECREF(po->pool);
2660 Py_XDECREF(po->result);
2661 PyMem_Free(po->indices);
2662 PyMem_Free(po->cycles);
2663 Py_TYPE(po)->tp_free(po);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002664}
2665
2666static int
2667permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2668{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002669 Py_VISIT(po->pool);
2670 Py_VISIT(po->result);
2671 return 0;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002672}
2673
2674static PyObject *
2675permutations_next(permutationsobject *po)
2676{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002677 PyObject *elem;
2678 PyObject *oldelem;
2679 PyObject *pool = po->pool;
2680 Py_ssize_t *indices = po->indices;
2681 Py_ssize_t *cycles = po->cycles;
2682 PyObject *result = po->result;
2683 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2684 Py_ssize_t r = po->r;
2685 Py_ssize_t i, j, k, index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002686
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002687 if (po->stopped)
2688 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002689
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002690 if (result == NULL) {
2691 /* On the first pass, initialize result tuple using the indices */
2692 result = PyTuple_New(r);
2693 if (result == NULL)
2694 goto empty;
2695 po->result = result;
2696 for (i=0; i<r ; i++) {
2697 index = indices[i];
2698 elem = PyTuple_GET_ITEM(pool, index);
2699 Py_INCREF(elem);
2700 PyTuple_SET_ITEM(result, i, elem);
2701 }
2702 } else {
2703 if (n == 0)
2704 goto empty;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002705
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002706 /* Copy the previous result tuple or re-use it if available */
2707 if (Py_REFCNT(result) > 1) {
2708 PyObject *old_result = result;
2709 result = PyTuple_New(r);
2710 if (result == NULL)
2711 goto empty;
2712 po->result = result;
2713 for (i=0; i<r ; i++) {
2714 elem = PyTuple_GET_ITEM(old_result, i);
2715 Py_INCREF(elem);
2716 PyTuple_SET_ITEM(result, i, elem);
2717 }
2718 Py_DECREF(old_result);
2719 }
2720 /* Now, we've got the only copy so we can update it in-place */
2721 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002722
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002723 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2724 for (i=r-1 ; i>=0 ; i--) {
2725 cycles[i] -= 1;
2726 if (cycles[i] == 0) {
2727 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2728 index = indices[i];
2729 for (j=i ; j<n-1 ; j++)
2730 indices[j] = indices[j+1];
2731 indices[n-1] = index;
2732 cycles[i] = n - i;
2733 } else {
2734 j = cycles[i];
2735 index = indices[i];
2736 indices[i] = indices[n-j];
2737 indices[n-j] = index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002738
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002739 for (k=i; k<r ; k++) {
2740 /* start with i, the leftmost element that changed */
2741 /* yield tuple(pool[k] for k in indices[:r]) */
2742 index = indices[k];
2743 elem = PyTuple_GET_ITEM(pool, index);
2744 Py_INCREF(elem);
2745 oldelem = PyTuple_GET_ITEM(result, k);
2746 PyTuple_SET_ITEM(result, k, elem);
2747 Py_DECREF(oldelem);
2748 }
2749 break;
2750 }
2751 }
2752 /* If i is negative, then the cycles have all
2753 rolled-over and we're done. */
2754 if (i < 0)
2755 goto empty;
2756 }
2757 Py_INCREF(result);
2758 return result;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002759
2760empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002761 po->stopped = 1;
2762 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002763}
2764
2765PyDoc_STRVAR(permutations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +00002766"permutations(iterable[, r]) --> permutations object\n\
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002767\n\
2768Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl8550d692008-06-10 12:46:39 +00002769permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002770
2771static PyTypeObject permutations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002772 PyVarObject_HEAD_INIT(NULL, 0)
2773 "itertools.permutations", /* tp_name */
2774 sizeof(permutationsobject), /* tp_basicsize */
2775 0, /* tp_itemsize */
2776 /* methods */
2777 (destructor)permutations_dealloc, /* tp_dealloc */
2778 0, /* tp_print */
2779 0, /* tp_getattr */
2780 0, /* tp_setattr */
2781 0, /* tp_compare */
2782 0, /* tp_repr */
2783 0, /* tp_as_number */
2784 0, /* tp_as_sequence */
2785 0, /* tp_as_mapping */
2786 0, /* tp_hash */
2787 0, /* tp_call */
2788 0, /* tp_str */
2789 PyObject_GenericGetAttr, /* tp_getattro */
2790 0, /* tp_setattro */
2791 0, /* tp_as_buffer */
2792 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2793 Py_TPFLAGS_BASETYPE, /* tp_flags */
2794 permutations_doc, /* tp_doc */
2795 (traverseproc)permutations_traverse, /* tp_traverse */
2796 0, /* tp_clear */
2797 0, /* tp_richcompare */
2798 0, /* tp_weaklistoffset */
2799 PyObject_SelfIter, /* tp_iter */
2800 (iternextfunc)permutations_next, /* tp_iternext */
2801 0, /* tp_methods */
2802 0, /* tp_members */
2803 0, /* tp_getset */
2804 0, /* tp_base */
2805 0, /* tp_dict */
2806 0, /* tp_descr_get */
2807 0, /* tp_descr_set */
2808 0, /* tp_dictoffset */
2809 0, /* tp_init */
2810 0, /* tp_alloc */
2811 permutations_new, /* tp_new */
2812 PyObject_GC_Del, /* tp_free */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002813};
2814
2815
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002816/* compress object ************************************************************/
2817
2818/* Equivalent to:
2819
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002820 def compress(data, selectors):
2821 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2822 return (d for d, s in izip(data, selectors) if s)
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002823*/
2824
2825typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002826 PyObject_HEAD
2827 PyObject *data;
2828 PyObject *selectors;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002829} compressobject;
2830
2831static PyTypeObject compress_type;
2832
2833static PyObject *
2834compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2835{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002836 PyObject *seq1, *seq2;
2837 PyObject *data=NULL, *selectors=NULL;
2838 compressobject *lz;
2839 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002840
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002841 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2842 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002843
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002844 data = PyObject_GetIter(seq1);
2845 if (data == NULL)
2846 goto fail;
2847 selectors = PyObject_GetIter(seq2);
2848 if (selectors == NULL)
2849 goto fail;
2850
2851 /* create compressobject structure */
2852 lz = (compressobject *)type->tp_alloc(type, 0);
2853 if (lz == NULL)
2854 goto fail;
2855 lz->data = data;
2856 lz->selectors = selectors;
2857 return (PyObject *)lz;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002858
2859fail:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002860 Py_XDECREF(data);
2861 Py_XDECREF(selectors);
2862 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002863}
2864
2865static void
2866compress_dealloc(compressobject *lz)
2867{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002868 PyObject_GC_UnTrack(lz);
2869 Py_XDECREF(lz->data);
2870 Py_XDECREF(lz->selectors);
2871 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002872}
2873
2874static int
2875compress_traverse(compressobject *lz, visitproc visit, void *arg)
2876{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002877 Py_VISIT(lz->data);
2878 Py_VISIT(lz->selectors);
2879 return 0;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002880}
2881
2882static PyObject *
2883compress_next(compressobject *lz)
2884{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002885 PyObject *data = lz->data, *selectors = lz->selectors;
2886 PyObject *datum, *selector;
2887 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2888 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2889 int ok;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002890
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002891 while (1) {
2892 /* Steps: get datum, get selector, evaluate selector.
2893 Order is important (to match the pure python version
2894 in terms of which input gets a chance to raise an
2895 exception first).
2896 */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002897
Serhiy Storchakabb845652013-04-06 22:04:10 +03002898 datum = datanext(data);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03002899 if (datum == NULL)
Serhiy Storchakabb845652013-04-06 22:04:10 +03002900 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002901
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002902 selector = selectornext(selectors);
2903 if (selector == NULL) {
2904 Py_DECREF(datum);
2905 return NULL;
2906 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002907
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002908 ok = PyObject_IsTrue(selector);
2909 Py_DECREF(selector);
2910 if (ok == 1)
2911 return datum;
2912 Py_DECREF(datum);
2913 if (ok == -1)
2914 return NULL;
2915 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002916}
2917
2918PyDoc_STRVAR(compress_doc,
Raymond Hettinger2e2909f2009-02-19 02:15:14 +00002919"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002920\n\
2921Return data elements corresponding to true selector elements.\n\
2922Forms a shorter iterator from selected data elements using the\n\
2923selectors to choose the data elements.");
2924
2925static PyTypeObject compress_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002926 PyVarObject_HEAD_INIT(NULL, 0)
2927 "itertools.compress", /* tp_name */
2928 sizeof(compressobject), /* tp_basicsize */
2929 0, /* tp_itemsize */
2930 /* methods */
2931 (destructor)compress_dealloc, /* tp_dealloc */
2932 0, /* tp_print */
2933 0, /* tp_getattr */
2934 0, /* tp_setattr */
2935 0, /* tp_compare */
2936 0, /* tp_repr */
2937 0, /* tp_as_number */
2938 0, /* tp_as_sequence */
2939 0, /* tp_as_mapping */
2940 0, /* tp_hash */
2941 0, /* tp_call */
2942 0, /* tp_str */
2943 PyObject_GenericGetAttr, /* tp_getattro */
2944 0, /* tp_setattro */
2945 0, /* tp_as_buffer */
2946 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2947 Py_TPFLAGS_BASETYPE, /* tp_flags */
2948 compress_doc, /* tp_doc */
2949 (traverseproc)compress_traverse, /* tp_traverse */
2950 0, /* tp_clear */
2951 0, /* tp_richcompare */
2952 0, /* tp_weaklistoffset */
2953 PyObject_SelfIter, /* tp_iter */
2954 (iternextfunc)compress_next, /* tp_iternext */
2955 0, /* tp_methods */
2956 0, /* tp_members */
2957 0, /* tp_getset */
2958 0, /* tp_base */
2959 0, /* tp_dict */
2960 0, /* tp_descr_get */
2961 0, /* tp_descr_set */
2962 0, /* tp_dictoffset */
2963 0, /* tp_init */
2964 0, /* tp_alloc */
2965 compress_new, /* tp_new */
2966 PyObject_GC_Del, /* tp_free */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002967};
2968
2969
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002970/* ifilter object ************************************************************/
2971
2972typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002973 PyObject_HEAD
2974 PyObject *func;
2975 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002976} ifilterobject;
2977
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002978static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002979
2980static PyObject *
2981ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2982{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002983 PyObject *func, *seq;
2984 PyObject *it;
2985 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002986
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002987 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
2988 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002989
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002990 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
2991 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002992
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002993 /* Get iterator. */
2994 it = PyObject_GetIter(seq);
2995 if (it == NULL)
2996 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002997
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002998 /* create ifilterobject structure */
2999 lz = (ifilterobject *)type->tp_alloc(type, 0);
3000 if (lz == NULL) {
3001 Py_DECREF(it);
3002 return NULL;
3003 }
3004 Py_INCREF(func);
3005 lz->func = func;
3006 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003007
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003008 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003009}
3010
3011static void
3012ifilter_dealloc(ifilterobject *lz)
3013{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003014 PyObject_GC_UnTrack(lz);
3015 Py_XDECREF(lz->func);
3016 Py_XDECREF(lz->it);
3017 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003018}
3019
3020static int
3021ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
3022{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003023 Py_VISIT(lz->it);
3024 Py_VISIT(lz->func);
3025 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003026}
3027
3028static PyObject *
3029ifilter_next(ifilterobject *lz)
3030{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003031 PyObject *item;
3032 PyObject *it = lz->it;
3033 long ok;
3034 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003035
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003036 iternext = *Py_TYPE(it)->tp_iternext;
3037 for (;;) {
3038 item = iternext(it);
3039 if (item == NULL)
3040 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003041
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003042 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3043 ok = PyObject_IsTrue(item);
3044 } else {
3045 PyObject *good;
3046 good = PyObject_CallFunctionObjArgs(lz->func,
3047 item, NULL);
3048 if (good == NULL) {
3049 Py_DECREF(item);
3050 return NULL;
3051 }
3052 ok = PyObject_IsTrue(good);
3053 Py_DECREF(good);
3054 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003055 if (ok > 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003056 return item;
3057 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003058 if (ok < 0)
3059 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003060 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003061}
3062
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003063PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003064"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003065\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003066Return those items of sequence for which function(item) is true.\n\
3067If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003068
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003069static PyTypeObject ifilter_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003070 PyVarObject_HEAD_INIT(NULL, 0)
3071 "itertools.ifilter", /* tp_name */
3072 sizeof(ifilterobject), /* tp_basicsize */
3073 0, /* tp_itemsize */
3074 /* methods */
3075 (destructor)ifilter_dealloc, /* tp_dealloc */
3076 0, /* tp_print */
3077 0, /* tp_getattr */
3078 0, /* tp_setattr */
3079 0, /* tp_compare */
3080 0, /* tp_repr */
3081 0, /* tp_as_number */
3082 0, /* tp_as_sequence */
3083 0, /* tp_as_mapping */
3084 0, /* tp_hash */
3085 0, /* tp_call */
3086 0, /* tp_str */
3087 PyObject_GenericGetAttr, /* tp_getattro */
3088 0, /* tp_setattro */
3089 0, /* tp_as_buffer */
3090 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3091 Py_TPFLAGS_BASETYPE, /* tp_flags */
3092 ifilter_doc, /* tp_doc */
3093 (traverseproc)ifilter_traverse, /* tp_traverse */
3094 0, /* tp_clear */
3095 0, /* tp_richcompare */
3096 0, /* tp_weaklistoffset */
3097 PyObject_SelfIter, /* tp_iter */
3098 (iternextfunc)ifilter_next, /* tp_iternext */
3099 0, /* tp_methods */
3100 0, /* tp_members */
3101 0, /* tp_getset */
3102 0, /* tp_base */
3103 0, /* tp_dict */
3104 0, /* tp_descr_get */
3105 0, /* tp_descr_set */
3106 0, /* tp_dictoffset */
3107 0, /* tp_init */
3108 0, /* tp_alloc */
3109 ifilter_new, /* tp_new */
3110 PyObject_GC_Del, /* tp_free */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003111};
3112
3113
3114/* ifilterfalse object ************************************************************/
3115
3116typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003117 PyObject_HEAD
3118 PyObject *func;
3119 PyObject *it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003120} ifilterfalseobject;
3121
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003122static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003123
3124static PyObject *
3125ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3126{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003127 PyObject *func, *seq;
3128 PyObject *it;
3129 ifilterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003130
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003131 if (type == &ifilterfalse_type &&
3132 !_PyArg_NoKeywords("ifilterfalse()", kwds))
3133 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003134
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003135 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
3136 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003137
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003138 /* Get iterator. */
3139 it = PyObject_GetIter(seq);
3140 if (it == NULL)
3141 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003142
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003143 /* create ifilterfalseobject structure */
3144 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
3145 if (lz == NULL) {
3146 Py_DECREF(it);
3147 return NULL;
3148 }
3149 Py_INCREF(func);
3150 lz->func = func;
3151 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003152
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003153 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003154}
3155
3156static void
3157ifilterfalse_dealloc(ifilterfalseobject *lz)
3158{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003159 PyObject_GC_UnTrack(lz);
3160 Py_XDECREF(lz->func);
3161 Py_XDECREF(lz->it);
3162 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003163}
3164
3165static int
3166ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
3167{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003168 Py_VISIT(lz->it);
3169 Py_VISIT(lz->func);
3170 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003171}
3172
3173static PyObject *
3174ifilterfalse_next(ifilterfalseobject *lz)
3175{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003176 PyObject *item;
3177 PyObject *it = lz->it;
3178 long ok;
3179 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003180
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003181 iternext = *Py_TYPE(it)->tp_iternext;
3182 for (;;) {
3183 item = iternext(it);
3184 if (item == NULL)
3185 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003186
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003187 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3188 ok = PyObject_IsTrue(item);
3189 } else {
3190 PyObject *good;
3191 good = PyObject_CallFunctionObjArgs(lz->func,
3192 item, NULL);
3193 if (good == NULL) {
3194 Py_DECREF(item);
3195 return NULL;
3196 }
3197 ok = PyObject_IsTrue(good);
3198 Py_DECREF(good);
3199 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003200 if (ok == 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003201 return item;
3202 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003203 if (ok < 0)
3204 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003205 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003206}
3207
Raymond Hettinger60eca932003-02-09 06:40:58 +00003208PyDoc_STRVAR(ifilterfalse_doc,
3209"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
3210\n\
3211Return those items of sequence for which function(item) is false.\n\
3212If function is None, return the items that are false.");
3213
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003214static PyTypeObject ifilterfalse_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003215 PyVarObject_HEAD_INIT(NULL, 0)
3216 "itertools.ifilterfalse", /* tp_name */
3217 sizeof(ifilterfalseobject), /* tp_basicsize */
3218 0, /* tp_itemsize */
3219 /* methods */
3220 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
3221 0, /* tp_print */
3222 0, /* tp_getattr */
3223 0, /* tp_setattr */
3224 0, /* tp_compare */
3225 0, /* tp_repr */
3226 0, /* tp_as_number */
3227 0, /* tp_as_sequence */
3228 0, /* tp_as_mapping */
3229 0, /* tp_hash */
3230 0, /* tp_call */
3231 0, /* tp_str */
3232 PyObject_GenericGetAttr, /* tp_getattro */
3233 0, /* tp_setattro */
3234 0, /* tp_as_buffer */
3235 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3236 Py_TPFLAGS_BASETYPE, /* tp_flags */
3237 ifilterfalse_doc, /* tp_doc */
3238 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
3239 0, /* tp_clear */
3240 0, /* tp_richcompare */
3241 0, /* tp_weaklistoffset */
3242 PyObject_SelfIter, /* tp_iter */
3243 (iternextfunc)ifilterfalse_next, /* tp_iternext */
3244 0, /* tp_methods */
3245 0, /* tp_members */
3246 0, /* tp_getset */
3247 0, /* tp_base */
3248 0, /* tp_dict */
3249 0, /* tp_descr_get */
3250 0, /* tp_descr_set */
3251 0, /* tp_dictoffset */
3252 0, /* tp_init */
3253 0, /* tp_alloc */
3254 ifilterfalse_new, /* tp_new */
3255 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003256};
3257
3258
3259/* count object ************************************************************/
3260
3261typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003262 PyObject_HEAD
3263 Py_ssize_t cnt;
3264 PyObject *long_cnt;
3265 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003266} countobject;
3267
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003268/* Counting logic and invariants:
3269
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003270fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003271
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003272 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3273 Advances with: cnt += 1
3274 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003275
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003276slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003277
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003278 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3279 All counting is done with python objects (no overflows or underflows).
3280 Advances with: long_cnt += long_step
3281 Step may be zero -- effectively a slow version of repeat(cnt).
3282 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003283*/
3284
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003285static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003286
3287static PyObject *
3288count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3289{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003290 countobject *lz;
3291 int slow_mode = 0;
3292 Py_ssize_t cnt = 0;
3293 PyObject *long_cnt = NULL;
3294 PyObject *long_step = NULL;
3295 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003296
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003297 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3298 kwlist, &long_cnt, &long_step))
3299 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003300
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003301 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3302 (long_step != NULL && !PyNumber_Check(long_step))) {
3303 PyErr_SetString(PyExc_TypeError, "a number is required");
3304 return NULL;
3305 }
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003306
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003307 if (long_cnt != NULL) {
3308 cnt = PyInt_AsSsize_t(long_cnt);
3309 if ((cnt == -1 && PyErr_Occurred()) || !PyInt_Check(long_cnt)) {
3310 PyErr_Clear();
3311 slow_mode = 1;
3312 }
3313 Py_INCREF(long_cnt);
3314 } else {
3315 cnt = 0;
3316 long_cnt = PyInt_FromLong(0);
3317 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003318
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003319 /* If not specified, step defaults to 1 */
3320 if (long_step == NULL) {
3321 long_step = PyInt_FromLong(1);
3322 if (long_step == NULL) {
3323 Py_DECREF(long_cnt);
3324 return NULL;
3325 }
3326 } else
3327 Py_INCREF(long_step);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003328
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003329 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003330
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003331 /* Fast mode only works when the step is 1 */
3332 if (!PyInt_Check(long_step) ||
3333 PyInt_AS_LONG(long_step) != 1) {
3334 slow_mode = 1;
3335 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003336
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003337 if (slow_mode)
3338 cnt = PY_SSIZE_T_MAX;
3339 else
3340 Py_CLEAR(long_cnt);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003341
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003342 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3343 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3344 assert(slow_mode ||
3345 (PyInt_Check(long_step) && PyInt_AS_LONG(long_step) == 1));
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003346
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003347 /* create countobject structure */
3348 lz = (countobject *)type->tp_alloc(type, 0);
3349 if (lz == NULL) {
3350 Py_XDECREF(long_cnt);
3351 return NULL;
3352 }
3353 lz->cnt = cnt;
3354 lz->long_cnt = long_cnt;
3355 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003356
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003357 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003358}
3359
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003360static void
3361count_dealloc(countobject *lz)
3362{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003363 PyObject_GC_UnTrack(lz);
3364 Py_XDECREF(lz->long_cnt);
3365 Py_XDECREF(lz->long_step);
3366 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003367}
3368
Benjamin Peterson062a7c32009-02-16 21:07:52 +00003369static int
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003370count_traverse(countobject *lz, visitproc visit, void *arg)
3371{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003372 Py_VISIT(lz->long_cnt);
3373 Py_VISIT(lz->long_step);
3374 return 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003375}
3376
3377static PyObject *
3378count_nextlong(countobject *lz)
3379{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003380 PyObject *long_cnt;
3381 PyObject *stepped_up;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003382
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003383 long_cnt = lz->long_cnt;
3384 if (long_cnt == NULL) {
3385 /* Switch to slow_mode */
3386 long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
3387 if (long_cnt == NULL)
3388 return NULL;
3389 }
3390 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003391
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003392 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3393 if (stepped_up == NULL)
3394 return NULL;
3395 lz->long_cnt = stepped_up;
3396 return long_cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003397}
3398
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003399static PyObject *
3400count_next(countobject *lz)
3401{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003402 if (lz->cnt == PY_SSIZE_T_MAX)
3403 return count_nextlong(lz);
3404 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003405}
3406
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003407static PyObject *
3408count_repr(countobject *lz)
3409{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003410 PyObject *cnt_repr, *step_repr = NULL;
3411 PyObject *result = NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003412
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003413 if (lz->cnt != PY_SSIZE_T_MAX)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003414 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003415
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003416 cnt_repr = PyObject_Repr(lz->long_cnt);
3417 if (cnt_repr == NULL)
3418 return NULL;
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003419
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003420 if (PyInt_Check(lz->long_step) && PyInt_AS_LONG(lz->long_step) == 1) {
3421 /* Don't display step when it is an integer equal to 1 */
3422 result = PyString_FromFormat("count(%s)",
3423 PyString_AS_STRING(cnt_repr));
3424 } else {
3425 step_repr = PyObject_Repr(lz->long_step);
3426 if (step_repr != NULL)
3427 result = PyString_FromFormat("count(%s, %s)",
3428 PyString_AS_STRING(cnt_repr),
3429 PyString_AS_STRING(step_repr));
3430 }
3431 Py_DECREF(cnt_repr);
3432 Py_XDECREF(step_repr);
3433 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003434}
3435
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003436static PyObject *
3437count_reduce(countobject *lz)
3438{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003439 if (lz->cnt == PY_SSIZE_T_MAX)
3440 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3441 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003442}
3443
3444PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3445
3446static PyMethodDef count_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003447 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3448 count_reduce_doc},
3449 {NULL, NULL} /* sentinel */
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003450};
3451
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003452PyDoc_STRVAR(count_doc,
Ezio Melottifdc1e0d2010-06-11 02:21:25 +00003453 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003454\n\
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003455Return a count object whose .next() method returns consecutive values.\n\
3456Equivalent to:\n\n\
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003457 def count(firstval=0, step=1):\n\
Eli Benderskye6802db2013-09-03 06:41:58 -07003458 x = firstval\n\
3459 while 1:\n\
3460 yield x\n\
3461 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003462
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003463static PyTypeObject count_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003464 PyVarObject_HEAD_INIT(NULL, 0)
3465 "itertools.count", /* tp_name */
3466 sizeof(countobject), /* tp_basicsize */
3467 0, /* tp_itemsize */
3468 /* methods */
3469 (destructor)count_dealloc, /* tp_dealloc */
3470 0, /* tp_print */
3471 0, /* tp_getattr */
3472 0, /* tp_setattr */
3473 0, /* tp_compare */
3474 (reprfunc)count_repr, /* tp_repr */
3475 0, /* tp_as_number */
3476 0, /* tp_as_sequence */
3477 0, /* tp_as_mapping */
3478 0, /* tp_hash */
3479 0, /* tp_call */
3480 0, /* tp_str */
3481 PyObject_GenericGetAttr, /* tp_getattro */
3482 0, /* tp_setattro */
3483 0, /* tp_as_buffer */
3484 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3485 Py_TPFLAGS_BASETYPE, /* tp_flags */
3486 count_doc, /* tp_doc */
3487 (traverseproc)count_traverse, /* tp_traverse */
3488 0, /* tp_clear */
3489 0, /* tp_richcompare */
3490 0, /* tp_weaklistoffset */
3491 PyObject_SelfIter, /* tp_iter */
3492 (iternextfunc)count_next, /* tp_iternext */
3493 count_methods, /* tp_methods */
3494 0, /* tp_members */
3495 0, /* tp_getset */
3496 0, /* tp_base */
3497 0, /* tp_dict */
3498 0, /* tp_descr_get */
3499 0, /* tp_descr_set */
3500 0, /* tp_dictoffset */
3501 0, /* tp_init */
3502 0, /* tp_alloc */
3503 count_new, /* tp_new */
3504 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003505};
3506
3507
3508/* izip object ************************************************************/
3509
3510#include "Python.h"
3511
3512typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003513 PyObject_HEAD
3514 Py_ssize_t tuplesize;
3515 PyObject *ittuple; /* tuple of iterators */
3516 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003517} izipobject;
3518
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003519static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003520
3521static PyObject *
3522izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3523{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003524 izipobject *lz;
3525 Py_ssize_t i;
3526 PyObject *ittuple; /* tuple of iterators */
3527 PyObject *result;
3528 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003529
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003530 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
3531 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003532
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003533 /* args must be a tuple */
3534 assert(PyTuple_Check(args));
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003535
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003536 /* obtain iterators */
3537 ittuple = PyTuple_New(tuplesize);
3538 if (ittuple == NULL)
3539 return NULL;
3540 for (i=0; i < tuplesize; ++i) {
3541 PyObject *item = PyTuple_GET_ITEM(args, i);
3542 PyObject *it = PyObject_GetIter(item);
3543 if (it == NULL) {
3544 if (PyErr_ExceptionMatches(PyExc_TypeError))
3545 PyErr_Format(PyExc_TypeError,
3546 "izip argument #%zd must support iteration",
3547 i+1);
3548 Py_DECREF(ittuple);
3549 return NULL;
3550 }
3551 PyTuple_SET_ITEM(ittuple, i, it);
3552 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003553
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003554 /* create a result holder */
3555 result = PyTuple_New(tuplesize);
3556 if (result == NULL) {
3557 Py_DECREF(ittuple);
3558 return NULL;
3559 }
3560 for (i=0 ; i < tuplesize ; i++) {
3561 Py_INCREF(Py_None);
3562 PyTuple_SET_ITEM(result, i, Py_None);
3563 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003564
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003565 /* create izipobject structure */
3566 lz = (izipobject *)type->tp_alloc(type, 0);
3567 if (lz == NULL) {
3568 Py_DECREF(ittuple);
3569 Py_DECREF(result);
3570 return NULL;
3571 }
3572 lz->ittuple = ittuple;
3573 lz->tuplesize = tuplesize;
3574 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003575
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003576 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003577}
3578
3579static void
3580izip_dealloc(izipobject *lz)
3581{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003582 PyObject_GC_UnTrack(lz);
3583 Py_XDECREF(lz->ittuple);
3584 Py_XDECREF(lz->result);
3585 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003586}
3587
3588static int
3589izip_traverse(izipobject *lz, visitproc visit, void *arg)
3590{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003591 Py_VISIT(lz->ittuple);
3592 Py_VISIT(lz->result);
3593 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003594}
3595
3596static PyObject *
3597izip_next(izipobject *lz)
3598{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003599 Py_ssize_t i;
3600 Py_ssize_t tuplesize = lz->tuplesize;
3601 PyObject *result = lz->result;
3602 PyObject *it;
3603 PyObject *item;
3604 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003605
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003606 if (tuplesize == 0)
3607 return NULL;
3608 if (Py_REFCNT(result) == 1) {
3609 Py_INCREF(result);
3610 for (i=0 ; i < tuplesize ; i++) {
3611 it = PyTuple_GET_ITEM(lz->ittuple, i);
3612 item = (*Py_TYPE(it)->tp_iternext)(it);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003613 if (item == NULL) {
3614 Py_DECREF(result);
3615 return NULL;
3616 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003617 olditem = PyTuple_GET_ITEM(result, i);
3618 PyTuple_SET_ITEM(result, i, item);
3619 Py_DECREF(olditem);
3620 }
3621 } else {
3622 result = PyTuple_New(tuplesize);
3623 if (result == NULL)
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003624 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003625 for (i=0 ; i < tuplesize ; i++) {
3626 it = PyTuple_GET_ITEM(lz->ittuple, i);
3627 item = (*Py_TYPE(it)->tp_iternext)(it);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003628 if (item == NULL) {
3629 Py_DECREF(result);
3630 return NULL;
3631 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003632 PyTuple_SET_ITEM(result, i, item);
3633 }
3634 }
3635 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003636}
3637
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003638PyDoc_STRVAR(izip_doc,
3639"izip(iter1 [,iter2 [...]]) --> izip object\n\
3640\n\
3641Return a izip object whose .next() method returns a tuple where\n\
3642the i-th element comes from the i-th iterable argument. The .next()\n\
3643method continues until the shortest iterable in the argument sequence\n\
3644is exhausted and then it raises StopIteration. Works like the zip()\n\
3645function but consumes less memory by returning an iterator instead of\n\
3646a list.");
3647
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003648static PyTypeObject izip_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003649 PyVarObject_HEAD_INIT(NULL, 0)
3650 "itertools.izip", /* tp_name */
3651 sizeof(izipobject), /* tp_basicsize */
3652 0, /* tp_itemsize */
3653 /* methods */
3654 (destructor)izip_dealloc, /* tp_dealloc */
3655 0, /* tp_print */
3656 0, /* tp_getattr */
3657 0, /* tp_setattr */
3658 0, /* tp_compare */
3659 0, /* tp_repr */
3660 0, /* tp_as_number */
3661 0, /* tp_as_sequence */
3662 0, /* tp_as_mapping */
3663 0, /* tp_hash */
3664 0, /* tp_call */
3665 0, /* tp_str */
3666 PyObject_GenericGetAttr, /* tp_getattro */
3667 0, /* tp_setattro */
3668 0, /* tp_as_buffer */
3669 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3670 Py_TPFLAGS_BASETYPE, /* tp_flags */
3671 izip_doc, /* tp_doc */
3672 (traverseproc)izip_traverse, /* tp_traverse */
3673 0, /* tp_clear */
3674 0, /* tp_richcompare */
3675 0, /* tp_weaklistoffset */
3676 PyObject_SelfIter, /* tp_iter */
3677 (iternextfunc)izip_next, /* tp_iternext */
3678 0, /* tp_methods */
3679 0, /* tp_members */
3680 0, /* tp_getset */
3681 0, /* tp_base */
3682 0, /* tp_dict */
3683 0, /* tp_descr_get */
3684 0, /* tp_descr_set */
3685 0, /* tp_dictoffset */
3686 0, /* tp_init */
3687 0, /* tp_alloc */
3688 izip_new, /* tp_new */
3689 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003690};
3691
3692
3693/* repeat object ************************************************************/
3694
3695typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003696 PyObject_HEAD
3697 PyObject *element;
3698 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003699} repeatobject;
3700
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003701static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003702
3703static PyObject *
3704repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3705{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003706 repeatobject *ro;
3707 PyObject *element;
Raymond Hettinger58ad2452014-06-24 21:53:45 -07003708 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003709 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003710
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003711 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3712 &element, &cnt))
3713 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003714
Raymond Hettinger58ad2452014-06-24 21:53:45 -07003715 if (kwds != NULL)
3716 n_kwds = PyDict_Size(kwds);
3717 /* Does user supply times argument? */
3718 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003719 cnt = 0;
3720
3721 ro = (repeatobject *)type->tp_alloc(type, 0);
3722 if (ro == NULL)
3723 return NULL;
3724 Py_INCREF(element);
3725 ro->element = element;
3726 ro->cnt = cnt;
3727 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003728}
3729
3730static void
3731repeat_dealloc(repeatobject *ro)
3732{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003733 PyObject_GC_UnTrack(ro);
3734 Py_XDECREF(ro->element);
3735 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003736}
3737
3738static int
3739repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3740{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003741 Py_VISIT(ro->element);
3742 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003743}
3744
3745static PyObject *
3746repeat_next(repeatobject *ro)
3747{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003748 if (ro->cnt == 0)
3749 return NULL;
3750 if (ro->cnt > 0)
3751 ro->cnt--;
3752 Py_INCREF(ro->element);
3753 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003754}
3755
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003756static PyObject *
3757repeat_repr(repeatobject *ro)
3758{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003759 PyObject *result, *objrepr;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003760
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003761 objrepr = PyObject_Repr(ro->element);
3762 if (objrepr == NULL)
3763 return NULL;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003764
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003765 if (ro->cnt == -1)
3766 result = PyString_FromFormat("repeat(%s)",
3767 PyString_AS_STRING(objrepr));
3768 else
3769 result = PyString_FromFormat("repeat(%s, %zd)",
3770 PyString_AS_STRING(objrepr), ro->cnt);
3771 Py_DECREF(objrepr);
3772 return result;
3773}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003774
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003775static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003776repeat_len(repeatobject *ro)
3777{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003778 if (ro->cnt == -1) {
3779 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3780 return NULL;
3781 }
3782 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003783}
3784
Armin Rigof5b3e362006-02-11 21:32:43 +00003785PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003786
3787static PyMethodDef repeat_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003788 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3789 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003790};
3791
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003792PyDoc_STRVAR(repeat_doc,
Raymond Hettinger182edae2009-02-19 02:38:25 +00003793"repeat(object [,times]) -> create an iterator which returns the object\n\
3794for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003795endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003796
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003797static PyTypeObject repeat_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003798 PyVarObject_HEAD_INIT(NULL, 0)
3799 "itertools.repeat", /* tp_name */
3800 sizeof(repeatobject), /* tp_basicsize */
3801 0, /* tp_itemsize */
3802 /* methods */
3803 (destructor)repeat_dealloc, /* tp_dealloc */
3804 0, /* tp_print */
3805 0, /* tp_getattr */
3806 0, /* tp_setattr */
3807 0, /* tp_compare */
3808 (reprfunc)repeat_repr, /* tp_repr */
3809 0, /* tp_as_number */
3810 0, /* tp_as_sequence */
3811 0, /* tp_as_mapping */
3812 0, /* tp_hash */
3813 0, /* tp_call */
3814 0, /* tp_str */
3815 PyObject_GenericGetAttr, /* tp_getattro */
3816 0, /* tp_setattro */
3817 0, /* tp_as_buffer */
3818 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3819 Py_TPFLAGS_BASETYPE, /* tp_flags */
3820 repeat_doc, /* tp_doc */
3821 (traverseproc)repeat_traverse, /* tp_traverse */
3822 0, /* tp_clear */
3823 0, /* tp_richcompare */
3824 0, /* tp_weaklistoffset */
3825 PyObject_SelfIter, /* tp_iter */
3826 (iternextfunc)repeat_next, /* tp_iternext */
3827 repeat_methods, /* tp_methods */
3828 0, /* tp_members */
3829 0, /* tp_getset */
3830 0, /* tp_base */
3831 0, /* tp_dict */
3832 0, /* tp_descr_get */
3833 0, /* tp_descr_set */
3834 0, /* tp_dictoffset */
3835 0, /* tp_init */
3836 0, /* tp_alloc */
3837 repeat_new, /* tp_new */
3838 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003839};
3840
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003841/* iziplongest object ************************************************************/
3842
3843#include "Python.h"
3844
3845typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003846 PyObject_HEAD
3847 Py_ssize_t tuplesize;
3848 Py_ssize_t numactive;
3849 PyObject *ittuple; /* tuple of iterators */
3850 PyObject *result;
3851 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003852} iziplongestobject;
3853
3854static PyTypeObject iziplongest_type;
3855
3856static PyObject *
3857izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3858{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003859 iziplongestobject *lz;
3860 Py_ssize_t i;
3861 PyObject *ittuple; /* tuple of iterators */
3862 PyObject *result;
3863 PyObject *fillvalue = Py_None;
3864 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003865
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003866 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3867 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3868 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3869 PyErr_SetString(PyExc_TypeError,
3870 "izip_longest() got an unexpected keyword argument");
3871 return NULL;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003872 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003873 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003874
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003875 /* args must be a tuple */
3876 assert(PyTuple_Check(args));
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003877
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003878 /* obtain iterators */
3879 ittuple = PyTuple_New(tuplesize);
3880 if (ittuple == NULL)
3881 return NULL;
3882 for (i=0; i < tuplesize; ++i) {
3883 PyObject *item = PyTuple_GET_ITEM(args, i);
3884 PyObject *it = PyObject_GetIter(item);
3885 if (it == NULL) {
3886 if (PyErr_ExceptionMatches(PyExc_TypeError))
3887 PyErr_Format(PyExc_TypeError,
3888 "izip_longest argument #%zd must support iteration",
3889 i+1);
3890 Py_DECREF(ittuple);
3891 return NULL;
3892 }
3893 PyTuple_SET_ITEM(ittuple, i, it);
3894 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003895
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003896 /* create a result holder */
3897 result = PyTuple_New(tuplesize);
3898 if (result == NULL) {
3899 Py_DECREF(ittuple);
3900 return NULL;
3901 }
3902 for (i=0 ; i < tuplesize ; i++) {
3903 Py_INCREF(Py_None);
3904 PyTuple_SET_ITEM(result, i, Py_None);
3905 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003906
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003907 /* create iziplongestobject structure */
3908 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3909 if (lz == NULL) {
3910 Py_DECREF(ittuple);
3911 Py_DECREF(result);
3912 return NULL;
3913 }
3914 lz->ittuple = ittuple;
3915 lz->tuplesize = tuplesize;
3916 lz->numactive = tuplesize;
3917 lz->result = result;
3918 Py_INCREF(fillvalue);
3919 lz->fillvalue = fillvalue;
3920 return (PyObject *)lz;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003921}
3922
3923static void
3924izip_longest_dealloc(iziplongestobject *lz)
3925{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003926 PyObject_GC_UnTrack(lz);
3927 Py_XDECREF(lz->ittuple);
3928 Py_XDECREF(lz->result);
3929 Py_XDECREF(lz->fillvalue);
3930 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003931}
3932
3933static int
3934izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3935{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003936 Py_VISIT(lz->ittuple);
3937 Py_VISIT(lz->result);
3938 Py_VISIT(lz->fillvalue);
3939 return 0;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003940}
3941
3942static PyObject *
3943izip_longest_next(iziplongestobject *lz)
3944{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003945 Py_ssize_t i;
3946 Py_ssize_t tuplesize = lz->tuplesize;
3947 PyObject *result = lz->result;
3948 PyObject *it;
3949 PyObject *item;
3950 PyObject *olditem;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003951
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003952 if (tuplesize == 0)
3953 return NULL;
3954 if (lz->numactive == 0)
3955 return NULL;
3956 if (Py_REFCNT(result) == 1) {
3957 Py_INCREF(result);
3958 for (i=0 ; i < tuplesize ; i++) {
3959 it = PyTuple_GET_ITEM(lz->ittuple, i);
3960 if (it == NULL) {
3961 Py_INCREF(lz->fillvalue);
3962 item = lz->fillvalue;
3963 } else {
3964 item = PyIter_Next(it);
3965 if (item == NULL) {
3966 lz->numactive -= 1;
3967 if (lz->numactive == 0 || PyErr_Occurred()) {
3968 lz->numactive = 0;
3969 Py_DECREF(result);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003970 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003971 } else {
3972 Py_INCREF(lz->fillvalue);
3973 item = lz->fillvalue;
3974 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3975 Py_DECREF(it);
3976 }
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003977 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003978 }
3979 olditem = PyTuple_GET_ITEM(result, i);
3980 PyTuple_SET_ITEM(result, i, item);
3981 Py_DECREF(olditem);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003982 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003983 } else {
3984 result = PyTuple_New(tuplesize);
3985 if (result == NULL)
3986 return NULL;
3987 for (i=0 ; i < tuplesize ; i++) {
3988 it = PyTuple_GET_ITEM(lz->ittuple, i);
3989 if (it == NULL) {
3990 Py_INCREF(lz->fillvalue);
3991 item = lz->fillvalue;
3992 } else {
3993 item = PyIter_Next(it);
3994 if (item == NULL) {
3995 lz->numactive -= 1;
3996 if (lz->numactive == 0 || PyErr_Occurred()) {
3997 lz->numactive = 0;
3998 Py_DECREF(result);
3999 return NULL;
4000 } else {
4001 Py_INCREF(lz->fillvalue);
4002 item = lz->fillvalue;
4003 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4004 Py_DECREF(it);
4005 }
4006 }
4007 }
4008 PyTuple_SET_ITEM(result, i, item);
4009 }
4010 }
4011 return result;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00004012}
4013
4014PyDoc_STRVAR(izip_longest_doc,
4015"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
4016\n\
4017Return an izip_longest object whose .next() method returns a tuple where\n\
4018the i-th element comes from the i-th iterable argument. The .next()\n\
4019method continues until the longest iterable in the argument sequence\n\
4020is exhausted and then it raises StopIteration. When the shorter iterables\n\
4021are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4022defaults to None or can be specified by a keyword argument.\n\
4023");
4024
4025static PyTypeObject iziplongest_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004026 PyVarObject_HEAD_INIT(NULL, 0)
4027 "itertools.izip_longest", /* tp_name */
4028 sizeof(iziplongestobject), /* tp_basicsize */
4029 0, /* tp_itemsize */
4030 /* methods */
4031 (destructor)izip_longest_dealloc, /* tp_dealloc */
4032 0, /* tp_print */
4033 0, /* tp_getattr */
4034 0, /* tp_setattr */
4035 0, /* tp_compare */
4036 0, /* tp_repr */
4037 0, /* tp_as_number */
4038 0, /* tp_as_sequence */
4039 0, /* tp_as_mapping */
4040 0, /* tp_hash */
4041 0, /* tp_call */
4042 0, /* tp_str */
4043 PyObject_GenericGetAttr, /* tp_getattro */
4044 0, /* tp_setattro */
4045 0, /* tp_as_buffer */
4046 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4047 Py_TPFLAGS_BASETYPE, /* tp_flags */
4048 izip_longest_doc, /* tp_doc */
4049 (traverseproc)izip_longest_traverse, /* tp_traverse */
4050 0, /* tp_clear */
4051 0, /* tp_richcompare */
4052 0, /* tp_weaklistoffset */
4053 PyObject_SelfIter, /* tp_iter */
4054 (iternextfunc)izip_longest_next, /* tp_iternext */
4055 0, /* tp_methods */
4056 0, /* tp_members */
4057 0, /* tp_getset */
4058 0, /* tp_base */
4059 0, /* tp_dict */
4060 0, /* tp_descr_get */
4061 0, /* tp_descr_set */
4062 0, /* tp_dictoffset */
4063 0, /* tp_init */
4064 0, /* tp_alloc */
4065 izip_longest_new, /* tp_new */
4066 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd36862c2007-02-21 05:20:38 +00004067};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004068
4069/* module level code ********************************************************/
4070
4071PyDoc_STRVAR(module_doc,
4072"Functional tools for creating and using iterators.\n\
4073\n\
4074Infinite iterators:\n\
4075count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004076cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004077repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004078\n\
4079Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004080chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettingere76816b2009-01-29 03:43:44 +00004081compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004082dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4083groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00004084ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
4085ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004086islice(seq, [start,] stop [, step]) --> elements from\n\
4087 seq[start:stop:step]\n\
4088imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
4089starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004090tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004091takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004092izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4093izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4094\n\
4095Combinatoric generators:\n\
4096product(p, q, ... [repeat=1]) --> cartesian product\n\
4097permutations(p[, r])\n\
Raymond Hettinger9eac1192009-11-19 01:22:04 +00004098combinations(p, r)\n\
4099combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004100");
4101
4102
Raymond Hettingerad983e72003-11-12 14:32:26 +00004103static PyMethodDef module_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004104 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4105 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004106};
4107
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004108PyMODINIT_FUNC
4109inititertools(void)
4110{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004111 int i;
4112 PyObject *m;
4113 char *name;
4114 PyTypeObject *typelist[] = {
4115 &combinations_type,
4116 &cwr_type,
4117 &cycle_type,
4118 &dropwhile_type,
4119 &takewhile_type,
4120 &islice_type,
4121 &starmap_type,
4122 &imap_type,
4123 &chain_type,
4124 &compress_type,
4125 &ifilter_type,
4126 &ifilterfalse_type,
4127 &count_type,
4128 &izip_type,
4129 &iziplongest_type,
4130 &permutations_type,
4131 &product_type,
4132 &repeat_type,
4133 &groupby_type,
4134 NULL
4135 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004136
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004137 Py_TYPE(&teedataobject_type) = &PyType_Type;
4138 m = Py_InitModule3("itertools", module_methods, module_doc);
4139 if (m == NULL)
4140 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004141
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004142 for (i=0 ; typelist[i] != NULL ; i++) {
4143 if (PyType_Ready(typelist[i]) < 0)
4144 return;
4145 name = strchr(typelist[i]->tp_name, '.');
4146 assert (name != NULL);
4147 Py_INCREF(typelist[i]);
4148 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4149 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004150
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004151 if (PyType_Ready(&teedataobject_type) < 0)
4152 return;
4153 if (PyType_Ready(&tee_type) < 0)
4154 return;
4155 if (PyType_Ready(&_grouper_type) < 0)
4156 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004157}