blob: e5b2b7443954b20f46e234b007c57a07270c9c03 [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);
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02001850 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Petersondda91212015-02-01 21:34:07 -05001851 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
1852 return NULL;
1853 }
1854 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001855 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001856
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02001857 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001858 if (indices == NULL) {
1859 PyErr_NoMemory();
1860 goto error;
1861 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001862
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001863 pools = PyTuple_New(npools);
1864 if (pools == NULL)
1865 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001866
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001867 for (i=0; i < nargs ; ++i) {
1868 PyObject *item = PyTuple_GET_ITEM(args, i);
1869 PyObject *pool = PySequence_Tuple(item);
1870 if (pool == NULL)
1871 goto error;
1872 PyTuple_SET_ITEM(pools, i, pool);
1873 indices[i] = 0;
1874 }
1875 for ( ; i < npools; ++i) {
1876 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1877 Py_INCREF(pool);
1878 PyTuple_SET_ITEM(pools, i, pool);
1879 indices[i] = 0;
1880 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001881
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001882 /* create productobject structure */
1883 lz = (productobject *)type->tp_alloc(type, 0);
1884 if (lz == NULL)
1885 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001886
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001887 lz->pools = pools;
1888 lz->indices = indices;
1889 lz->result = NULL;
1890 lz->stopped = 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001891
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001892 return (PyObject *)lz;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001893
1894error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001895 if (indices != NULL)
1896 PyMem_Free(indices);
1897 Py_XDECREF(pools);
1898 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001899}
1900
1901static void
1902product_dealloc(productobject *lz)
1903{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001904 PyObject_GC_UnTrack(lz);
1905 Py_XDECREF(lz->pools);
1906 Py_XDECREF(lz->result);
1907 if (lz->indices != NULL)
1908 PyMem_Free(lz->indices);
1909 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001910}
1911
1912static int
1913product_traverse(productobject *lz, visitproc visit, void *arg)
1914{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001915 Py_VISIT(lz->pools);
1916 Py_VISIT(lz->result);
1917 return 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001918}
1919
1920static PyObject *
1921product_next(productobject *lz)
1922{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001923 PyObject *pool;
1924 PyObject *elem;
1925 PyObject *oldelem;
1926 PyObject *pools = lz->pools;
1927 PyObject *result = lz->result;
1928 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1929 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001930
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001931 if (lz->stopped)
1932 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001933
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001934 if (result == NULL) {
1935 /* On the first pass, return an initial tuple filled with the
1936 first element from each pool. */
1937 result = PyTuple_New(npools);
1938 if (result == NULL)
1939 goto empty;
1940 lz->result = result;
1941 for (i=0; i < npools; i++) {
1942 pool = PyTuple_GET_ITEM(pools, i);
1943 if (PyTuple_GET_SIZE(pool) == 0)
1944 goto empty;
1945 elem = PyTuple_GET_ITEM(pool, 0);
1946 Py_INCREF(elem);
1947 PyTuple_SET_ITEM(result, i, elem);
1948 }
1949 } else {
1950 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001951
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001952 /* Copy the previous result tuple or re-use it if available */
1953 if (Py_REFCNT(result) > 1) {
1954 PyObject *old_result = result;
1955 result = PyTuple_New(npools);
1956 if (result == NULL)
1957 goto empty;
1958 lz->result = result;
1959 for (i=0; i < npools; i++) {
1960 elem = PyTuple_GET_ITEM(old_result, i);
1961 Py_INCREF(elem);
1962 PyTuple_SET_ITEM(result, i, elem);
1963 }
1964 Py_DECREF(old_result);
1965 }
1966 /* Now, we've got the only copy so we can update it in-place */
1967 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001968
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001969 /* Update the pool indices right-to-left. Only advance to the
1970 next pool when the previous one rolls-over */
1971 for (i=npools-1 ; i >= 0 ; i--) {
1972 pool = PyTuple_GET_ITEM(pools, i);
1973 indices[i]++;
1974 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1975 /* Roll-over and advance to next pool */
1976 indices[i] = 0;
1977 elem = PyTuple_GET_ITEM(pool, 0);
1978 Py_INCREF(elem);
1979 oldelem = PyTuple_GET_ITEM(result, i);
1980 PyTuple_SET_ITEM(result, i, elem);
1981 Py_DECREF(oldelem);
1982 } else {
1983 /* No rollover. Just increment and stop here. */
1984 elem = PyTuple_GET_ITEM(pool, indices[i]);
1985 Py_INCREF(elem);
1986 oldelem = PyTuple_GET_ITEM(result, i);
1987 PyTuple_SET_ITEM(result, i, elem);
1988 Py_DECREF(oldelem);
1989 break;
1990 }
1991 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00001992
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001993 /* If i is negative, then the indices have all rolled-over
1994 and we're done. */
1995 if (i < 0)
1996 goto empty;
1997 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001998
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001999 Py_INCREF(result);
2000 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002001
2002empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002003 lz->stopped = 1;
2004 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002005}
2006
2007PyDoc_STRVAR(product_doc,
2008"product(*iterables) --> product object\n\
2009\n\
Raymond Hettinger73d79632008-02-23 02:20:41 +00002010Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002011For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2012The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2013cycle in a manner similar to an odometer (with the rightmost element changing\n\
2014on every iteration).\n\n\
Raymond Hettinger94034ea2009-02-09 18:39:41 +00002015To compute the product of an iterable with itself, specify the number\n\
2016of repetitions with the optional repeat keyword argument. For example,\n\
2017product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002018product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2019product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2020
2021static PyTypeObject product_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002022 PyVarObject_HEAD_INIT(NULL, 0)
2023 "itertools.product", /* tp_name */
2024 sizeof(productobject), /* tp_basicsize */
2025 0, /* tp_itemsize */
2026 /* methods */
2027 (destructor)product_dealloc, /* tp_dealloc */
2028 0, /* tp_print */
2029 0, /* tp_getattr */
2030 0, /* tp_setattr */
2031 0, /* tp_compare */
2032 0, /* tp_repr */
2033 0, /* tp_as_number */
2034 0, /* tp_as_sequence */
2035 0, /* tp_as_mapping */
2036 0, /* tp_hash */
2037 0, /* tp_call */
2038 0, /* tp_str */
2039 PyObject_GenericGetAttr, /* tp_getattro */
2040 0, /* tp_setattro */
2041 0, /* tp_as_buffer */
2042 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2043 Py_TPFLAGS_BASETYPE, /* tp_flags */
2044 product_doc, /* tp_doc */
2045 (traverseproc)product_traverse, /* tp_traverse */
2046 0, /* tp_clear */
2047 0, /* tp_richcompare */
2048 0, /* tp_weaklistoffset */
2049 PyObject_SelfIter, /* tp_iter */
2050 (iternextfunc)product_next, /* tp_iternext */
2051 0, /* tp_methods */
2052 0, /* tp_members */
2053 0, /* tp_getset */
2054 0, /* tp_base */
2055 0, /* tp_dict */
2056 0, /* tp_descr_get */
2057 0, /* tp_descr_set */
2058 0, /* tp_dictoffset */
2059 0, /* tp_init */
2060 0, /* tp_alloc */
2061 product_new, /* tp_new */
2062 PyObject_GC_Del, /* tp_free */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002063};
2064
2065
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002066/* combinations object ************************************************************/
2067
2068typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002069 PyObject_HEAD
2070 PyObject *pool; /* input converted to a tuple */
2071 Py_ssize_t *indices; /* one index per result element */
2072 PyObject *result; /* most recently returned result tuple */
2073 Py_ssize_t r; /* size of result tuple */
2074 int stopped; /* set to 1 when the combinations iterator is exhausted */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002075} combinationsobject;
2076
2077static PyTypeObject combinations_type;
2078
2079static PyObject *
2080combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2081{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002082 combinationsobject *co;
2083 Py_ssize_t n;
2084 Py_ssize_t r;
2085 PyObject *pool = NULL;
2086 PyObject *iterable = NULL;
2087 Py_ssize_t *indices = NULL;
2088 Py_ssize_t i;
2089 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002090
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002091 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2092 &iterable, &r))
2093 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002094
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002095 pool = PySequence_Tuple(iterable);
2096 if (pool == NULL)
2097 goto error;
2098 n = PyTuple_GET_SIZE(pool);
2099 if (r < 0) {
2100 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2101 goto error;
2102 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002103
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02002104 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002105 if (indices == NULL) {
2106 PyErr_NoMemory();
2107 goto error;
2108 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002109
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002110 for (i=0 ; i<r ; i++)
2111 indices[i] = i;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002112
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002113 /* create combinationsobject structure */
2114 co = (combinationsobject *)type->tp_alloc(type, 0);
2115 if (co == NULL)
2116 goto error;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002117
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002118 co->pool = pool;
2119 co->indices = indices;
2120 co->result = NULL;
2121 co->r = r;
2122 co->stopped = r > n ? 1 : 0;
2123
2124 return (PyObject *)co;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002125
2126error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002127 if (indices != NULL)
2128 PyMem_Free(indices);
2129 Py_XDECREF(pool);
2130 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002131}
2132
2133static void
2134combinations_dealloc(combinationsobject *co)
2135{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002136 PyObject_GC_UnTrack(co);
2137 Py_XDECREF(co->pool);
2138 Py_XDECREF(co->result);
2139 if (co->indices != NULL)
2140 PyMem_Free(co->indices);
2141 Py_TYPE(co)->tp_free(co);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002142}
2143
2144static int
2145combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2146{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002147 Py_VISIT(co->pool);
2148 Py_VISIT(co->result);
2149 return 0;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002150}
2151
2152static PyObject *
2153combinations_next(combinationsobject *co)
2154{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002155 PyObject *elem;
2156 PyObject *oldelem;
2157 PyObject *pool = co->pool;
2158 Py_ssize_t *indices = co->indices;
2159 PyObject *result = co->result;
2160 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2161 Py_ssize_t r = co->r;
2162 Py_ssize_t i, j, index;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002163
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002164 if (co->stopped)
2165 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002166
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002167 if (result == NULL) {
2168 /* On the first pass, initialize result tuple using the indices */
2169 result = PyTuple_New(r);
2170 if (result == NULL)
2171 goto empty;
2172 co->result = result;
2173 for (i=0; i<r ; i++) {
2174 index = indices[i];
2175 elem = PyTuple_GET_ITEM(pool, index);
2176 Py_INCREF(elem);
2177 PyTuple_SET_ITEM(result, i, elem);
2178 }
2179 } else {
2180 /* Copy the previous result tuple or re-use it if available */
2181 if (Py_REFCNT(result) > 1) {
2182 PyObject *old_result = result;
2183 result = PyTuple_New(r);
2184 if (result == NULL)
2185 goto empty;
2186 co->result = result;
2187 for (i=0; i<r ; i++) {
2188 elem = PyTuple_GET_ITEM(old_result, i);
2189 Py_INCREF(elem);
2190 PyTuple_SET_ITEM(result, i, elem);
2191 }
2192 Py_DECREF(old_result);
2193 }
2194 /* Now, we've got the only copy so we can update it in-place
2195 * CPython's empty tuple is a singleton and cached in
2196 * PyTuple's freelist.
2197 */
2198 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002199
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002200 /* Scan indices right-to-left until finding one that is not
2201 at its maximum (i + n - r). */
2202 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2203 ;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002204
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002205 /* If i is negative, then the indices are all at
2206 their maximum value and we're done. */
2207 if (i < 0)
2208 goto empty;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002209
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002210 /* Increment the current index which we know is not at its
2211 maximum. Then move back to the right setting each index
2212 to its lowest possible value (one higher than the index
2213 to its left -- this maintains the sort order invariant). */
2214 indices[i]++;
2215 for (j=i+1 ; j<r ; j++)
2216 indices[j] = indices[j-1] + 1;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002217
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002218 /* Update the result tuple for the new indices
2219 starting with i, the leftmost index that changed */
2220 for ( ; i<r ; i++) {
2221 index = indices[i];
2222 elem = PyTuple_GET_ITEM(pool, index);
2223 Py_INCREF(elem);
2224 oldelem = PyTuple_GET_ITEM(result, i);
2225 PyTuple_SET_ITEM(result, i, elem);
2226 Py_DECREF(oldelem);
2227 }
2228 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002229
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002230 Py_INCREF(result);
2231 return result;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002232
2233empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002234 co->stopped = 1;
2235 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002236}
2237
2238PyDoc_STRVAR(combinations_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002239"combinations(iterable, r) --> combinations object\n\
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002240\n\
2241Return successive r-length combinations of elements in the iterable.\n\n\
2242combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2243
2244static PyTypeObject combinations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002245 PyVarObject_HEAD_INIT(NULL, 0)
2246 "itertools.combinations", /* tp_name */
2247 sizeof(combinationsobject), /* tp_basicsize */
2248 0, /* tp_itemsize */
2249 /* methods */
2250 (destructor)combinations_dealloc, /* tp_dealloc */
2251 0, /* tp_print */
2252 0, /* tp_getattr */
2253 0, /* tp_setattr */
2254 0, /* tp_compare */
2255 0, /* tp_repr */
2256 0, /* tp_as_number */
2257 0, /* tp_as_sequence */
2258 0, /* tp_as_mapping */
2259 0, /* tp_hash */
2260 0, /* tp_call */
2261 0, /* tp_str */
2262 PyObject_GenericGetAttr, /* tp_getattro */
2263 0, /* tp_setattro */
2264 0, /* tp_as_buffer */
2265 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2266 Py_TPFLAGS_BASETYPE, /* tp_flags */
2267 combinations_doc, /* tp_doc */
2268 (traverseproc)combinations_traverse, /* tp_traverse */
2269 0, /* tp_clear */
2270 0, /* tp_richcompare */
2271 0, /* tp_weaklistoffset */
2272 PyObject_SelfIter, /* tp_iter */
2273 (iternextfunc)combinations_next, /* tp_iternext */
2274 0, /* tp_methods */
2275 0, /* tp_members */
2276 0, /* tp_getset */
2277 0, /* tp_base */
2278 0, /* tp_dict */
2279 0, /* tp_descr_get */
2280 0, /* tp_descr_set */
2281 0, /* tp_dictoffset */
2282 0, /* tp_init */
2283 0, /* tp_alloc */
2284 combinations_new, /* tp_new */
2285 PyObject_GC_Del, /* tp_free */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002286};
2287
2288
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002289/* combinations with replacement object *******************************************/
2290
2291/* Equivalent to:
2292
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002293 def combinations_with_replacement(iterable, r):
2294 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2295 # number items returned: (n+r-1)! / r! / (n-1)!
2296 pool = tuple(iterable)
2297 n = len(pool)
2298 indices = [0] * r
2299 yield tuple(pool[i] for i in indices)
2300 while 1:
2301 for i in reversed(range(r)):
2302 if indices[i] != n - 1:
2303 break
2304 else:
2305 return
2306 indices[i:] = [indices[i] + 1] * (r - i)
2307 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002308
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002309 def combinations_with_replacement2(iterable, r):
2310 'Alternate version that filters from product()'
2311 pool = tuple(iterable)
2312 n = len(pool)
2313 for indices in product(range(n), repeat=r):
2314 if sorted(indices) == list(indices):
2315 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002316*/
2317typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002318 PyObject_HEAD
2319 PyObject *pool; /* input converted to a tuple */
2320 Py_ssize_t *indices; /* one index per result element */
2321 PyObject *result; /* most recently returned result tuple */
2322 Py_ssize_t r; /* size of result tuple */
2323 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002324} cwrobject;
2325
2326static PyTypeObject cwr_type;
2327
2328static PyObject *
2329cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2330{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002331 cwrobject *co;
2332 Py_ssize_t n;
2333 Py_ssize_t r;
2334 PyObject *pool = NULL;
2335 PyObject *iterable = NULL;
2336 Py_ssize_t *indices = NULL;
2337 Py_ssize_t i;
2338 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002339
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002340 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2341 &iterable, &r))
2342 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002343
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002344 pool = PySequence_Tuple(iterable);
2345 if (pool == NULL)
2346 goto error;
2347 n = PyTuple_GET_SIZE(pool);
2348 if (r < 0) {
2349 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2350 goto error;
2351 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002352
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02002353 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002354 if (indices == NULL) {
2355 PyErr_NoMemory();
2356 goto error;
2357 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002358
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002359 for (i=0 ; i<r ; i++)
2360 indices[i] = 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002361
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002362 /* create cwrobject structure */
2363 co = (cwrobject *)type->tp_alloc(type, 0);
2364 if (co == NULL)
2365 goto error;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002366
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002367 co->pool = pool;
2368 co->indices = indices;
2369 co->result = NULL;
2370 co->r = r;
2371 co->stopped = !n && r;
2372
2373 return (PyObject *)co;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002374
2375error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002376 if (indices != NULL)
2377 PyMem_Free(indices);
2378 Py_XDECREF(pool);
2379 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002380}
2381
2382static void
2383cwr_dealloc(cwrobject *co)
2384{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002385 PyObject_GC_UnTrack(co);
2386 Py_XDECREF(co->pool);
2387 Py_XDECREF(co->result);
2388 if (co->indices != NULL)
2389 PyMem_Free(co->indices);
2390 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002391}
2392
2393static int
2394cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2395{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002396 Py_VISIT(co->pool);
2397 Py_VISIT(co->result);
2398 return 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002399}
2400
2401static PyObject *
2402cwr_next(cwrobject *co)
2403{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002404 PyObject *elem;
2405 PyObject *oldelem;
2406 PyObject *pool = co->pool;
2407 Py_ssize_t *indices = co->indices;
2408 PyObject *result = co->result;
2409 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2410 Py_ssize_t r = co->r;
2411 Py_ssize_t i, j, index;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002412
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002413 if (co->stopped)
2414 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002415
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002416 if (result == NULL) {
2417 /* On the first pass, initialize result tuple using the indices */
2418 result = PyTuple_New(r);
2419 if (result == NULL)
2420 goto empty;
2421 co->result = result;
2422 for (i=0; i<r ; i++) {
2423 index = indices[i];
2424 elem = PyTuple_GET_ITEM(pool, index);
2425 Py_INCREF(elem);
2426 PyTuple_SET_ITEM(result, i, elem);
2427 }
2428 } else {
2429 /* Copy the previous result tuple or re-use it if available */
2430 if (Py_REFCNT(result) > 1) {
2431 PyObject *old_result = result;
2432 result = PyTuple_New(r);
2433 if (result == NULL)
2434 goto empty;
2435 co->result = result;
2436 for (i=0; i<r ; i++) {
2437 elem = PyTuple_GET_ITEM(old_result, i);
2438 Py_INCREF(elem);
2439 PyTuple_SET_ITEM(result, i, elem);
2440 }
2441 Py_DECREF(old_result);
2442 }
2443 /* Now, we've got the only copy so we can update it in-place CPython's
2444 empty tuple is a singleton and cached in PyTuple's freelist. */
2445 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002446
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002447 /* Scan indices right-to-left until finding one that is not
2448 * at its maximum (n-1). */
2449 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2450 ;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002451
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002452 /* If i is negative, then the indices are all at
2453 their maximum value and we're done. */
2454 if (i < 0)
2455 goto empty;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002456
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002457 /* Increment the current index which we know is not at its
2458 maximum. Then set all to the right to the same value. */
2459 indices[i]++;
2460 for (j=i+1 ; j<r ; j++)
2461 indices[j] = indices[j-1];
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002462
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002463 /* Update the result tuple for the new indices
2464 starting with i, the leftmost index that changed */
2465 for ( ; i<r ; i++) {
2466 index = indices[i];
2467 elem = PyTuple_GET_ITEM(pool, index);
2468 Py_INCREF(elem);
2469 oldelem = PyTuple_GET_ITEM(result, i);
2470 PyTuple_SET_ITEM(result, i, elem);
2471 Py_DECREF(oldelem);
2472 }
2473 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002474
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002475 Py_INCREF(result);
2476 return result;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002477
2478empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002479 co->stopped = 1;
2480 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002481}
2482
2483PyDoc_STRVAR(cwr_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002484"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002485\n\
2486Return successive r-length combinations of elements in the iterable\n\
2487allowing individual elements to have successive repeats.\n\
2488combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2489
2490static PyTypeObject cwr_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002491 PyVarObject_HEAD_INIT(NULL, 0)
2492 "itertools.combinations_with_replacement", /* tp_name */
2493 sizeof(cwrobject), /* tp_basicsize */
2494 0, /* tp_itemsize */
2495 /* methods */
2496 (destructor)cwr_dealloc, /* tp_dealloc */
2497 0, /* tp_print */
2498 0, /* tp_getattr */
2499 0, /* tp_setattr */
2500 0, /* tp_compare */
2501 0, /* tp_repr */
2502 0, /* tp_as_number */
2503 0, /* tp_as_sequence */
2504 0, /* tp_as_mapping */
2505 0, /* tp_hash */
2506 0, /* tp_call */
2507 0, /* tp_str */
2508 PyObject_GenericGetAttr, /* tp_getattro */
2509 0, /* tp_setattro */
2510 0, /* tp_as_buffer */
2511 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2512 Py_TPFLAGS_BASETYPE, /* tp_flags */
2513 cwr_doc, /* tp_doc */
2514 (traverseproc)cwr_traverse, /* tp_traverse */
2515 0, /* tp_clear */
2516 0, /* tp_richcompare */
2517 0, /* tp_weaklistoffset */
2518 PyObject_SelfIter, /* tp_iter */
2519 (iternextfunc)cwr_next, /* tp_iternext */
2520 0, /* tp_methods */
2521 0, /* tp_members */
2522 0, /* tp_getset */
2523 0, /* tp_base */
2524 0, /* tp_dict */
2525 0, /* tp_descr_get */
2526 0, /* tp_descr_set */
2527 0, /* tp_dictoffset */
2528 0, /* tp_init */
2529 0, /* tp_alloc */
2530 cwr_new, /* tp_new */
2531 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002532};
2533
2534
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002535/* permutations object ************************************************************
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002536
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002537def permutations(iterable, r=None):
2538 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2539 pool = tuple(iterable)
2540 n = len(pool)
2541 r = n if r is None else r
2542 indices = range(n)
2543 cycles = range(n-r+1, n+1)[::-1]
2544 yield tuple(pool[i] for i in indices[:r])
2545 while n:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002546 for i in reversed(range(r)):
2547 cycles[i] -= 1
2548 if cycles[i] == 0:
2549 indices[i:] = indices[i+1:] + indices[i:i+1]
2550 cycles[i] = n - i
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002551 else:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002552 j = cycles[i]
2553 indices[i], indices[-j] = indices[-j], indices[i]
2554 yield tuple(pool[i] for i in indices[:r])
2555 break
2556 else:
2557 return
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002558*/
2559
2560typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002561 PyObject_HEAD
2562 PyObject *pool; /* input converted to a tuple */
2563 Py_ssize_t *indices; /* one index per element in the pool */
2564 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2565 PyObject *result; /* most recently returned result tuple */
2566 Py_ssize_t r; /* size of result tuple */
2567 int stopped; /* set to 1 when the permutations iterator is exhausted */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002568} permutationsobject;
2569
2570static PyTypeObject permutations_type;
2571
2572static PyObject *
2573permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2574{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002575 permutationsobject *po;
2576 Py_ssize_t n;
2577 Py_ssize_t r;
2578 PyObject *robj = Py_None;
2579 PyObject *pool = NULL;
2580 PyObject *iterable = NULL;
2581 Py_ssize_t *indices = NULL;
2582 Py_ssize_t *cycles = NULL;
2583 Py_ssize_t i;
2584 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002585
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002586 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2587 &iterable, &robj))
2588 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002589
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002590 pool = PySequence_Tuple(iterable);
2591 if (pool == NULL)
2592 goto error;
2593 n = PyTuple_GET_SIZE(pool);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002594
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002595 r = n;
2596 if (robj != Py_None) {
2597 r = PyInt_AsSsize_t(robj);
2598 if (r == -1 && PyErr_Occurred())
2599 goto error;
2600 }
2601 if (r < 0) {
2602 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2603 goto error;
2604 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002605
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02002606 indices = PyMem_New(Py_ssize_t, n);
2607 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002608 if (indices == NULL || cycles == NULL) {
2609 PyErr_NoMemory();
2610 goto error;
2611 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002612
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002613 for (i=0 ; i<n ; i++)
2614 indices[i] = i;
2615 for (i=0 ; i<r ; i++)
2616 cycles[i] = n - i;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002617
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002618 /* create permutationsobject structure */
2619 po = (permutationsobject *)type->tp_alloc(type, 0);
2620 if (po == NULL)
2621 goto error;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002622
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002623 po->pool = pool;
2624 po->indices = indices;
2625 po->cycles = cycles;
2626 po->result = NULL;
2627 po->r = r;
2628 po->stopped = r > n ? 1 : 0;
2629
2630 return (PyObject *)po;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002631
2632error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002633 if (indices != NULL)
2634 PyMem_Free(indices);
2635 if (cycles != NULL)
2636 PyMem_Free(cycles);
2637 Py_XDECREF(pool);
2638 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002639}
2640
2641static void
2642permutations_dealloc(permutationsobject *po)
2643{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002644 PyObject_GC_UnTrack(po);
2645 Py_XDECREF(po->pool);
2646 Py_XDECREF(po->result);
2647 PyMem_Free(po->indices);
2648 PyMem_Free(po->cycles);
2649 Py_TYPE(po)->tp_free(po);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002650}
2651
2652static int
2653permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2654{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002655 Py_VISIT(po->pool);
2656 Py_VISIT(po->result);
2657 return 0;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002658}
2659
2660static PyObject *
2661permutations_next(permutationsobject *po)
2662{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002663 PyObject *elem;
2664 PyObject *oldelem;
2665 PyObject *pool = po->pool;
2666 Py_ssize_t *indices = po->indices;
2667 Py_ssize_t *cycles = po->cycles;
2668 PyObject *result = po->result;
2669 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2670 Py_ssize_t r = po->r;
2671 Py_ssize_t i, j, k, index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002672
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002673 if (po->stopped)
2674 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002675
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002676 if (result == NULL) {
2677 /* On the first pass, initialize result tuple using the indices */
2678 result = PyTuple_New(r);
2679 if (result == NULL)
2680 goto empty;
2681 po->result = result;
2682 for (i=0; i<r ; i++) {
2683 index = indices[i];
2684 elem = PyTuple_GET_ITEM(pool, index);
2685 Py_INCREF(elem);
2686 PyTuple_SET_ITEM(result, i, elem);
2687 }
2688 } else {
2689 if (n == 0)
2690 goto empty;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002691
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002692 /* Copy the previous result tuple or re-use it if available */
2693 if (Py_REFCNT(result) > 1) {
2694 PyObject *old_result = result;
2695 result = PyTuple_New(r);
2696 if (result == NULL)
2697 goto empty;
2698 po->result = result;
2699 for (i=0; i<r ; i++) {
2700 elem = PyTuple_GET_ITEM(old_result, i);
2701 Py_INCREF(elem);
2702 PyTuple_SET_ITEM(result, i, elem);
2703 }
2704 Py_DECREF(old_result);
2705 }
2706 /* Now, we've got the only copy so we can update it in-place */
2707 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002708
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002709 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2710 for (i=r-1 ; i>=0 ; i--) {
2711 cycles[i] -= 1;
2712 if (cycles[i] == 0) {
2713 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2714 index = indices[i];
2715 for (j=i ; j<n-1 ; j++)
2716 indices[j] = indices[j+1];
2717 indices[n-1] = index;
2718 cycles[i] = n - i;
2719 } else {
2720 j = cycles[i];
2721 index = indices[i];
2722 indices[i] = indices[n-j];
2723 indices[n-j] = index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002724
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002725 for (k=i; k<r ; k++) {
2726 /* start with i, the leftmost element that changed */
2727 /* yield tuple(pool[k] for k in indices[:r]) */
2728 index = indices[k];
2729 elem = PyTuple_GET_ITEM(pool, index);
2730 Py_INCREF(elem);
2731 oldelem = PyTuple_GET_ITEM(result, k);
2732 PyTuple_SET_ITEM(result, k, elem);
2733 Py_DECREF(oldelem);
2734 }
2735 break;
2736 }
2737 }
2738 /* If i is negative, then the cycles have all
2739 rolled-over and we're done. */
2740 if (i < 0)
2741 goto empty;
2742 }
2743 Py_INCREF(result);
2744 return result;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002745
2746empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002747 po->stopped = 1;
2748 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002749}
2750
2751PyDoc_STRVAR(permutations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +00002752"permutations(iterable[, r]) --> permutations object\n\
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002753\n\
2754Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl8550d692008-06-10 12:46:39 +00002755permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002756
2757static PyTypeObject permutations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002758 PyVarObject_HEAD_INIT(NULL, 0)
2759 "itertools.permutations", /* tp_name */
2760 sizeof(permutationsobject), /* tp_basicsize */
2761 0, /* tp_itemsize */
2762 /* methods */
2763 (destructor)permutations_dealloc, /* tp_dealloc */
2764 0, /* tp_print */
2765 0, /* tp_getattr */
2766 0, /* tp_setattr */
2767 0, /* tp_compare */
2768 0, /* tp_repr */
2769 0, /* tp_as_number */
2770 0, /* tp_as_sequence */
2771 0, /* tp_as_mapping */
2772 0, /* tp_hash */
2773 0, /* tp_call */
2774 0, /* tp_str */
2775 PyObject_GenericGetAttr, /* tp_getattro */
2776 0, /* tp_setattro */
2777 0, /* tp_as_buffer */
2778 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2779 Py_TPFLAGS_BASETYPE, /* tp_flags */
2780 permutations_doc, /* tp_doc */
2781 (traverseproc)permutations_traverse, /* tp_traverse */
2782 0, /* tp_clear */
2783 0, /* tp_richcompare */
2784 0, /* tp_weaklistoffset */
2785 PyObject_SelfIter, /* tp_iter */
2786 (iternextfunc)permutations_next, /* tp_iternext */
2787 0, /* tp_methods */
2788 0, /* tp_members */
2789 0, /* tp_getset */
2790 0, /* tp_base */
2791 0, /* tp_dict */
2792 0, /* tp_descr_get */
2793 0, /* tp_descr_set */
2794 0, /* tp_dictoffset */
2795 0, /* tp_init */
2796 0, /* tp_alloc */
2797 permutations_new, /* tp_new */
2798 PyObject_GC_Del, /* tp_free */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002799};
2800
2801
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002802/* compress object ************************************************************/
2803
2804/* Equivalent to:
2805
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002806 def compress(data, selectors):
2807 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2808 return (d for d, s in izip(data, selectors) if s)
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002809*/
2810
2811typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002812 PyObject_HEAD
2813 PyObject *data;
2814 PyObject *selectors;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002815} compressobject;
2816
2817static PyTypeObject compress_type;
2818
2819static PyObject *
2820compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2821{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002822 PyObject *seq1, *seq2;
2823 PyObject *data=NULL, *selectors=NULL;
2824 compressobject *lz;
2825 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002826
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002827 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2828 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002829
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002830 data = PyObject_GetIter(seq1);
2831 if (data == NULL)
2832 goto fail;
2833 selectors = PyObject_GetIter(seq2);
2834 if (selectors == NULL)
2835 goto fail;
2836
2837 /* create compressobject structure */
2838 lz = (compressobject *)type->tp_alloc(type, 0);
2839 if (lz == NULL)
2840 goto fail;
2841 lz->data = data;
2842 lz->selectors = selectors;
2843 return (PyObject *)lz;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002844
2845fail:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002846 Py_XDECREF(data);
2847 Py_XDECREF(selectors);
2848 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002849}
2850
2851static void
2852compress_dealloc(compressobject *lz)
2853{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002854 PyObject_GC_UnTrack(lz);
2855 Py_XDECREF(lz->data);
2856 Py_XDECREF(lz->selectors);
2857 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002858}
2859
2860static int
2861compress_traverse(compressobject *lz, visitproc visit, void *arg)
2862{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002863 Py_VISIT(lz->data);
2864 Py_VISIT(lz->selectors);
2865 return 0;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002866}
2867
2868static PyObject *
2869compress_next(compressobject *lz)
2870{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002871 PyObject *data = lz->data, *selectors = lz->selectors;
2872 PyObject *datum, *selector;
2873 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2874 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2875 int ok;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002876
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002877 while (1) {
2878 /* Steps: get datum, get selector, evaluate selector.
2879 Order is important (to match the pure python version
2880 in terms of which input gets a chance to raise an
2881 exception first).
2882 */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002883
Serhiy Storchakabb845652013-04-06 22:04:10 +03002884 datum = datanext(data);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03002885 if (datum == NULL)
Serhiy Storchakabb845652013-04-06 22:04:10 +03002886 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002887
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002888 selector = selectornext(selectors);
2889 if (selector == NULL) {
2890 Py_DECREF(datum);
2891 return NULL;
2892 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002893
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002894 ok = PyObject_IsTrue(selector);
2895 Py_DECREF(selector);
2896 if (ok == 1)
2897 return datum;
2898 Py_DECREF(datum);
2899 if (ok == -1)
2900 return NULL;
2901 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002902}
2903
2904PyDoc_STRVAR(compress_doc,
Raymond Hettinger2e2909f2009-02-19 02:15:14 +00002905"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002906\n\
2907Return data elements corresponding to true selector elements.\n\
2908Forms a shorter iterator from selected data elements using the\n\
2909selectors to choose the data elements.");
2910
2911static PyTypeObject compress_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002912 PyVarObject_HEAD_INIT(NULL, 0)
2913 "itertools.compress", /* tp_name */
2914 sizeof(compressobject), /* tp_basicsize */
2915 0, /* tp_itemsize */
2916 /* methods */
2917 (destructor)compress_dealloc, /* tp_dealloc */
2918 0, /* tp_print */
2919 0, /* tp_getattr */
2920 0, /* tp_setattr */
2921 0, /* tp_compare */
2922 0, /* tp_repr */
2923 0, /* tp_as_number */
2924 0, /* tp_as_sequence */
2925 0, /* tp_as_mapping */
2926 0, /* tp_hash */
2927 0, /* tp_call */
2928 0, /* tp_str */
2929 PyObject_GenericGetAttr, /* tp_getattro */
2930 0, /* tp_setattro */
2931 0, /* tp_as_buffer */
2932 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2933 Py_TPFLAGS_BASETYPE, /* tp_flags */
2934 compress_doc, /* tp_doc */
2935 (traverseproc)compress_traverse, /* tp_traverse */
2936 0, /* tp_clear */
2937 0, /* tp_richcompare */
2938 0, /* tp_weaklistoffset */
2939 PyObject_SelfIter, /* tp_iter */
2940 (iternextfunc)compress_next, /* tp_iternext */
2941 0, /* tp_methods */
2942 0, /* tp_members */
2943 0, /* tp_getset */
2944 0, /* tp_base */
2945 0, /* tp_dict */
2946 0, /* tp_descr_get */
2947 0, /* tp_descr_set */
2948 0, /* tp_dictoffset */
2949 0, /* tp_init */
2950 0, /* tp_alloc */
2951 compress_new, /* tp_new */
2952 PyObject_GC_Del, /* tp_free */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002953};
2954
2955
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002956/* ifilter object ************************************************************/
2957
2958typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002959 PyObject_HEAD
2960 PyObject *func;
2961 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002962} ifilterobject;
2963
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002964static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002965
2966static PyObject *
2967ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2968{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002969 PyObject *func, *seq;
2970 PyObject *it;
2971 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002972
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002973 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
2974 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002975
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002976 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
2977 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002978
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002979 /* Get iterator. */
2980 it = PyObject_GetIter(seq);
2981 if (it == NULL)
2982 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002983
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002984 /* create ifilterobject structure */
2985 lz = (ifilterobject *)type->tp_alloc(type, 0);
2986 if (lz == NULL) {
2987 Py_DECREF(it);
2988 return NULL;
2989 }
2990 Py_INCREF(func);
2991 lz->func = func;
2992 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002993
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002994 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002995}
2996
2997static void
2998ifilter_dealloc(ifilterobject *lz)
2999{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003000 PyObject_GC_UnTrack(lz);
3001 Py_XDECREF(lz->func);
3002 Py_XDECREF(lz->it);
3003 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003004}
3005
3006static int
3007ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
3008{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003009 Py_VISIT(lz->it);
3010 Py_VISIT(lz->func);
3011 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003012}
3013
3014static PyObject *
3015ifilter_next(ifilterobject *lz)
3016{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003017 PyObject *item;
3018 PyObject *it = lz->it;
3019 long ok;
3020 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003021
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003022 iternext = *Py_TYPE(it)->tp_iternext;
3023 for (;;) {
3024 item = iternext(it);
3025 if (item == NULL)
3026 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003027
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003028 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3029 ok = PyObject_IsTrue(item);
3030 } else {
3031 PyObject *good;
3032 good = PyObject_CallFunctionObjArgs(lz->func,
3033 item, NULL);
3034 if (good == NULL) {
3035 Py_DECREF(item);
3036 return NULL;
3037 }
3038 ok = PyObject_IsTrue(good);
3039 Py_DECREF(good);
3040 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003041 if (ok > 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003042 return item;
3043 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003044 if (ok < 0)
3045 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003046 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003047}
3048
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003049PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003050"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003051\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003052Return those items of sequence for which function(item) is true.\n\
3053If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003054
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003055static PyTypeObject ifilter_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003056 PyVarObject_HEAD_INIT(NULL, 0)
3057 "itertools.ifilter", /* tp_name */
3058 sizeof(ifilterobject), /* tp_basicsize */
3059 0, /* tp_itemsize */
3060 /* methods */
3061 (destructor)ifilter_dealloc, /* tp_dealloc */
3062 0, /* tp_print */
3063 0, /* tp_getattr */
3064 0, /* tp_setattr */
3065 0, /* tp_compare */
3066 0, /* tp_repr */
3067 0, /* tp_as_number */
3068 0, /* tp_as_sequence */
3069 0, /* tp_as_mapping */
3070 0, /* tp_hash */
3071 0, /* tp_call */
3072 0, /* tp_str */
3073 PyObject_GenericGetAttr, /* tp_getattro */
3074 0, /* tp_setattro */
3075 0, /* tp_as_buffer */
3076 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3077 Py_TPFLAGS_BASETYPE, /* tp_flags */
3078 ifilter_doc, /* tp_doc */
3079 (traverseproc)ifilter_traverse, /* tp_traverse */
3080 0, /* tp_clear */
3081 0, /* tp_richcompare */
3082 0, /* tp_weaklistoffset */
3083 PyObject_SelfIter, /* tp_iter */
3084 (iternextfunc)ifilter_next, /* tp_iternext */
3085 0, /* tp_methods */
3086 0, /* tp_members */
3087 0, /* tp_getset */
3088 0, /* tp_base */
3089 0, /* tp_dict */
3090 0, /* tp_descr_get */
3091 0, /* tp_descr_set */
3092 0, /* tp_dictoffset */
3093 0, /* tp_init */
3094 0, /* tp_alloc */
3095 ifilter_new, /* tp_new */
3096 PyObject_GC_Del, /* tp_free */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003097};
3098
3099
3100/* ifilterfalse object ************************************************************/
3101
3102typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003103 PyObject_HEAD
3104 PyObject *func;
3105 PyObject *it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003106} ifilterfalseobject;
3107
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003108static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003109
3110static PyObject *
3111ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3112{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003113 PyObject *func, *seq;
3114 PyObject *it;
3115 ifilterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003116
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003117 if (type == &ifilterfalse_type &&
3118 !_PyArg_NoKeywords("ifilterfalse()", kwds))
3119 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003120
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003121 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
3122 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003123
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003124 /* Get iterator. */
3125 it = PyObject_GetIter(seq);
3126 if (it == NULL)
3127 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003128
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003129 /* create ifilterfalseobject structure */
3130 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
3131 if (lz == NULL) {
3132 Py_DECREF(it);
3133 return NULL;
3134 }
3135 Py_INCREF(func);
3136 lz->func = func;
3137 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003138
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003139 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003140}
3141
3142static void
3143ifilterfalse_dealloc(ifilterfalseobject *lz)
3144{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003145 PyObject_GC_UnTrack(lz);
3146 Py_XDECREF(lz->func);
3147 Py_XDECREF(lz->it);
3148 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003149}
3150
3151static int
3152ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
3153{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003154 Py_VISIT(lz->it);
3155 Py_VISIT(lz->func);
3156 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003157}
3158
3159static PyObject *
3160ifilterfalse_next(ifilterfalseobject *lz)
3161{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003162 PyObject *item;
3163 PyObject *it = lz->it;
3164 long ok;
3165 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003166
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003167 iternext = *Py_TYPE(it)->tp_iternext;
3168 for (;;) {
3169 item = iternext(it);
3170 if (item == NULL)
3171 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003172
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003173 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3174 ok = PyObject_IsTrue(item);
3175 } else {
3176 PyObject *good;
3177 good = PyObject_CallFunctionObjArgs(lz->func,
3178 item, NULL);
3179 if (good == NULL) {
3180 Py_DECREF(item);
3181 return NULL;
3182 }
3183 ok = PyObject_IsTrue(good);
3184 Py_DECREF(good);
3185 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003186 if (ok == 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003187 return item;
3188 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003189 if (ok < 0)
3190 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003191 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003192}
3193
Raymond Hettinger60eca932003-02-09 06:40:58 +00003194PyDoc_STRVAR(ifilterfalse_doc,
3195"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
3196\n\
3197Return those items of sequence for which function(item) is false.\n\
3198If function is None, return the items that are false.");
3199
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003200static PyTypeObject ifilterfalse_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003201 PyVarObject_HEAD_INIT(NULL, 0)
3202 "itertools.ifilterfalse", /* tp_name */
3203 sizeof(ifilterfalseobject), /* tp_basicsize */
3204 0, /* tp_itemsize */
3205 /* methods */
3206 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
3207 0, /* tp_print */
3208 0, /* tp_getattr */
3209 0, /* tp_setattr */
3210 0, /* tp_compare */
3211 0, /* tp_repr */
3212 0, /* tp_as_number */
3213 0, /* tp_as_sequence */
3214 0, /* tp_as_mapping */
3215 0, /* tp_hash */
3216 0, /* tp_call */
3217 0, /* tp_str */
3218 PyObject_GenericGetAttr, /* tp_getattro */
3219 0, /* tp_setattro */
3220 0, /* tp_as_buffer */
3221 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3222 Py_TPFLAGS_BASETYPE, /* tp_flags */
3223 ifilterfalse_doc, /* tp_doc */
3224 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
3225 0, /* tp_clear */
3226 0, /* tp_richcompare */
3227 0, /* tp_weaklistoffset */
3228 PyObject_SelfIter, /* tp_iter */
3229 (iternextfunc)ifilterfalse_next, /* tp_iternext */
3230 0, /* tp_methods */
3231 0, /* tp_members */
3232 0, /* tp_getset */
3233 0, /* tp_base */
3234 0, /* tp_dict */
3235 0, /* tp_descr_get */
3236 0, /* tp_descr_set */
3237 0, /* tp_dictoffset */
3238 0, /* tp_init */
3239 0, /* tp_alloc */
3240 ifilterfalse_new, /* tp_new */
3241 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003242};
3243
3244
3245/* count object ************************************************************/
3246
3247typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003248 PyObject_HEAD
3249 Py_ssize_t cnt;
3250 PyObject *long_cnt;
3251 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003252} countobject;
3253
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003254/* Counting logic and invariants:
3255
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003256fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003257
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003258 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3259 Advances with: cnt += 1
3260 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003261
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003262slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003263
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003264 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3265 All counting is done with python objects (no overflows or underflows).
3266 Advances with: long_cnt += long_step
3267 Step may be zero -- effectively a slow version of repeat(cnt).
3268 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003269*/
3270
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003271static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003272
3273static PyObject *
3274count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3275{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003276 countobject *lz;
3277 int slow_mode = 0;
3278 Py_ssize_t cnt = 0;
3279 PyObject *long_cnt = NULL;
3280 PyObject *long_step = NULL;
3281 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003282
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003283 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3284 kwlist, &long_cnt, &long_step))
3285 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003286
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003287 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3288 (long_step != NULL && !PyNumber_Check(long_step))) {
3289 PyErr_SetString(PyExc_TypeError, "a number is required");
3290 return NULL;
3291 }
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003292
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003293 if (long_cnt != NULL) {
3294 cnt = PyInt_AsSsize_t(long_cnt);
3295 if ((cnt == -1 && PyErr_Occurred()) || !PyInt_Check(long_cnt)) {
3296 PyErr_Clear();
3297 slow_mode = 1;
3298 }
3299 Py_INCREF(long_cnt);
3300 } else {
3301 cnt = 0;
3302 long_cnt = PyInt_FromLong(0);
3303 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003304
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003305 /* If not specified, step defaults to 1 */
3306 if (long_step == NULL) {
3307 long_step = PyInt_FromLong(1);
3308 if (long_step == NULL) {
3309 Py_DECREF(long_cnt);
3310 return NULL;
3311 }
3312 } else
3313 Py_INCREF(long_step);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003314
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003315 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003316
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003317 /* Fast mode only works when the step is 1 */
3318 if (!PyInt_Check(long_step) ||
3319 PyInt_AS_LONG(long_step) != 1) {
3320 slow_mode = 1;
3321 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003322
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003323 if (slow_mode)
3324 cnt = PY_SSIZE_T_MAX;
3325 else
3326 Py_CLEAR(long_cnt);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003327
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003328 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3329 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3330 assert(slow_mode ||
3331 (PyInt_Check(long_step) && PyInt_AS_LONG(long_step) == 1));
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003332
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003333 /* create countobject structure */
3334 lz = (countobject *)type->tp_alloc(type, 0);
3335 if (lz == NULL) {
3336 Py_XDECREF(long_cnt);
3337 return NULL;
3338 }
3339 lz->cnt = cnt;
3340 lz->long_cnt = long_cnt;
3341 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003342
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003343 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003344}
3345
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003346static void
3347count_dealloc(countobject *lz)
3348{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003349 PyObject_GC_UnTrack(lz);
3350 Py_XDECREF(lz->long_cnt);
3351 Py_XDECREF(lz->long_step);
3352 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003353}
3354
Benjamin Peterson062a7c32009-02-16 21:07:52 +00003355static int
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003356count_traverse(countobject *lz, visitproc visit, void *arg)
3357{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003358 Py_VISIT(lz->long_cnt);
3359 Py_VISIT(lz->long_step);
3360 return 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003361}
3362
3363static PyObject *
3364count_nextlong(countobject *lz)
3365{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003366 PyObject *long_cnt;
3367 PyObject *stepped_up;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003368
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003369 long_cnt = lz->long_cnt;
3370 if (long_cnt == NULL) {
3371 /* Switch to slow_mode */
3372 long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
3373 if (long_cnt == NULL)
3374 return NULL;
3375 }
3376 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003377
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003378 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3379 if (stepped_up == NULL)
3380 return NULL;
3381 lz->long_cnt = stepped_up;
3382 return long_cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003383}
3384
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003385static PyObject *
3386count_next(countobject *lz)
3387{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003388 if (lz->cnt == PY_SSIZE_T_MAX)
3389 return count_nextlong(lz);
3390 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003391}
3392
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003393static PyObject *
3394count_repr(countobject *lz)
3395{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003396 PyObject *cnt_repr, *step_repr = NULL;
3397 PyObject *result = NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003398
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003399 if (lz->cnt != PY_SSIZE_T_MAX)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003400 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003401
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003402 cnt_repr = PyObject_Repr(lz->long_cnt);
3403 if (cnt_repr == NULL)
3404 return NULL;
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003405
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003406 if (PyInt_Check(lz->long_step) && PyInt_AS_LONG(lz->long_step) == 1) {
3407 /* Don't display step when it is an integer equal to 1 */
3408 result = PyString_FromFormat("count(%s)",
3409 PyString_AS_STRING(cnt_repr));
3410 } else {
3411 step_repr = PyObject_Repr(lz->long_step);
3412 if (step_repr != NULL)
3413 result = PyString_FromFormat("count(%s, %s)",
3414 PyString_AS_STRING(cnt_repr),
3415 PyString_AS_STRING(step_repr));
3416 }
3417 Py_DECREF(cnt_repr);
3418 Py_XDECREF(step_repr);
3419 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003420}
3421
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003422static PyObject *
3423count_reduce(countobject *lz)
3424{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003425 if (lz->cnt == PY_SSIZE_T_MAX)
3426 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3427 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003428}
3429
3430PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3431
3432static PyMethodDef count_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003433 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3434 count_reduce_doc},
3435 {NULL, NULL} /* sentinel */
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003436};
3437
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003438PyDoc_STRVAR(count_doc,
Ezio Melottifdc1e0d2010-06-11 02:21:25 +00003439 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003440\n\
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003441Return a count object whose .next() method returns consecutive values.\n\
3442Equivalent to:\n\n\
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003443 def count(firstval=0, step=1):\n\
Eli Benderskye6802db2013-09-03 06:41:58 -07003444 x = firstval\n\
3445 while 1:\n\
3446 yield x\n\
3447 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003448
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003449static PyTypeObject count_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003450 PyVarObject_HEAD_INIT(NULL, 0)
3451 "itertools.count", /* tp_name */
3452 sizeof(countobject), /* tp_basicsize */
3453 0, /* tp_itemsize */
3454 /* methods */
3455 (destructor)count_dealloc, /* tp_dealloc */
3456 0, /* tp_print */
3457 0, /* tp_getattr */
3458 0, /* tp_setattr */
3459 0, /* tp_compare */
3460 (reprfunc)count_repr, /* tp_repr */
3461 0, /* tp_as_number */
3462 0, /* tp_as_sequence */
3463 0, /* tp_as_mapping */
3464 0, /* tp_hash */
3465 0, /* tp_call */
3466 0, /* tp_str */
3467 PyObject_GenericGetAttr, /* tp_getattro */
3468 0, /* tp_setattro */
3469 0, /* tp_as_buffer */
3470 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3471 Py_TPFLAGS_BASETYPE, /* tp_flags */
3472 count_doc, /* tp_doc */
3473 (traverseproc)count_traverse, /* tp_traverse */
3474 0, /* tp_clear */
3475 0, /* tp_richcompare */
3476 0, /* tp_weaklistoffset */
3477 PyObject_SelfIter, /* tp_iter */
3478 (iternextfunc)count_next, /* tp_iternext */
3479 count_methods, /* tp_methods */
3480 0, /* tp_members */
3481 0, /* tp_getset */
3482 0, /* tp_base */
3483 0, /* tp_dict */
3484 0, /* tp_descr_get */
3485 0, /* tp_descr_set */
3486 0, /* tp_dictoffset */
3487 0, /* tp_init */
3488 0, /* tp_alloc */
3489 count_new, /* tp_new */
3490 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003491};
3492
3493
3494/* izip object ************************************************************/
3495
3496#include "Python.h"
3497
3498typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003499 PyObject_HEAD
3500 Py_ssize_t tuplesize;
3501 PyObject *ittuple; /* tuple of iterators */
3502 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003503} izipobject;
3504
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003505static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003506
3507static PyObject *
3508izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3509{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003510 izipobject *lz;
3511 Py_ssize_t i;
3512 PyObject *ittuple; /* tuple of iterators */
3513 PyObject *result;
3514 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003515
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003516 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
3517 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003518
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003519 /* args must be a tuple */
3520 assert(PyTuple_Check(args));
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003521
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003522 /* obtain iterators */
3523 ittuple = PyTuple_New(tuplesize);
3524 if (ittuple == NULL)
3525 return NULL;
3526 for (i=0; i < tuplesize; ++i) {
3527 PyObject *item = PyTuple_GET_ITEM(args, i);
3528 PyObject *it = PyObject_GetIter(item);
3529 if (it == NULL) {
3530 if (PyErr_ExceptionMatches(PyExc_TypeError))
3531 PyErr_Format(PyExc_TypeError,
3532 "izip argument #%zd must support iteration",
3533 i+1);
3534 Py_DECREF(ittuple);
3535 return NULL;
3536 }
3537 PyTuple_SET_ITEM(ittuple, i, it);
3538 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003539
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003540 /* create a result holder */
3541 result = PyTuple_New(tuplesize);
3542 if (result == NULL) {
3543 Py_DECREF(ittuple);
3544 return NULL;
3545 }
3546 for (i=0 ; i < tuplesize ; i++) {
3547 Py_INCREF(Py_None);
3548 PyTuple_SET_ITEM(result, i, Py_None);
3549 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003550
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003551 /* create izipobject structure */
3552 lz = (izipobject *)type->tp_alloc(type, 0);
3553 if (lz == NULL) {
3554 Py_DECREF(ittuple);
3555 Py_DECREF(result);
3556 return NULL;
3557 }
3558 lz->ittuple = ittuple;
3559 lz->tuplesize = tuplesize;
3560 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003561
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003562 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003563}
3564
3565static void
3566izip_dealloc(izipobject *lz)
3567{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003568 PyObject_GC_UnTrack(lz);
3569 Py_XDECREF(lz->ittuple);
3570 Py_XDECREF(lz->result);
3571 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003572}
3573
3574static int
3575izip_traverse(izipobject *lz, visitproc visit, void *arg)
3576{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003577 Py_VISIT(lz->ittuple);
3578 Py_VISIT(lz->result);
3579 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003580}
3581
3582static PyObject *
3583izip_next(izipobject *lz)
3584{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003585 Py_ssize_t i;
3586 Py_ssize_t tuplesize = lz->tuplesize;
3587 PyObject *result = lz->result;
3588 PyObject *it;
3589 PyObject *item;
3590 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003591
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003592 if (tuplesize == 0)
3593 return NULL;
3594 if (Py_REFCNT(result) == 1) {
3595 Py_INCREF(result);
3596 for (i=0 ; i < tuplesize ; i++) {
3597 it = PyTuple_GET_ITEM(lz->ittuple, i);
3598 item = (*Py_TYPE(it)->tp_iternext)(it);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003599 if (item == NULL) {
3600 Py_DECREF(result);
3601 return NULL;
3602 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003603 olditem = PyTuple_GET_ITEM(result, i);
3604 PyTuple_SET_ITEM(result, i, item);
3605 Py_DECREF(olditem);
3606 }
3607 } else {
3608 result = PyTuple_New(tuplesize);
3609 if (result == NULL)
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003610 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003611 for (i=0 ; i < tuplesize ; i++) {
3612 it = PyTuple_GET_ITEM(lz->ittuple, i);
3613 item = (*Py_TYPE(it)->tp_iternext)(it);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003614 if (item == NULL) {
3615 Py_DECREF(result);
3616 return NULL;
3617 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003618 PyTuple_SET_ITEM(result, i, item);
3619 }
3620 }
3621 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003622}
3623
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003624PyDoc_STRVAR(izip_doc,
3625"izip(iter1 [,iter2 [...]]) --> izip object\n\
3626\n\
3627Return a izip object whose .next() method returns a tuple where\n\
3628the i-th element comes from the i-th iterable argument. The .next()\n\
3629method continues until the shortest iterable in the argument sequence\n\
3630is exhausted and then it raises StopIteration. Works like the zip()\n\
3631function but consumes less memory by returning an iterator instead of\n\
3632a list.");
3633
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003634static PyTypeObject izip_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003635 PyVarObject_HEAD_INIT(NULL, 0)
3636 "itertools.izip", /* tp_name */
3637 sizeof(izipobject), /* tp_basicsize */
3638 0, /* tp_itemsize */
3639 /* methods */
3640 (destructor)izip_dealloc, /* tp_dealloc */
3641 0, /* tp_print */
3642 0, /* tp_getattr */
3643 0, /* tp_setattr */
3644 0, /* tp_compare */
3645 0, /* tp_repr */
3646 0, /* tp_as_number */
3647 0, /* tp_as_sequence */
3648 0, /* tp_as_mapping */
3649 0, /* tp_hash */
3650 0, /* tp_call */
3651 0, /* tp_str */
3652 PyObject_GenericGetAttr, /* tp_getattro */
3653 0, /* tp_setattro */
3654 0, /* tp_as_buffer */
3655 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3656 Py_TPFLAGS_BASETYPE, /* tp_flags */
3657 izip_doc, /* tp_doc */
3658 (traverseproc)izip_traverse, /* tp_traverse */
3659 0, /* tp_clear */
3660 0, /* tp_richcompare */
3661 0, /* tp_weaklistoffset */
3662 PyObject_SelfIter, /* tp_iter */
3663 (iternextfunc)izip_next, /* tp_iternext */
3664 0, /* tp_methods */
3665 0, /* tp_members */
3666 0, /* tp_getset */
3667 0, /* tp_base */
3668 0, /* tp_dict */
3669 0, /* tp_descr_get */
3670 0, /* tp_descr_set */
3671 0, /* tp_dictoffset */
3672 0, /* tp_init */
3673 0, /* tp_alloc */
3674 izip_new, /* tp_new */
3675 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003676};
3677
3678
3679/* repeat object ************************************************************/
3680
3681typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003682 PyObject_HEAD
3683 PyObject *element;
3684 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003685} repeatobject;
3686
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003687static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003688
3689static PyObject *
3690repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3691{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003692 repeatobject *ro;
3693 PyObject *element;
Raymond Hettinger58ad2452014-06-24 21:53:45 -07003694 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003695 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003696
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003697 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3698 &element, &cnt))
3699 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003700
Raymond Hettinger58ad2452014-06-24 21:53:45 -07003701 if (kwds != NULL)
3702 n_kwds = PyDict_Size(kwds);
3703 /* Does user supply times argument? */
3704 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003705 cnt = 0;
3706
3707 ro = (repeatobject *)type->tp_alloc(type, 0);
3708 if (ro == NULL)
3709 return NULL;
3710 Py_INCREF(element);
3711 ro->element = element;
3712 ro->cnt = cnt;
3713 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003714}
3715
3716static void
3717repeat_dealloc(repeatobject *ro)
3718{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003719 PyObject_GC_UnTrack(ro);
3720 Py_XDECREF(ro->element);
3721 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003722}
3723
3724static int
3725repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3726{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003727 Py_VISIT(ro->element);
3728 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003729}
3730
3731static PyObject *
3732repeat_next(repeatobject *ro)
3733{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003734 if (ro->cnt == 0)
3735 return NULL;
3736 if (ro->cnt > 0)
3737 ro->cnt--;
3738 Py_INCREF(ro->element);
3739 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003740}
3741
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003742static PyObject *
3743repeat_repr(repeatobject *ro)
3744{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003745 PyObject *result, *objrepr;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003746
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003747 objrepr = PyObject_Repr(ro->element);
3748 if (objrepr == NULL)
3749 return NULL;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003750
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003751 if (ro->cnt == -1)
3752 result = PyString_FromFormat("repeat(%s)",
3753 PyString_AS_STRING(objrepr));
3754 else
3755 result = PyString_FromFormat("repeat(%s, %zd)",
3756 PyString_AS_STRING(objrepr), ro->cnt);
3757 Py_DECREF(objrepr);
3758 return result;
3759}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003760
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003761static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003762repeat_len(repeatobject *ro)
3763{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003764 if (ro->cnt == -1) {
3765 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3766 return NULL;
3767 }
3768 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003769}
3770
Armin Rigof5b3e362006-02-11 21:32:43 +00003771PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003772
3773static PyMethodDef repeat_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003774 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3775 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003776};
3777
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003778PyDoc_STRVAR(repeat_doc,
Raymond Hettinger182edae2009-02-19 02:38:25 +00003779"repeat(object [,times]) -> create an iterator which returns the object\n\
3780for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003781endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003782
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003783static PyTypeObject repeat_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003784 PyVarObject_HEAD_INIT(NULL, 0)
3785 "itertools.repeat", /* tp_name */
3786 sizeof(repeatobject), /* tp_basicsize */
3787 0, /* tp_itemsize */
3788 /* methods */
3789 (destructor)repeat_dealloc, /* tp_dealloc */
3790 0, /* tp_print */
3791 0, /* tp_getattr */
3792 0, /* tp_setattr */
3793 0, /* tp_compare */
3794 (reprfunc)repeat_repr, /* tp_repr */
3795 0, /* tp_as_number */
3796 0, /* tp_as_sequence */
3797 0, /* tp_as_mapping */
3798 0, /* tp_hash */
3799 0, /* tp_call */
3800 0, /* tp_str */
3801 PyObject_GenericGetAttr, /* tp_getattro */
3802 0, /* tp_setattro */
3803 0, /* tp_as_buffer */
3804 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3805 Py_TPFLAGS_BASETYPE, /* tp_flags */
3806 repeat_doc, /* tp_doc */
3807 (traverseproc)repeat_traverse, /* tp_traverse */
3808 0, /* tp_clear */
3809 0, /* tp_richcompare */
3810 0, /* tp_weaklistoffset */
3811 PyObject_SelfIter, /* tp_iter */
3812 (iternextfunc)repeat_next, /* tp_iternext */
3813 repeat_methods, /* tp_methods */
3814 0, /* tp_members */
3815 0, /* tp_getset */
3816 0, /* tp_base */
3817 0, /* tp_dict */
3818 0, /* tp_descr_get */
3819 0, /* tp_descr_set */
3820 0, /* tp_dictoffset */
3821 0, /* tp_init */
3822 0, /* tp_alloc */
3823 repeat_new, /* tp_new */
3824 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003825};
3826
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003827/* iziplongest object ************************************************************/
3828
3829#include "Python.h"
3830
3831typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003832 PyObject_HEAD
3833 Py_ssize_t tuplesize;
3834 Py_ssize_t numactive;
3835 PyObject *ittuple; /* tuple of iterators */
3836 PyObject *result;
3837 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003838} iziplongestobject;
3839
3840static PyTypeObject iziplongest_type;
3841
3842static PyObject *
3843izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3844{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003845 iziplongestobject *lz;
3846 Py_ssize_t i;
3847 PyObject *ittuple; /* tuple of iterators */
3848 PyObject *result;
3849 PyObject *fillvalue = Py_None;
3850 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003851
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003852 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3853 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3854 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3855 PyErr_SetString(PyExc_TypeError,
3856 "izip_longest() got an unexpected keyword argument");
3857 return NULL;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003858 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003859 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003860
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003861 /* args must be a tuple */
3862 assert(PyTuple_Check(args));
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003863
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003864 /* obtain iterators */
3865 ittuple = PyTuple_New(tuplesize);
3866 if (ittuple == NULL)
3867 return NULL;
3868 for (i=0; i < tuplesize; ++i) {
3869 PyObject *item = PyTuple_GET_ITEM(args, i);
3870 PyObject *it = PyObject_GetIter(item);
3871 if (it == NULL) {
3872 if (PyErr_ExceptionMatches(PyExc_TypeError))
3873 PyErr_Format(PyExc_TypeError,
3874 "izip_longest argument #%zd must support iteration",
3875 i+1);
3876 Py_DECREF(ittuple);
3877 return NULL;
3878 }
3879 PyTuple_SET_ITEM(ittuple, i, it);
3880 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003881
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003882 /* create a result holder */
3883 result = PyTuple_New(tuplesize);
3884 if (result == NULL) {
3885 Py_DECREF(ittuple);
3886 return NULL;
3887 }
3888 for (i=0 ; i < tuplesize ; i++) {
3889 Py_INCREF(Py_None);
3890 PyTuple_SET_ITEM(result, i, Py_None);
3891 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003892
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003893 /* create iziplongestobject structure */
3894 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3895 if (lz == NULL) {
3896 Py_DECREF(ittuple);
3897 Py_DECREF(result);
3898 return NULL;
3899 }
3900 lz->ittuple = ittuple;
3901 lz->tuplesize = tuplesize;
3902 lz->numactive = tuplesize;
3903 lz->result = result;
3904 Py_INCREF(fillvalue);
3905 lz->fillvalue = fillvalue;
3906 return (PyObject *)lz;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003907}
3908
3909static void
3910izip_longest_dealloc(iziplongestobject *lz)
3911{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003912 PyObject_GC_UnTrack(lz);
3913 Py_XDECREF(lz->ittuple);
3914 Py_XDECREF(lz->result);
3915 Py_XDECREF(lz->fillvalue);
3916 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003917}
3918
3919static int
3920izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3921{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003922 Py_VISIT(lz->ittuple);
3923 Py_VISIT(lz->result);
3924 Py_VISIT(lz->fillvalue);
3925 return 0;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003926}
3927
3928static PyObject *
3929izip_longest_next(iziplongestobject *lz)
3930{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003931 Py_ssize_t i;
3932 Py_ssize_t tuplesize = lz->tuplesize;
3933 PyObject *result = lz->result;
3934 PyObject *it;
3935 PyObject *item;
3936 PyObject *olditem;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003937
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003938 if (tuplesize == 0)
3939 return NULL;
3940 if (lz->numactive == 0)
3941 return NULL;
3942 if (Py_REFCNT(result) == 1) {
3943 Py_INCREF(result);
3944 for (i=0 ; i < tuplesize ; i++) {
3945 it = PyTuple_GET_ITEM(lz->ittuple, i);
3946 if (it == NULL) {
3947 Py_INCREF(lz->fillvalue);
3948 item = lz->fillvalue;
3949 } else {
3950 item = PyIter_Next(it);
3951 if (item == NULL) {
3952 lz->numactive -= 1;
3953 if (lz->numactive == 0 || PyErr_Occurred()) {
3954 lz->numactive = 0;
3955 Py_DECREF(result);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003956 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003957 } else {
3958 Py_INCREF(lz->fillvalue);
3959 item = lz->fillvalue;
3960 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3961 Py_DECREF(it);
3962 }
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003963 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003964 }
3965 olditem = PyTuple_GET_ITEM(result, i);
3966 PyTuple_SET_ITEM(result, i, item);
3967 Py_DECREF(olditem);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003968 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003969 } else {
3970 result = PyTuple_New(tuplesize);
3971 if (result == NULL)
3972 return NULL;
3973 for (i=0 ; i < tuplesize ; i++) {
3974 it = PyTuple_GET_ITEM(lz->ittuple, i);
3975 if (it == NULL) {
3976 Py_INCREF(lz->fillvalue);
3977 item = lz->fillvalue;
3978 } else {
3979 item = PyIter_Next(it);
3980 if (item == NULL) {
3981 lz->numactive -= 1;
3982 if (lz->numactive == 0 || PyErr_Occurred()) {
3983 lz->numactive = 0;
3984 Py_DECREF(result);
3985 return NULL;
3986 } else {
3987 Py_INCREF(lz->fillvalue);
3988 item = lz->fillvalue;
3989 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3990 Py_DECREF(it);
3991 }
3992 }
3993 }
3994 PyTuple_SET_ITEM(result, i, item);
3995 }
3996 }
3997 return result;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003998}
3999
4000PyDoc_STRVAR(izip_longest_doc,
4001"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
4002\n\
4003Return an izip_longest object whose .next() method returns a tuple where\n\
4004the i-th element comes from the i-th iterable argument. The .next()\n\
4005method continues until the longest iterable in the argument sequence\n\
4006is exhausted and then it raises StopIteration. When the shorter iterables\n\
4007are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4008defaults to None or can be specified by a keyword argument.\n\
4009");
4010
4011static PyTypeObject iziplongest_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004012 PyVarObject_HEAD_INIT(NULL, 0)
4013 "itertools.izip_longest", /* tp_name */
4014 sizeof(iziplongestobject), /* tp_basicsize */
4015 0, /* tp_itemsize */
4016 /* methods */
4017 (destructor)izip_longest_dealloc, /* tp_dealloc */
4018 0, /* tp_print */
4019 0, /* tp_getattr */
4020 0, /* tp_setattr */
4021 0, /* tp_compare */
4022 0, /* tp_repr */
4023 0, /* tp_as_number */
4024 0, /* tp_as_sequence */
4025 0, /* tp_as_mapping */
4026 0, /* tp_hash */
4027 0, /* tp_call */
4028 0, /* tp_str */
4029 PyObject_GenericGetAttr, /* tp_getattro */
4030 0, /* tp_setattro */
4031 0, /* tp_as_buffer */
4032 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4033 Py_TPFLAGS_BASETYPE, /* tp_flags */
4034 izip_longest_doc, /* tp_doc */
4035 (traverseproc)izip_longest_traverse, /* tp_traverse */
4036 0, /* tp_clear */
4037 0, /* tp_richcompare */
4038 0, /* tp_weaklistoffset */
4039 PyObject_SelfIter, /* tp_iter */
4040 (iternextfunc)izip_longest_next, /* tp_iternext */
4041 0, /* tp_methods */
4042 0, /* tp_members */
4043 0, /* tp_getset */
4044 0, /* tp_base */
4045 0, /* tp_dict */
4046 0, /* tp_descr_get */
4047 0, /* tp_descr_set */
4048 0, /* tp_dictoffset */
4049 0, /* tp_init */
4050 0, /* tp_alloc */
4051 izip_longest_new, /* tp_new */
4052 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd36862c2007-02-21 05:20:38 +00004053};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004054
4055/* module level code ********************************************************/
4056
4057PyDoc_STRVAR(module_doc,
4058"Functional tools for creating and using iterators.\n\
4059\n\
4060Infinite iterators:\n\
4061count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004062cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004063repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004064\n\
4065Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004066chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettingere76816b2009-01-29 03:43:44 +00004067compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004068dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4069groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00004070ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
4071ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004072islice(seq, [start,] stop [, step]) --> elements from\n\
4073 seq[start:stop:step]\n\
4074imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
4075starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004076tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004077takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004078izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4079izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4080\n\
4081Combinatoric generators:\n\
4082product(p, q, ... [repeat=1]) --> cartesian product\n\
4083permutations(p[, r])\n\
Raymond Hettinger9eac1192009-11-19 01:22:04 +00004084combinations(p, r)\n\
4085combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004086");
4087
4088
Raymond Hettingerad983e72003-11-12 14:32:26 +00004089static PyMethodDef module_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004090 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4091 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004092};
4093
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004094PyMODINIT_FUNC
4095inititertools(void)
4096{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004097 int i;
4098 PyObject *m;
4099 char *name;
4100 PyTypeObject *typelist[] = {
4101 &combinations_type,
4102 &cwr_type,
4103 &cycle_type,
4104 &dropwhile_type,
4105 &takewhile_type,
4106 &islice_type,
4107 &starmap_type,
4108 &imap_type,
4109 &chain_type,
4110 &compress_type,
4111 &ifilter_type,
4112 &ifilterfalse_type,
4113 &count_type,
4114 &izip_type,
4115 &iziplongest_type,
4116 &permutations_type,
4117 &product_type,
4118 &repeat_type,
4119 &groupby_type,
4120 NULL
4121 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004122
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004123 Py_TYPE(&teedataobject_type) = &PyType_Type;
4124 m = Py_InitModule3("itertools", module_methods, module_doc);
4125 if (m == NULL)
4126 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004127
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004128 for (i=0 ; typelist[i] != NULL ; i++) {
4129 if (PyType_Ready(typelist[i]) < 0)
4130 return;
4131 name = strchr(typelist[i]->tp_name, '.');
4132 assert (name != NULL);
4133 Py_INCREF(typelist[i]);
4134 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4135 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004136
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004137 if (PyType_Ready(&teedataobject_type) < 0)
4138 return;
4139 if (PyType_Ready(&tee_type) < 0)
4140 return;
4141 if (PyType_Ready(&_grouper_type) < 0)
4142 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004143}