blob: 9b9e1daedf8636579949fee786cf8c86beaab64d [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;
Serhiy Storchaka5951f232015-12-24 10:35:35 +0200497 Py_SETREF(to->dataobj, (teedataobject *)link);
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000498 to->index = 0;
499 }
500 value = teedataobject_getitem(to->dataobj, to->index);
501 if (value == NULL)
502 return NULL;
503 to->index++;
504 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000505}
506
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000507static int
508tee_traverse(teeobject *to, visitproc visit, void *arg)
509{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000510 Py_VISIT((PyObject *)to->dataobj);
511 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000512}
513
Raymond Hettingerad983e72003-11-12 14:32:26 +0000514static PyObject *
515tee_copy(teeobject *to)
516{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000517 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000518
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000519 newto = PyObject_GC_New(teeobject, &tee_type);
520 if (newto == NULL)
521 return NULL;
522 Py_INCREF(to->dataobj);
523 newto->dataobj = to->dataobj;
524 newto->index = to->index;
525 newto->weakreflist = NULL;
526 PyObject_GC_Track(newto);
527 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000528}
529
530PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
531
532static PyObject *
533tee_fromiterable(PyObject *iterable)
534{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000535 teeobject *to;
536 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000537
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000538 it = PyObject_GetIter(iterable);
539 if (it == NULL)
540 return NULL;
541 if (PyObject_TypeCheck(it, &tee_type)) {
542 to = (teeobject *)tee_copy((teeobject *)it);
543 goto done;
544 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000545
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000546 to = PyObject_GC_New(teeobject, &tee_type);
547 if (to == NULL)
548 goto done;
549 to->dataobj = (teedataobject *)teedataobject_new(it);
550 if (!to->dataobj) {
551 PyObject_GC_Del(to);
552 to = NULL;
553 goto done;
554 }
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000555
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000556 to->index = 0;
557 to->weakreflist = NULL;
558 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000559done:
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000560 Py_XDECREF(it);
561 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000562}
563
564static PyObject *
565tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
566{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000567 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000568
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000569 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
570 return NULL;
571 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000572}
573
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000574static int
575tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000576{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000577 if (to->weakreflist != NULL)
578 PyObject_ClearWeakRefs((PyObject *) to);
579 Py_CLEAR(to->dataobj);
580 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000581}
582
583static void
584tee_dealloc(teeobject *to)
585{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000586 PyObject_GC_UnTrack(to);
587 tee_clear(to);
588 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000589}
590
Raymond Hettingerad983e72003-11-12 14:32:26 +0000591PyDoc_STRVAR(teeobject_doc,
592"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000593
Raymond Hettingerad983e72003-11-12 14:32:26 +0000594static PyMethodDef tee_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000595 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
596 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000597};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000598
599static PyTypeObject tee_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000600 PyVarObject_HEAD_INIT(NULL, 0)
601 "itertools.tee", /* tp_name */
602 sizeof(teeobject), /* tp_basicsize */
603 0, /* tp_itemsize */
604 /* methods */
605 (destructor)tee_dealloc, /* tp_dealloc */
606 0, /* tp_print */
607 0, /* tp_getattr */
608 0, /* tp_setattr */
609 0, /* tp_compare */
610 0, /* tp_repr */
611 0, /* tp_as_number */
612 0, /* tp_as_sequence */
613 0, /* tp_as_mapping */
614 0, /* tp_hash */
615 0, /* tp_call */
616 0, /* tp_str */
617 0, /* tp_getattro */
618 0, /* tp_setattro */
619 0, /* tp_as_buffer */
620 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
621 teeobject_doc, /* tp_doc */
622 (traverseproc)tee_traverse, /* tp_traverse */
623 (inquiry)tee_clear, /* tp_clear */
624 0, /* tp_richcompare */
625 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
626 PyObject_SelfIter, /* tp_iter */
627 (iternextfunc)tee_next, /* tp_iternext */
628 tee_methods, /* tp_methods */
629 0, /* tp_members */
630 0, /* tp_getset */
631 0, /* tp_base */
632 0, /* tp_dict */
633 0, /* tp_descr_get */
634 0, /* tp_descr_set */
635 0, /* tp_dictoffset */
636 0, /* tp_init */
637 0, /* tp_alloc */
638 tee_new, /* tp_new */
639 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000640};
641
Raymond Hettingerad983e72003-11-12 14:32:26 +0000642static PyObject *
643tee(PyObject *self, PyObject *args)
644{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000645 Py_ssize_t i, n=2;
646 PyObject *it, *iterable, *copyable, *result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000647
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000648 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
649 return NULL;
650 if (n < 0) {
651 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
652 return NULL;
653 }
654 result = PyTuple_New(n);
655 if (result == NULL)
656 return NULL;
657 if (n == 0)
658 return result;
659 it = PyObject_GetIter(iterable);
660 if (it == NULL) {
661 Py_DECREF(result);
662 return NULL;
663 }
664 if (!PyObject_HasAttrString(it, "__copy__")) {
665 copyable = tee_fromiterable(it);
666 Py_DECREF(it);
667 if (copyable == NULL) {
668 Py_DECREF(result);
669 return NULL;
670 }
671 } else
672 copyable = it;
673 PyTuple_SET_ITEM(result, 0, copyable);
674 for (i=1 ; i<n ; i++) {
675 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
676 if (copyable == NULL) {
677 Py_DECREF(result);
678 return NULL;
679 }
680 PyTuple_SET_ITEM(result, i, copyable);
681 }
682 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000683}
684
685PyDoc_STRVAR(tee_doc,
686"tee(iterable, n=2) --> tuple of n independent iterators.");
687
688
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000689/* cycle object **********************************************************/
690
691typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000692 PyObject_HEAD
693 PyObject *it;
694 PyObject *saved;
695 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000696} cycleobject;
697
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000698static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000699
700static PyObject *
701cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
702{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000703 PyObject *it;
704 PyObject *iterable;
705 PyObject *saved;
706 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000707
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000708 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
709 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000710
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000711 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
712 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000713
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000714 /* Get iterator. */
715 it = PyObject_GetIter(iterable);
716 if (it == NULL)
717 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000718
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000719 saved = PyList_New(0);
720 if (saved == NULL) {
721 Py_DECREF(it);
722 return NULL;
723 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000724
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000725 /* create cycleobject structure */
726 lz = (cycleobject *)type->tp_alloc(type, 0);
727 if (lz == NULL) {
728 Py_DECREF(it);
729 Py_DECREF(saved);
730 return NULL;
731 }
732 lz->it = it;
733 lz->saved = saved;
734 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000735
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000736 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000737}
738
739static void
740cycle_dealloc(cycleobject *lz)
741{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000742 PyObject_GC_UnTrack(lz);
743 Py_XDECREF(lz->saved);
744 Py_XDECREF(lz->it);
745 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000746}
747
748static int
749cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
750{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000751 Py_VISIT(lz->it);
752 Py_VISIT(lz->saved);
753 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000754}
755
756static PyObject *
757cycle_next(cycleobject *lz)
758{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000759 PyObject *item;
760 PyObject *it;
761 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000762
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000763 while (1) {
764 item = PyIter_Next(lz->it);
765 if (item != NULL) {
766 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
767 Py_DECREF(item);
768 return NULL;
769 }
770 return item;
771 }
772 if (PyErr_Occurred()) {
773 if (PyErr_ExceptionMatches(PyExc_StopIteration))
774 PyErr_Clear();
775 else
776 return NULL;
777 }
778 if (PyList_Size(lz->saved) == 0)
779 return NULL;
780 it = PyObject_GetIter(lz->saved);
781 if (it == NULL)
782 return NULL;
783 tmp = lz->it;
784 lz->it = it;
785 lz->firstpass = 1;
786 Py_DECREF(tmp);
787 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000788}
789
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000790PyDoc_STRVAR(cycle_doc,
791"cycle(iterable) --> cycle object\n\
792\n\
793Return elements from the iterable until it is exhausted.\n\
794Then repeat the sequence indefinitely.");
795
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000796static PyTypeObject cycle_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000797 PyVarObject_HEAD_INIT(NULL, 0)
798 "itertools.cycle", /* tp_name */
799 sizeof(cycleobject), /* tp_basicsize */
800 0, /* tp_itemsize */
801 /* methods */
802 (destructor)cycle_dealloc, /* tp_dealloc */
803 0, /* tp_print */
804 0, /* tp_getattr */
805 0, /* tp_setattr */
806 0, /* tp_compare */
807 0, /* tp_repr */
808 0, /* tp_as_number */
809 0, /* tp_as_sequence */
810 0, /* tp_as_mapping */
811 0, /* tp_hash */
812 0, /* tp_call */
813 0, /* tp_str */
814 PyObject_GenericGetAttr, /* tp_getattro */
815 0, /* tp_setattro */
816 0, /* tp_as_buffer */
817 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
818 Py_TPFLAGS_BASETYPE, /* tp_flags */
819 cycle_doc, /* tp_doc */
820 (traverseproc)cycle_traverse, /* tp_traverse */
821 0, /* tp_clear */
822 0, /* tp_richcompare */
823 0, /* tp_weaklistoffset */
824 PyObject_SelfIter, /* tp_iter */
825 (iternextfunc)cycle_next, /* tp_iternext */
826 0, /* tp_methods */
827 0, /* tp_members */
828 0, /* tp_getset */
829 0, /* tp_base */
830 0, /* tp_dict */
831 0, /* tp_descr_get */
832 0, /* tp_descr_set */
833 0, /* tp_dictoffset */
834 0, /* tp_init */
835 0, /* tp_alloc */
836 cycle_new, /* tp_new */
837 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000838};
839
840
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000841/* dropwhile object **********************************************************/
842
843typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000844 PyObject_HEAD
845 PyObject *func;
846 PyObject *it;
847 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000848} dropwhileobject;
849
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000850static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000851
852static PyObject *
853dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
854{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000855 PyObject *func, *seq;
856 PyObject *it;
857 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000858
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000859 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
860 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000861
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000862 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
863 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000864
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000865 /* Get iterator. */
866 it = PyObject_GetIter(seq);
867 if (it == NULL)
868 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000869
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000870 /* create dropwhileobject structure */
871 lz = (dropwhileobject *)type->tp_alloc(type, 0);
872 if (lz == NULL) {
873 Py_DECREF(it);
874 return NULL;
875 }
876 Py_INCREF(func);
877 lz->func = func;
878 lz->it = it;
879 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000880
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000881 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000882}
883
884static void
885dropwhile_dealloc(dropwhileobject *lz)
886{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000887 PyObject_GC_UnTrack(lz);
888 Py_XDECREF(lz->func);
889 Py_XDECREF(lz->it);
890 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000891}
892
893static int
894dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
895{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000896 Py_VISIT(lz->it);
897 Py_VISIT(lz->func);
898 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000899}
900
901static PyObject *
902dropwhile_next(dropwhileobject *lz)
903{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000904 PyObject *item, *good;
905 PyObject *it = lz->it;
906 long ok;
907 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000908
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000909 iternext = *Py_TYPE(it)->tp_iternext;
910 for (;;) {
911 item = iternext(it);
912 if (item == NULL)
913 return NULL;
914 if (lz->start == 1)
915 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000916
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000917 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
918 if (good == NULL) {
919 Py_DECREF(item);
920 return NULL;
921 }
922 ok = PyObject_IsTrue(good);
923 Py_DECREF(good);
Antoine Pitrouc5bef752012-08-15 23:16:51 +0200924 if (ok == 0) {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000925 lz->start = 1;
926 return item;
927 }
928 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +0200929 if (ok < 0)
930 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000931 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000932}
933
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000934PyDoc_STRVAR(dropwhile_doc,
935"dropwhile(predicate, iterable) --> dropwhile object\n\
936\n\
937Drop items from the iterable while predicate(item) is true.\n\
938Afterwards, return every element until the iterable is exhausted.");
939
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000940static PyTypeObject dropwhile_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000941 PyVarObject_HEAD_INIT(NULL, 0)
942 "itertools.dropwhile", /* tp_name */
943 sizeof(dropwhileobject), /* tp_basicsize */
944 0, /* tp_itemsize */
945 /* methods */
946 (destructor)dropwhile_dealloc, /* tp_dealloc */
947 0, /* tp_print */
948 0, /* tp_getattr */
949 0, /* tp_setattr */
950 0, /* tp_compare */
951 0, /* tp_repr */
952 0, /* tp_as_number */
953 0, /* tp_as_sequence */
954 0, /* tp_as_mapping */
955 0, /* tp_hash */
956 0, /* tp_call */
957 0, /* tp_str */
958 PyObject_GenericGetAttr, /* tp_getattro */
959 0, /* tp_setattro */
960 0, /* tp_as_buffer */
961 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
962 Py_TPFLAGS_BASETYPE, /* tp_flags */
963 dropwhile_doc, /* tp_doc */
964 (traverseproc)dropwhile_traverse, /* tp_traverse */
965 0, /* tp_clear */
966 0, /* tp_richcompare */
967 0, /* tp_weaklistoffset */
968 PyObject_SelfIter, /* tp_iter */
969 (iternextfunc)dropwhile_next, /* tp_iternext */
970 0, /* tp_methods */
971 0, /* tp_members */
972 0, /* tp_getset */
973 0, /* tp_base */
974 0, /* tp_dict */
975 0, /* tp_descr_get */
976 0, /* tp_descr_set */
977 0, /* tp_dictoffset */
978 0, /* tp_init */
979 0, /* tp_alloc */
980 dropwhile_new, /* tp_new */
981 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000982};
983
984
985/* takewhile object **********************************************************/
986
987typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000988 PyObject_HEAD
989 PyObject *func;
990 PyObject *it;
991 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000992} takewhileobject;
993
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000994static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000995
996static PyObject *
997takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
998{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000999 PyObject *func, *seq;
1000 PyObject *it;
1001 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001002
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001003 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1004 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001005
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001006 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1007 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001008
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001009 /* Get iterator. */
1010 it = PyObject_GetIter(seq);
1011 if (it == NULL)
1012 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001013
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001014 /* create takewhileobject structure */
1015 lz = (takewhileobject *)type->tp_alloc(type, 0);
1016 if (lz == NULL) {
1017 Py_DECREF(it);
1018 return NULL;
1019 }
1020 Py_INCREF(func);
1021 lz->func = func;
1022 lz->it = it;
1023 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001024
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001025 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001026}
1027
1028static void
1029takewhile_dealloc(takewhileobject *lz)
1030{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001031 PyObject_GC_UnTrack(lz);
1032 Py_XDECREF(lz->func);
1033 Py_XDECREF(lz->it);
1034 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001035}
1036
1037static int
1038takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1039{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001040 Py_VISIT(lz->it);
1041 Py_VISIT(lz->func);
1042 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001043}
1044
1045static PyObject *
1046takewhile_next(takewhileobject *lz)
1047{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001048 PyObject *item, *good;
1049 PyObject *it = lz->it;
1050 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001051
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001052 if (lz->stop == 1)
1053 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001054
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001055 item = (*Py_TYPE(it)->tp_iternext)(it);
1056 if (item == NULL)
1057 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001058
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001059 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1060 if (good == NULL) {
1061 Py_DECREF(item);
1062 return NULL;
1063 }
1064 ok = PyObject_IsTrue(good);
1065 Py_DECREF(good);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02001066 if (ok > 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001067 return item;
1068 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02001069 if (ok == 0)
1070 lz->stop = 1;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001071 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001072}
1073
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001074PyDoc_STRVAR(takewhile_doc,
1075"takewhile(predicate, iterable) --> takewhile object\n\
1076\n\
1077Return successive entries from an iterable as long as the \n\
1078predicate evaluates to true for each entry.");
1079
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001080static PyTypeObject takewhile_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001081 PyVarObject_HEAD_INIT(NULL, 0)
1082 "itertools.takewhile", /* tp_name */
1083 sizeof(takewhileobject), /* tp_basicsize */
1084 0, /* tp_itemsize */
1085 /* methods */
1086 (destructor)takewhile_dealloc, /* tp_dealloc */
1087 0, /* tp_print */
1088 0, /* tp_getattr */
1089 0, /* tp_setattr */
1090 0, /* tp_compare */
1091 0, /* tp_repr */
1092 0, /* tp_as_number */
1093 0, /* tp_as_sequence */
1094 0, /* tp_as_mapping */
1095 0, /* tp_hash */
1096 0, /* tp_call */
1097 0, /* tp_str */
1098 PyObject_GenericGetAttr, /* tp_getattro */
1099 0, /* tp_setattro */
1100 0, /* tp_as_buffer */
1101 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1102 Py_TPFLAGS_BASETYPE, /* tp_flags */
1103 takewhile_doc, /* tp_doc */
1104 (traverseproc)takewhile_traverse, /* tp_traverse */
1105 0, /* tp_clear */
1106 0, /* tp_richcompare */
1107 0, /* tp_weaklistoffset */
1108 PyObject_SelfIter, /* tp_iter */
1109 (iternextfunc)takewhile_next, /* tp_iternext */
1110 0, /* tp_methods */
1111 0, /* tp_members */
1112 0, /* tp_getset */
1113 0, /* tp_base */
1114 0, /* tp_dict */
1115 0, /* tp_descr_get */
1116 0, /* tp_descr_set */
1117 0, /* tp_dictoffset */
1118 0, /* tp_init */
1119 0, /* tp_alloc */
1120 takewhile_new, /* tp_new */
1121 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001122};
1123
1124
1125/* islice object ************************************************************/
1126
1127typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001128 PyObject_HEAD
1129 PyObject *it;
1130 Py_ssize_t next;
1131 Py_ssize_t stop;
1132 Py_ssize_t step;
1133 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001134} isliceobject;
1135
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001136static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001137
1138static PyObject *
1139islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1140{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001141 PyObject *seq;
1142 Py_ssize_t start=0, stop=-1, step=1;
1143 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1144 Py_ssize_t numargs;
1145 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001146
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001147 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1148 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001149
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001150 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1151 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001152
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001153 numargs = PyTuple_Size(args);
1154 if (numargs == 2) {
1155 if (a1 != Py_None) {
1156 stop = PyInt_AsSsize_t(a1);
1157 if (stop == -1) {
1158 if (PyErr_Occurred())
1159 PyErr_Clear();
1160 PyErr_SetString(PyExc_ValueError,
1161 "Stop argument for islice() must be None or an integer: 0 <= x <= maxint.");
1162 return NULL;
1163 }
1164 }
1165 } else {
1166 if (a1 != Py_None)
1167 start = PyInt_AsSsize_t(a1);
1168 if (start == -1 && PyErr_Occurred())
1169 PyErr_Clear();
1170 if (a2 != Py_None) {
1171 stop = PyInt_AsSsize_t(a2);
1172 if (stop == -1) {
1173 if (PyErr_Occurred())
1174 PyErr_Clear();
1175 PyErr_SetString(PyExc_ValueError,
1176 "Stop argument for islice() must be None or an integer: 0 <= x <= maxint.");
1177 return NULL;
1178 }
1179 }
1180 }
1181 if (start<0 || stop<-1) {
1182 PyErr_SetString(PyExc_ValueError,
1183 "Indices for islice() must be None or an integer: 0 <= x <= maxint.");
1184 return NULL;
1185 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001186
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001187 if (a3 != NULL) {
1188 if (a3 != Py_None)
1189 step = PyInt_AsSsize_t(a3);
1190 if (step == -1 && PyErr_Occurred())
1191 PyErr_Clear();
1192 }
1193 if (step<1) {
1194 PyErr_SetString(PyExc_ValueError,
1195 "Step for islice() must be a positive integer or None.");
1196 return NULL;
1197 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001198
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001199 /* Get iterator. */
1200 it = PyObject_GetIter(seq);
1201 if (it == NULL)
1202 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001203
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001204 /* create isliceobject structure */
1205 lz = (isliceobject *)type->tp_alloc(type, 0);
1206 if (lz == NULL) {
1207 Py_DECREF(it);
1208 return NULL;
1209 }
1210 lz->it = it;
1211 lz->next = start;
1212 lz->stop = stop;
1213 lz->step = step;
1214 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001215
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001216 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001217}
1218
1219static void
1220islice_dealloc(isliceobject *lz)
1221{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001222 PyObject_GC_UnTrack(lz);
1223 Py_XDECREF(lz->it);
1224 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001225}
1226
1227static int
1228islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1229{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001230 Py_VISIT(lz->it);
1231 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001232}
1233
1234static PyObject *
1235islice_next(isliceobject *lz)
1236{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001237 PyObject *item;
1238 PyObject *it = lz->it;
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001239 Py_ssize_t stop = lz->stop;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001240 Py_ssize_t oldnext;
1241 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001242
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001243 if (it == NULL)
1244 return NULL;
1245
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001246 iternext = *Py_TYPE(it)->tp_iternext;
1247 while (lz->cnt < lz->next) {
1248 item = iternext(it);
1249 if (item == NULL)
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001250 goto empty;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001251 Py_DECREF(item);
1252 lz->cnt++;
1253 }
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001254 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001255 goto empty;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001256 item = iternext(it);
1257 if (item == NULL)
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001258 goto empty;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001259 lz->cnt++;
1260 oldnext = lz->next;
Mark Dickinsona96b0d12011-09-24 09:01:16 +01001261 /* The (size_t) cast below avoids the danger of undefined
1262 behaviour from signed integer overflow. */
1263 lz->next += (size_t)lz->step;
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001264 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1265 lz->next = stop;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001266 return item;
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001267
1268empty:
1269 Py_CLEAR(lz->it);
1270 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001271}
1272
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001273PyDoc_STRVAR(islice_doc,
1274"islice(iterable, [start,] stop [, step]) --> islice object\n\
1275\n\
1276Return an iterator whose next() method returns selected values from an\n\
1277iterable. If start is specified, will skip all preceding elements;\n\
1278otherwise, start defaults to zero. Step defaults to one. If\n\
1279specified as another value, step determines how many values are \n\
1280skipped between successive calls. Works like a slice() on a list\n\
1281but returns an iterator.");
1282
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001283static PyTypeObject islice_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001284 PyVarObject_HEAD_INIT(NULL, 0)
1285 "itertools.islice", /* tp_name */
1286 sizeof(isliceobject), /* tp_basicsize */
1287 0, /* tp_itemsize */
1288 /* methods */
1289 (destructor)islice_dealloc, /* tp_dealloc */
1290 0, /* tp_print */
1291 0, /* tp_getattr */
1292 0, /* tp_setattr */
1293 0, /* tp_compare */
1294 0, /* tp_repr */
1295 0, /* tp_as_number */
1296 0, /* tp_as_sequence */
1297 0, /* tp_as_mapping */
1298 0, /* tp_hash */
1299 0, /* tp_call */
1300 0, /* tp_str */
1301 PyObject_GenericGetAttr, /* tp_getattro */
1302 0, /* tp_setattro */
1303 0, /* tp_as_buffer */
1304 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1305 Py_TPFLAGS_BASETYPE, /* tp_flags */
1306 islice_doc, /* tp_doc */
1307 (traverseproc)islice_traverse, /* tp_traverse */
1308 0, /* tp_clear */
1309 0, /* tp_richcompare */
1310 0, /* tp_weaklistoffset */
1311 PyObject_SelfIter, /* tp_iter */
1312 (iternextfunc)islice_next, /* tp_iternext */
1313 0, /* tp_methods */
1314 0, /* tp_members */
1315 0, /* tp_getset */
1316 0, /* tp_base */
1317 0, /* tp_dict */
1318 0, /* tp_descr_get */
1319 0, /* tp_descr_set */
1320 0, /* tp_dictoffset */
1321 0, /* tp_init */
1322 0, /* tp_alloc */
1323 islice_new, /* tp_new */
1324 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001325};
1326
1327
1328/* starmap object ************************************************************/
1329
1330typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001331 PyObject_HEAD
1332 PyObject *func;
1333 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001334} starmapobject;
1335
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001336static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001337
1338static PyObject *
1339starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1340{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001341 PyObject *func, *seq;
1342 PyObject *it;
1343 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001344
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001345 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1346 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001347
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001348 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1349 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001350
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001351 /* Get iterator. */
1352 it = PyObject_GetIter(seq);
1353 if (it == NULL)
1354 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001355
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001356 /* create starmapobject structure */
1357 lz = (starmapobject *)type->tp_alloc(type, 0);
1358 if (lz == NULL) {
1359 Py_DECREF(it);
1360 return NULL;
1361 }
1362 Py_INCREF(func);
1363 lz->func = func;
1364 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001365
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001366 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001367}
1368
1369static void
1370starmap_dealloc(starmapobject *lz)
1371{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001372 PyObject_GC_UnTrack(lz);
1373 Py_XDECREF(lz->func);
1374 Py_XDECREF(lz->it);
1375 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001376}
1377
1378static int
1379starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1380{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001381 Py_VISIT(lz->it);
1382 Py_VISIT(lz->func);
1383 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001384}
1385
1386static PyObject *
1387starmap_next(starmapobject *lz)
1388{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001389 PyObject *args;
1390 PyObject *result;
1391 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001392
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001393 args = (*Py_TYPE(it)->tp_iternext)(it);
1394 if (args == NULL)
1395 return NULL;
1396 if (!PyTuple_CheckExact(args)) {
1397 PyObject *newargs = PySequence_Tuple(args);
1398 Py_DECREF(args);
1399 if (newargs == NULL)
1400 return NULL;
1401 args = newargs;
1402 }
1403 result = PyObject_Call(lz->func, args, NULL);
1404 Py_DECREF(args);
1405 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001406}
1407
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001408PyDoc_STRVAR(starmap_doc,
1409"starmap(function, sequence) --> starmap object\n\
1410\n\
1411Return an iterator whose values are returned from the function evaluated\n\
Serhiy Storchakac72e66a2015-11-02 15:06:09 +02001412with an argument tuple taken from the given sequence.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001413
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001414static PyTypeObject starmap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001415 PyVarObject_HEAD_INIT(NULL, 0)
1416 "itertools.starmap", /* tp_name */
1417 sizeof(starmapobject), /* tp_basicsize */
1418 0, /* tp_itemsize */
1419 /* methods */
1420 (destructor)starmap_dealloc, /* tp_dealloc */
1421 0, /* tp_print */
1422 0, /* tp_getattr */
1423 0, /* tp_setattr */
1424 0, /* tp_compare */
1425 0, /* tp_repr */
1426 0, /* tp_as_number */
1427 0, /* tp_as_sequence */
1428 0, /* tp_as_mapping */
1429 0, /* tp_hash */
1430 0, /* tp_call */
1431 0, /* tp_str */
1432 PyObject_GenericGetAttr, /* tp_getattro */
1433 0, /* tp_setattro */
1434 0, /* tp_as_buffer */
1435 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1436 Py_TPFLAGS_BASETYPE, /* tp_flags */
1437 starmap_doc, /* tp_doc */
1438 (traverseproc)starmap_traverse, /* tp_traverse */
1439 0, /* tp_clear */
1440 0, /* tp_richcompare */
1441 0, /* tp_weaklistoffset */
1442 PyObject_SelfIter, /* tp_iter */
1443 (iternextfunc)starmap_next, /* tp_iternext */
1444 0, /* tp_methods */
1445 0, /* tp_members */
1446 0, /* tp_getset */
1447 0, /* tp_base */
1448 0, /* tp_dict */
1449 0, /* tp_descr_get */
1450 0, /* tp_descr_set */
1451 0, /* tp_dictoffset */
1452 0, /* tp_init */
1453 0, /* tp_alloc */
1454 starmap_new, /* tp_new */
1455 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001456};
1457
1458
1459/* imap object ************************************************************/
1460
1461typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001462 PyObject_HEAD
1463 PyObject *iters;
1464 PyObject *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001465} imapobject;
1466
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001467static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001468
1469static PyObject *
1470imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1471{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001472 PyObject *it, *iters, *func;
1473 imapobject *lz;
1474 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001475
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001476 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
1477 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001478
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001479 numargs = PyTuple_Size(args);
1480 if (numargs < 2) {
1481 PyErr_SetString(PyExc_TypeError,
1482 "imap() must have at least two arguments.");
1483 return NULL;
1484 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001485
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001486 iters = PyTuple_New(numargs-1);
1487 if (iters == NULL)
1488 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001489
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001490 for (i=1 ; i<numargs ; i++) {
1491 /* Get iterator. */
1492 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1493 if (it == NULL) {
1494 Py_DECREF(iters);
1495 return NULL;
1496 }
1497 PyTuple_SET_ITEM(iters, i-1, it);
1498 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001499
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001500 /* create imapobject structure */
1501 lz = (imapobject *)type->tp_alloc(type, 0);
1502 if (lz == NULL) {
1503 Py_DECREF(iters);
1504 return NULL;
1505 }
1506 lz->iters = iters;
1507 func = PyTuple_GET_ITEM(args, 0);
1508 Py_INCREF(func);
1509 lz->func = func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001510
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001511 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001512}
1513
1514static void
1515imap_dealloc(imapobject *lz)
1516{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001517 PyObject_GC_UnTrack(lz);
1518 Py_XDECREF(lz->iters);
1519 Py_XDECREF(lz->func);
1520 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001521}
1522
1523static int
1524imap_traverse(imapobject *lz, visitproc visit, void *arg)
1525{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001526 Py_VISIT(lz->iters);
1527 Py_VISIT(lz->func);
1528 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001529}
1530
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001531/*
Raymond Hettinger2012f172003-02-07 05:32:58 +00001532imap() is an iterator version of __builtins__.map() except that it does
1533not have the None fill-in feature. That was intentionally left out for
1534the following reasons:
1535
1536 1) Itertools are designed to be easily combined and chained together.
1537 Having all tools stop with the shortest input is a unifying principle
1538 that makes it easier to combine finite iterators (supplying data) with
1539 infinite iterators like count() and repeat() (for supplying sequential
1540 or constant arguments to a function).
1541
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001542 2) In typical use cases for combining itertools, having one finite data
Raymond Hettinger2012f172003-02-07 05:32:58 +00001543 supplier run out before another is likely to be an error condition which
1544 should not pass silently by automatically supplying None.
1545
1546 3) The use cases for automatic None fill-in are rare -- not many functions
1547 do something useful when a parameter suddenly switches type and becomes
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001548 None.
Raymond Hettinger2012f172003-02-07 05:32:58 +00001549
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001550 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001551 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001552
1553 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1554*/
1555
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001556static PyObject *
1557imap_next(imapobject *lz)
1558{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001559 PyObject *val;
1560 PyObject *argtuple;
1561 PyObject *result;
1562 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001563
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001564 numargs = PyTuple_Size(lz->iters);
1565 argtuple = PyTuple_New(numargs);
1566 if (argtuple == NULL)
1567 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001568
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001569 for (i=0 ; i<numargs ; i++) {
1570 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1571 if (val == NULL) {
1572 Py_DECREF(argtuple);
1573 return NULL;
1574 }
1575 PyTuple_SET_ITEM(argtuple, i, val);
1576 }
1577 if (lz->func == Py_None)
1578 return argtuple;
1579 result = PyObject_Call(lz->func, argtuple, NULL);
1580 Py_DECREF(argtuple);
1581 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001582}
1583
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001584PyDoc_STRVAR(imap_doc,
1585"imap(func, *iterables) --> imap object\n\
1586\n\
1587Make an iterator that computes the function using arguments from\n\
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001588each of the iterables. Like map() except that it returns\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001589an iterator instead of a list and that it stops when the shortest\n\
1590iterable is exhausted instead of filling in None for shorter\n\
1591iterables.");
1592
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001593static PyTypeObject imap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001594 PyVarObject_HEAD_INIT(NULL, 0)
1595 "itertools.imap", /* tp_name */
1596 sizeof(imapobject), /* tp_basicsize */
1597 0, /* tp_itemsize */
1598 /* methods */
1599 (destructor)imap_dealloc, /* tp_dealloc */
1600 0, /* tp_print */
1601 0, /* tp_getattr */
1602 0, /* tp_setattr */
1603 0, /* tp_compare */
1604 0, /* tp_repr */
1605 0, /* tp_as_number */
1606 0, /* tp_as_sequence */
1607 0, /* tp_as_mapping */
1608 0, /* tp_hash */
1609 0, /* tp_call */
1610 0, /* tp_str */
1611 PyObject_GenericGetAttr, /* tp_getattro */
1612 0, /* tp_setattro */
1613 0, /* tp_as_buffer */
1614 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1615 Py_TPFLAGS_BASETYPE, /* tp_flags */
1616 imap_doc, /* tp_doc */
1617 (traverseproc)imap_traverse, /* tp_traverse */
1618 0, /* tp_clear */
1619 0, /* tp_richcompare */
1620 0, /* tp_weaklistoffset */
1621 PyObject_SelfIter, /* tp_iter */
1622 (iternextfunc)imap_next, /* tp_iternext */
1623 0, /* tp_methods */
1624 0, /* tp_members */
1625 0, /* tp_getset */
1626 0, /* tp_base */
1627 0, /* tp_dict */
1628 0, /* tp_descr_get */
1629 0, /* tp_descr_set */
1630 0, /* tp_dictoffset */
1631 0, /* tp_init */
1632 0, /* tp_alloc */
1633 imap_new, /* tp_new */
1634 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001635};
1636
1637
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001638/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001639
1640typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001641 PyObject_HEAD
1642 PyObject *source; /* Iterator over input iterables */
1643 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001644} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001645
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001646static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001647
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001648static PyObject *
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001649chain_new_internal(PyTypeObject *type, PyObject *source)
1650{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001651 chainobject *lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001652
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001653 lz = (chainobject *)type->tp_alloc(type, 0);
1654 if (lz == NULL) {
1655 Py_DECREF(source);
1656 return NULL;
1657 }
1658
1659 lz->source = source;
1660 lz->active = NULL;
1661 return (PyObject *)lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001662}
1663
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001664static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001665chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001666{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001667 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001668
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001669 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1670 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001671
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001672 source = PyObject_GetIter(args);
1673 if (source == NULL)
1674 return NULL;
1675
1676 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001677}
1678
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001679static PyObject *
1680chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1681{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001682 PyObject *source;
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001683
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001684 source = PyObject_GetIter(arg);
1685 if (source == NULL)
1686 return NULL;
1687
1688 return chain_new_internal(type, source);
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001689}
1690
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001691static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001692chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001693{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001694 PyObject_GC_UnTrack(lz);
1695 Py_XDECREF(lz->active);
1696 Py_XDECREF(lz->source);
1697 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001698}
1699
Raymond Hettinger2012f172003-02-07 05:32:58 +00001700static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001701chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001702{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001703 Py_VISIT(lz->source);
1704 Py_VISIT(lz->active);
1705 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001706}
1707
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001708static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001709chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001710{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001711 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001712
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001713 if (lz->source == NULL)
1714 return NULL; /* already stopped */
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001715
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001716 if (lz->active == NULL) {
1717 PyObject *iterable = PyIter_Next(lz->source);
1718 if (iterable == NULL) {
1719 Py_CLEAR(lz->source);
1720 return NULL; /* no more input sources */
1721 }
1722 lz->active = PyObject_GetIter(iterable);
1723 Py_DECREF(iterable);
1724 if (lz->active == NULL) {
1725 Py_CLEAR(lz->source);
1726 return NULL; /* input not iterable */
1727 }
1728 }
1729 item = PyIter_Next(lz->active);
1730 if (item != NULL)
1731 return item;
1732 if (PyErr_Occurred()) {
1733 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1734 PyErr_Clear();
1735 else
1736 return NULL; /* input raised an exception */
1737 }
1738 Py_CLEAR(lz->active);
1739 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001740}
1741
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001742PyDoc_STRVAR(chain_doc,
1743"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001744\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001745Return a chain object whose .next() method returns elements from the\n\
1746first iterable until it is exhausted, then elements from the next\n\
1747iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001748
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001749PyDoc_STRVAR(chain_from_iterable_doc,
1750"chain.from_iterable(iterable) --> chain object\n\
1751\n\
1752Alternate chain() contructor taking a single iterable argument\n\
1753that evaluates lazily.");
1754
1755static PyMethodDef chain_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001756 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1757 chain_from_iterable_doc},
1758 {NULL, NULL} /* sentinel */
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001759};
1760
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001761static PyTypeObject chain_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001762 PyVarObject_HEAD_INIT(NULL, 0)
1763 "itertools.chain", /* tp_name */
1764 sizeof(chainobject), /* tp_basicsize */
1765 0, /* tp_itemsize */
1766 /* methods */
1767 (destructor)chain_dealloc, /* tp_dealloc */
1768 0, /* tp_print */
1769 0, /* tp_getattr */
1770 0, /* tp_setattr */
1771 0, /* tp_compare */
1772 0, /* tp_repr */
1773 0, /* tp_as_number */
1774 0, /* tp_as_sequence */
1775 0, /* tp_as_mapping */
1776 0, /* tp_hash */
1777 0, /* tp_call */
1778 0, /* tp_str */
1779 PyObject_GenericGetAttr, /* tp_getattro */
1780 0, /* tp_setattro */
1781 0, /* tp_as_buffer */
1782 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1783 Py_TPFLAGS_BASETYPE, /* tp_flags */
1784 chain_doc, /* tp_doc */
1785 (traverseproc)chain_traverse, /* tp_traverse */
1786 0, /* tp_clear */
1787 0, /* tp_richcompare */
1788 0, /* tp_weaklistoffset */
1789 PyObject_SelfIter, /* tp_iter */
1790 (iternextfunc)chain_next, /* tp_iternext */
1791 chain_methods, /* tp_methods */
1792 0, /* tp_members */
1793 0, /* tp_getset */
1794 0, /* tp_base */
1795 0, /* tp_dict */
1796 0, /* tp_descr_get */
1797 0, /* tp_descr_set */
1798 0, /* tp_dictoffset */
1799 0, /* tp_init */
1800 0, /* tp_alloc */
1801 chain_new, /* tp_new */
1802 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001803};
1804
1805
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001806/* product object ************************************************************/
1807
1808typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001809 PyObject_HEAD
1810 PyObject *pools; /* tuple of pool tuples */
1811 Py_ssize_t *indices; /* one index per pool */
1812 PyObject *result; /* most recently returned result tuple */
1813 int stopped; /* set to 1 when the product iterator is exhausted */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001814} productobject;
1815
1816static PyTypeObject product_type;
1817
1818static PyObject *
1819product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1820{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001821 productobject *lz;
1822 Py_ssize_t nargs, npools, repeat=1;
1823 PyObject *pools = NULL;
1824 Py_ssize_t *indices = NULL;
1825 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001826
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001827 if (kwds != NULL) {
1828 char *kwlist[] = {"repeat", 0};
1829 PyObject *tmpargs = PyTuple_New(0);
1830 if (tmpargs == NULL)
1831 return NULL;
1832 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1833 Py_DECREF(tmpargs);
1834 return NULL;
1835 }
1836 Py_DECREF(tmpargs);
1837 if (repeat < 0) {
1838 PyErr_SetString(PyExc_ValueError,
1839 "repeat argument cannot be negative");
1840 return NULL;
1841 }
1842 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001843
Benjamin Petersondda91212015-02-01 21:34:07 -05001844 assert(PyTuple_CheckExact(args));
1845 if (repeat == 0) {
1846 nargs = 0;
1847 } else {
1848 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02001849 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Petersondda91212015-02-01 21:34:07 -05001850 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
1851 return NULL;
1852 }
1853 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001854 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001855
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02001856 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001857 if (indices == NULL) {
1858 PyErr_NoMemory();
1859 goto error;
1860 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001861
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001862 pools = PyTuple_New(npools);
1863 if (pools == NULL)
1864 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001865
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001866 for (i=0; i < nargs ; ++i) {
1867 PyObject *item = PyTuple_GET_ITEM(args, i);
1868 PyObject *pool = PySequence_Tuple(item);
1869 if (pool == NULL)
1870 goto error;
1871 PyTuple_SET_ITEM(pools, i, pool);
1872 indices[i] = 0;
1873 }
1874 for ( ; i < npools; ++i) {
1875 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1876 Py_INCREF(pool);
1877 PyTuple_SET_ITEM(pools, i, pool);
1878 indices[i] = 0;
1879 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001880
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001881 /* create productobject structure */
1882 lz = (productobject *)type->tp_alloc(type, 0);
1883 if (lz == NULL)
1884 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001885
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001886 lz->pools = pools;
1887 lz->indices = indices;
1888 lz->result = NULL;
1889 lz->stopped = 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001890
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001891 return (PyObject *)lz;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001892
1893error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001894 if (indices != NULL)
1895 PyMem_Free(indices);
1896 Py_XDECREF(pools);
1897 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001898}
1899
1900static void
1901product_dealloc(productobject *lz)
1902{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001903 PyObject_GC_UnTrack(lz);
1904 Py_XDECREF(lz->pools);
1905 Py_XDECREF(lz->result);
1906 if (lz->indices != NULL)
1907 PyMem_Free(lz->indices);
1908 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001909}
1910
1911static int
1912product_traverse(productobject *lz, visitproc visit, void *arg)
1913{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001914 Py_VISIT(lz->pools);
1915 Py_VISIT(lz->result);
1916 return 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001917}
1918
1919static PyObject *
1920product_next(productobject *lz)
1921{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001922 PyObject *pool;
1923 PyObject *elem;
1924 PyObject *oldelem;
1925 PyObject *pools = lz->pools;
1926 PyObject *result = lz->result;
1927 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1928 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001929
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001930 if (lz->stopped)
1931 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001932
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001933 if (result == NULL) {
1934 /* On the first pass, return an initial tuple filled with the
1935 first element from each pool. */
1936 result = PyTuple_New(npools);
1937 if (result == NULL)
1938 goto empty;
1939 lz->result = result;
1940 for (i=0; i < npools; i++) {
1941 pool = PyTuple_GET_ITEM(pools, i);
1942 if (PyTuple_GET_SIZE(pool) == 0)
1943 goto empty;
1944 elem = PyTuple_GET_ITEM(pool, 0);
1945 Py_INCREF(elem);
1946 PyTuple_SET_ITEM(result, i, elem);
1947 }
1948 } else {
1949 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001950
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001951 /* Copy the previous result tuple or re-use it if available */
1952 if (Py_REFCNT(result) > 1) {
1953 PyObject *old_result = result;
1954 result = PyTuple_New(npools);
1955 if (result == NULL)
1956 goto empty;
1957 lz->result = result;
1958 for (i=0; i < npools; i++) {
1959 elem = PyTuple_GET_ITEM(old_result, i);
1960 Py_INCREF(elem);
1961 PyTuple_SET_ITEM(result, i, elem);
1962 }
1963 Py_DECREF(old_result);
1964 }
1965 /* Now, we've got the only copy so we can update it in-place */
1966 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001967
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001968 /* Update the pool indices right-to-left. Only advance to the
1969 next pool when the previous one rolls-over */
1970 for (i=npools-1 ; i >= 0 ; i--) {
1971 pool = PyTuple_GET_ITEM(pools, i);
1972 indices[i]++;
1973 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1974 /* Roll-over and advance to next pool */
1975 indices[i] = 0;
1976 elem = PyTuple_GET_ITEM(pool, 0);
1977 Py_INCREF(elem);
1978 oldelem = PyTuple_GET_ITEM(result, i);
1979 PyTuple_SET_ITEM(result, i, elem);
1980 Py_DECREF(oldelem);
1981 } else {
1982 /* No rollover. Just increment and stop here. */
1983 elem = PyTuple_GET_ITEM(pool, indices[i]);
1984 Py_INCREF(elem);
1985 oldelem = PyTuple_GET_ITEM(result, i);
1986 PyTuple_SET_ITEM(result, i, elem);
1987 Py_DECREF(oldelem);
1988 break;
1989 }
1990 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00001991
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001992 /* If i is negative, then the indices have all rolled-over
1993 and we're done. */
1994 if (i < 0)
1995 goto empty;
1996 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001997
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001998 Py_INCREF(result);
1999 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002000
2001empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002002 lz->stopped = 1;
2003 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002004}
2005
2006PyDoc_STRVAR(product_doc,
2007"product(*iterables) --> product object\n\
2008\n\
Raymond Hettinger73d79632008-02-23 02:20:41 +00002009Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002010For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2011The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2012cycle in a manner similar to an odometer (with the rightmost element changing\n\
2013on every iteration).\n\n\
Raymond Hettinger94034ea2009-02-09 18:39:41 +00002014To compute the product of an iterable with itself, specify the number\n\
2015of repetitions with the optional repeat keyword argument. For example,\n\
2016product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002017product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2018product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2019
2020static PyTypeObject product_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002021 PyVarObject_HEAD_INIT(NULL, 0)
2022 "itertools.product", /* tp_name */
2023 sizeof(productobject), /* tp_basicsize */
2024 0, /* tp_itemsize */
2025 /* methods */
2026 (destructor)product_dealloc, /* tp_dealloc */
2027 0, /* tp_print */
2028 0, /* tp_getattr */
2029 0, /* tp_setattr */
2030 0, /* tp_compare */
2031 0, /* tp_repr */
2032 0, /* tp_as_number */
2033 0, /* tp_as_sequence */
2034 0, /* tp_as_mapping */
2035 0, /* tp_hash */
2036 0, /* tp_call */
2037 0, /* tp_str */
2038 PyObject_GenericGetAttr, /* tp_getattro */
2039 0, /* tp_setattro */
2040 0, /* tp_as_buffer */
2041 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2042 Py_TPFLAGS_BASETYPE, /* tp_flags */
2043 product_doc, /* tp_doc */
2044 (traverseproc)product_traverse, /* tp_traverse */
2045 0, /* tp_clear */
2046 0, /* tp_richcompare */
2047 0, /* tp_weaklistoffset */
2048 PyObject_SelfIter, /* tp_iter */
2049 (iternextfunc)product_next, /* tp_iternext */
2050 0, /* tp_methods */
2051 0, /* tp_members */
2052 0, /* tp_getset */
2053 0, /* tp_base */
2054 0, /* tp_dict */
2055 0, /* tp_descr_get */
2056 0, /* tp_descr_set */
2057 0, /* tp_dictoffset */
2058 0, /* tp_init */
2059 0, /* tp_alloc */
2060 product_new, /* tp_new */
2061 PyObject_GC_Del, /* tp_free */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002062};
2063
2064
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002065/* combinations object ************************************************************/
2066
2067typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002068 PyObject_HEAD
2069 PyObject *pool; /* input converted to a tuple */
2070 Py_ssize_t *indices; /* one index per result element */
2071 PyObject *result; /* most recently returned result tuple */
2072 Py_ssize_t r; /* size of result tuple */
2073 int stopped; /* set to 1 when the combinations iterator is exhausted */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002074} combinationsobject;
2075
2076static PyTypeObject combinations_type;
2077
2078static PyObject *
2079combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2080{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002081 combinationsobject *co;
2082 Py_ssize_t n;
2083 Py_ssize_t r;
2084 PyObject *pool = NULL;
2085 PyObject *iterable = NULL;
2086 Py_ssize_t *indices = NULL;
2087 Py_ssize_t i;
2088 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002089
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002090 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2091 &iterable, &r))
2092 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002093
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002094 pool = PySequence_Tuple(iterable);
2095 if (pool == NULL)
2096 goto error;
2097 n = PyTuple_GET_SIZE(pool);
2098 if (r < 0) {
2099 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2100 goto error;
2101 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002102
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02002103 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002104 if (indices == NULL) {
2105 PyErr_NoMemory();
2106 goto error;
2107 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002108
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002109 for (i=0 ; i<r ; i++)
2110 indices[i] = i;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002111
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002112 /* create combinationsobject structure */
2113 co = (combinationsobject *)type->tp_alloc(type, 0);
2114 if (co == NULL)
2115 goto error;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002116
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002117 co->pool = pool;
2118 co->indices = indices;
2119 co->result = NULL;
2120 co->r = r;
2121 co->stopped = r > n ? 1 : 0;
2122
2123 return (PyObject *)co;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002124
2125error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002126 if (indices != NULL)
2127 PyMem_Free(indices);
2128 Py_XDECREF(pool);
2129 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002130}
2131
2132static void
2133combinations_dealloc(combinationsobject *co)
2134{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002135 PyObject_GC_UnTrack(co);
2136 Py_XDECREF(co->pool);
2137 Py_XDECREF(co->result);
2138 if (co->indices != NULL)
2139 PyMem_Free(co->indices);
2140 Py_TYPE(co)->tp_free(co);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002141}
2142
2143static int
2144combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2145{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002146 Py_VISIT(co->pool);
2147 Py_VISIT(co->result);
2148 return 0;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002149}
2150
2151static PyObject *
2152combinations_next(combinationsobject *co)
2153{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002154 PyObject *elem;
2155 PyObject *oldelem;
2156 PyObject *pool = co->pool;
2157 Py_ssize_t *indices = co->indices;
2158 PyObject *result = co->result;
2159 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2160 Py_ssize_t r = co->r;
2161 Py_ssize_t i, j, index;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002162
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002163 if (co->stopped)
2164 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002165
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002166 if (result == NULL) {
2167 /* On the first pass, initialize result tuple using the indices */
2168 result = PyTuple_New(r);
2169 if (result == NULL)
2170 goto empty;
2171 co->result = result;
2172 for (i=0; i<r ; i++) {
2173 index = indices[i];
2174 elem = PyTuple_GET_ITEM(pool, index);
2175 Py_INCREF(elem);
2176 PyTuple_SET_ITEM(result, i, elem);
2177 }
2178 } else {
2179 /* Copy the previous result tuple or re-use it if available */
2180 if (Py_REFCNT(result) > 1) {
2181 PyObject *old_result = result;
2182 result = PyTuple_New(r);
2183 if (result == NULL)
2184 goto empty;
2185 co->result = result;
2186 for (i=0; i<r ; i++) {
2187 elem = PyTuple_GET_ITEM(old_result, i);
2188 Py_INCREF(elem);
2189 PyTuple_SET_ITEM(result, i, elem);
2190 }
2191 Py_DECREF(old_result);
2192 }
2193 /* Now, we've got the only copy so we can update it in-place
2194 * CPython's empty tuple is a singleton and cached in
2195 * PyTuple's freelist.
2196 */
2197 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002198
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002199 /* Scan indices right-to-left until finding one that is not
2200 at its maximum (i + n - r). */
2201 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2202 ;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002203
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002204 /* If i is negative, then the indices are all at
2205 their maximum value and we're done. */
2206 if (i < 0)
2207 goto empty;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002208
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002209 /* Increment the current index which we know is not at its
2210 maximum. Then move back to the right setting each index
2211 to its lowest possible value (one higher than the index
2212 to its left -- this maintains the sort order invariant). */
2213 indices[i]++;
2214 for (j=i+1 ; j<r ; j++)
2215 indices[j] = indices[j-1] + 1;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002216
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002217 /* Update the result tuple for the new indices
2218 starting with i, the leftmost index that changed */
2219 for ( ; i<r ; i++) {
2220 index = indices[i];
2221 elem = PyTuple_GET_ITEM(pool, index);
2222 Py_INCREF(elem);
2223 oldelem = PyTuple_GET_ITEM(result, i);
2224 PyTuple_SET_ITEM(result, i, elem);
2225 Py_DECREF(oldelem);
2226 }
2227 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002228
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002229 Py_INCREF(result);
2230 return result;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002231
2232empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002233 co->stopped = 1;
2234 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002235}
2236
2237PyDoc_STRVAR(combinations_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002238"combinations(iterable, r) --> combinations object\n\
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002239\n\
2240Return successive r-length combinations of elements in the iterable.\n\n\
2241combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2242
2243static PyTypeObject combinations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002244 PyVarObject_HEAD_INIT(NULL, 0)
2245 "itertools.combinations", /* tp_name */
2246 sizeof(combinationsobject), /* tp_basicsize */
2247 0, /* tp_itemsize */
2248 /* methods */
2249 (destructor)combinations_dealloc, /* tp_dealloc */
2250 0, /* tp_print */
2251 0, /* tp_getattr */
2252 0, /* tp_setattr */
2253 0, /* tp_compare */
2254 0, /* tp_repr */
2255 0, /* tp_as_number */
2256 0, /* tp_as_sequence */
2257 0, /* tp_as_mapping */
2258 0, /* tp_hash */
2259 0, /* tp_call */
2260 0, /* tp_str */
2261 PyObject_GenericGetAttr, /* tp_getattro */
2262 0, /* tp_setattro */
2263 0, /* tp_as_buffer */
2264 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2265 Py_TPFLAGS_BASETYPE, /* tp_flags */
2266 combinations_doc, /* tp_doc */
2267 (traverseproc)combinations_traverse, /* tp_traverse */
2268 0, /* tp_clear */
2269 0, /* tp_richcompare */
2270 0, /* tp_weaklistoffset */
2271 PyObject_SelfIter, /* tp_iter */
2272 (iternextfunc)combinations_next, /* tp_iternext */
2273 0, /* tp_methods */
2274 0, /* tp_members */
2275 0, /* tp_getset */
2276 0, /* tp_base */
2277 0, /* tp_dict */
2278 0, /* tp_descr_get */
2279 0, /* tp_descr_set */
2280 0, /* tp_dictoffset */
2281 0, /* tp_init */
2282 0, /* tp_alloc */
2283 combinations_new, /* tp_new */
2284 PyObject_GC_Del, /* tp_free */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002285};
2286
2287
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002288/* combinations with replacement object *******************************************/
2289
2290/* Equivalent to:
2291
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002292 def combinations_with_replacement(iterable, r):
2293 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2294 # number items returned: (n+r-1)! / r! / (n-1)!
2295 pool = tuple(iterable)
2296 n = len(pool)
2297 indices = [0] * r
2298 yield tuple(pool[i] for i in indices)
2299 while 1:
2300 for i in reversed(range(r)):
2301 if indices[i] != n - 1:
2302 break
2303 else:
2304 return
2305 indices[i:] = [indices[i] + 1] * (r - i)
2306 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002307
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002308 def combinations_with_replacement2(iterable, r):
2309 'Alternate version that filters from product()'
2310 pool = tuple(iterable)
2311 n = len(pool)
2312 for indices in product(range(n), repeat=r):
2313 if sorted(indices) == list(indices):
2314 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002315*/
2316typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002317 PyObject_HEAD
2318 PyObject *pool; /* input converted to a tuple */
2319 Py_ssize_t *indices; /* one index per result element */
2320 PyObject *result; /* most recently returned result tuple */
2321 Py_ssize_t r; /* size of result tuple */
2322 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002323} cwrobject;
2324
2325static PyTypeObject cwr_type;
2326
2327static PyObject *
2328cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2329{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002330 cwrobject *co;
2331 Py_ssize_t n;
2332 Py_ssize_t r;
2333 PyObject *pool = NULL;
2334 PyObject *iterable = NULL;
2335 Py_ssize_t *indices = NULL;
2336 Py_ssize_t i;
2337 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002338
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002339 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2340 &iterable, &r))
2341 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002342
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002343 pool = PySequence_Tuple(iterable);
2344 if (pool == NULL)
2345 goto error;
2346 n = PyTuple_GET_SIZE(pool);
2347 if (r < 0) {
2348 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2349 goto error;
2350 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002351
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02002352 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002353 if (indices == NULL) {
2354 PyErr_NoMemory();
2355 goto error;
2356 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002357
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002358 for (i=0 ; i<r ; i++)
2359 indices[i] = 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002360
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002361 /* create cwrobject structure */
2362 co = (cwrobject *)type->tp_alloc(type, 0);
2363 if (co == NULL)
2364 goto error;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002365
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002366 co->pool = pool;
2367 co->indices = indices;
2368 co->result = NULL;
2369 co->r = r;
2370 co->stopped = !n && r;
2371
2372 return (PyObject *)co;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002373
2374error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002375 if (indices != NULL)
2376 PyMem_Free(indices);
2377 Py_XDECREF(pool);
2378 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002379}
2380
2381static void
2382cwr_dealloc(cwrobject *co)
2383{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002384 PyObject_GC_UnTrack(co);
2385 Py_XDECREF(co->pool);
2386 Py_XDECREF(co->result);
2387 if (co->indices != NULL)
2388 PyMem_Free(co->indices);
2389 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002390}
2391
2392static int
2393cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2394{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002395 Py_VISIT(co->pool);
2396 Py_VISIT(co->result);
2397 return 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002398}
2399
2400static PyObject *
2401cwr_next(cwrobject *co)
2402{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002403 PyObject *elem;
2404 PyObject *oldelem;
2405 PyObject *pool = co->pool;
2406 Py_ssize_t *indices = co->indices;
2407 PyObject *result = co->result;
2408 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2409 Py_ssize_t r = co->r;
2410 Py_ssize_t i, j, index;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002411
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002412 if (co->stopped)
2413 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002414
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002415 if (result == NULL) {
2416 /* On the first pass, initialize result tuple using the indices */
2417 result = PyTuple_New(r);
2418 if (result == NULL)
2419 goto empty;
2420 co->result = result;
2421 for (i=0; i<r ; i++) {
2422 index = indices[i];
2423 elem = PyTuple_GET_ITEM(pool, index);
2424 Py_INCREF(elem);
2425 PyTuple_SET_ITEM(result, i, elem);
2426 }
2427 } else {
2428 /* Copy the previous result tuple or re-use it if available */
2429 if (Py_REFCNT(result) > 1) {
2430 PyObject *old_result = result;
2431 result = PyTuple_New(r);
2432 if (result == NULL)
2433 goto empty;
2434 co->result = result;
2435 for (i=0; i<r ; i++) {
2436 elem = PyTuple_GET_ITEM(old_result, i);
2437 Py_INCREF(elem);
2438 PyTuple_SET_ITEM(result, i, elem);
2439 }
2440 Py_DECREF(old_result);
2441 }
2442 /* Now, we've got the only copy so we can update it in-place CPython's
2443 empty tuple is a singleton and cached in PyTuple's freelist. */
2444 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002445
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002446 /* Scan indices right-to-left until finding one that is not
2447 * at its maximum (n-1). */
2448 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2449 ;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002450
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002451 /* If i is negative, then the indices are all at
2452 their maximum value and we're done. */
2453 if (i < 0)
2454 goto empty;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002455
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002456 /* Increment the current index which we know is not at its
2457 maximum. Then set all to the right to the same value. */
2458 indices[i]++;
2459 for (j=i+1 ; j<r ; j++)
2460 indices[j] = indices[j-1];
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002461
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002462 /* Update the result tuple for the new indices
2463 starting with i, the leftmost index that changed */
2464 for ( ; i<r ; i++) {
2465 index = indices[i];
2466 elem = PyTuple_GET_ITEM(pool, index);
2467 Py_INCREF(elem);
2468 oldelem = PyTuple_GET_ITEM(result, i);
2469 PyTuple_SET_ITEM(result, i, elem);
2470 Py_DECREF(oldelem);
2471 }
2472 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002473
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002474 Py_INCREF(result);
2475 return result;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002476
2477empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002478 co->stopped = 1;
2479 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002480}
2481
2482PyDoc_STRVAR(cwr_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002483"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002484\n\
2485Return successive r-length combinations of elements in the iterable\n\
2486allowing individual elements to have successive repeats.\n\
2487combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2488
2489static PyTypeObject cwr_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002490 PyVarObject_HEAD_INIT(NULL, 0)
2491 "itertools.combinations_with_replacement", /* tp_name */
2492 sizeof(cwrobject), /* tp_basicsize */
2493 0, /* tp_itemsize */
2494 /* methods */
2495 (destructor)cwr_dealloc, /* tp_dealloc */
2496 0, /* tp_print */
2497 0, /* tp_getattr */
2498 0, /* tp_setattr */
2499 0, /* tp_compare */
2500 0, /* tp_repr */
2501 0, /* tp_as_number */
2502 0, /* tp_as_sequence */
2503 0, /* tp_as_mapping */
2504 0, /* tp_hash */
2505 0, /* tp_call */
2506 0, /* tp_str */
2507 PyObject_GenericGetAttr, /* tp_getattro */
2508 0, /* tp_setattro */
2509 0, /* tp_as_buffer */
2510 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2511 Py_TPFLAGS_BASETYPE, /* tp_flags */
2512 cwr_doc, /* tp_doc */
2513 (traverseproc)cwr_traverse, /* tp_traverse */
2514 0, /* tp_clear */
2515 0, /* tp_richcompare */
2516 0, /* tp_weaklistoffset */
2517 PyObject_SelfIter, /* tp_iter */
2518 (iternextfunc)cwr_next, /* tp_iternext */
2519 0, /* tp_methods */
2520 0, /* tp_members */
2521 0, /* tp_getset */
2522 0, /* tp_base */
2523 0, /* tp_dict */
2524 0, /* tp_descr_get */
2525 0, /* tp_descr_set */
2526 0, /* tp_dictoffset */
2527 0, /* tp_init */
2528 0, /* tp_alloc */
2529 cwr_new, /* tp_new */
2530 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002531};
2532
2533
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002534/* permutations object ************************************************************
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002535
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002536def permutations(iterable, r=None):
2537 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2538 pool = tuple(iterable)
2539 n = len(pool)
2540 r = n if r is None else r
2541 indices = range(n)
2542 cycles = range(n-r+1, n+1)[::-1]
2543 yield tuple(pool[i] for i in indices[:r])
2544 while n:
Serhiy Storchakaf0aa88f2015-06-11 00:06:27 +03002545 for i in reversed(range(r)):
2546 cycles[i] -= 1
2547 if cycles[i] == 0:
2548 indices[i:] = indices[i+1:] + indices[i:i+1]
2549 cycles[i] = n - i
2550 else:
2551 j = cycles[i]
2552 indices[i], indices[-j] = indices[-j], indices[i]
2553 yield tuple(pool[i] for i in indices[:r])
2554 break
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002555 else:
Serhiy Storchakaf0aa88f2015-06-11 00:06:27 +03002556 return
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002557*/
2558
2559typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002560 PyObject_HEAD
2561 PyObject *pool; /* input converted to a tuple */
2562 Py_ssize_t *indices; /* one index per element in the pool */
2563 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2564 PyObject *result; /* most recently returned result tuple */
2565 Py_ssize_t r; /* size of result tuple */
2566 int stopped; /* set to 1 when the permutations iterator is exhausted */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002567} permutationsobject;
2568
2569static PyTypeObject permutations_type;
2570
2571static PyObject *
2572permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2573{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002574 permutationsobject *po;
2575 Py_ssize_t n;
2576 Py_ssize_t r;
2577 PyObject *robj = Py_None;
2578 PyObject *pool = NULL;
2579 PyObject *iterable = NULL;
2580 Py_ssize_t *indices = NULL;
2581 Py_ssize_t *cycles = NULL;
2582 Py_ssize_t i;
2583 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002584
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002585 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2586 &iterable, &robj))
2587 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002588
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002589 pool = PySequence_Tuple(iterable);
2590 if (pool == NULL)
2591 goto error;
2592 n = PyTuple_GET_SIZE(pool);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002593
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002594 r = n;
2595 if (robj != Py_None) {
2596 r = PyInt_AsSsize_t(robj);
2597 if (r == -1 && PyErr_Occurred())
2598 goto error;
2599 }
2600 if (r < 0) {
2601 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2602 goto error;
2603 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002604
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02002605 indices = PyMem_New(Py_ssize_t, n);
2606 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002607 if (indices == NULL || cycles == NULL) {
2608 PyErr_NoMemory();
2609 goto error;
2610 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002611
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002612 for (i=0 ; i<n ; i++)
2613 indices[i] = i;
2614 for (i=0 ; i<r ; i++)
2615 cycles[i] = n - i;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002616
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002617 /* create permutationsobject structure */
2618 po = (permutationsobject *)type->tp_alloc(type, 0);
2619 if (po == NULL)
2620 goto error;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002621
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002622 po->pool = pool;
2623 po->indices = indices;
2624 po->cycles = cycles;
2625 po->result = NULL;
2626 po->r = r;
2627 po->stopped = r > n ? 1 : 0;
2628
2629 return (PyObject *)po;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002630
2631error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002632 if (indices != NULL)
2633 PyMem_Free(indices);
2634 if (cycles != NULL)
2635 PyMem_Free(cycles);
2636 Py_XDECREF(pool);
2637 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002638}
2639
2640static void
2641permutations_dealloc(permutationsobject *po)
2642{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002643 PyObject_GC_UnTrack(po);
2644 Py_XDECREF(po->pool);
2645 Py_XDECREF(po->result);
2646 PyMem_Free(po->indices);
2647 PyMem_Free(po->cycles);
2648 Py_TYPE(po)->tp_free(po);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002649}
2650
2651static int
2652permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2653{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002654 Py_VISIT(po->pool);
2655 Py_VISIT(po->result);
2656 return 0;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002657}
2658
2659static PyObject *
2660permutations_next(permutationsobject *po)
2661{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002662 PyObject *elem;
2663 PyObject *oldelem;
2664 PyObject *pool = po->pool;
2665 Py_ssize_t *indices = po->indices;
2666 Py_ssize_t *cycles = po->cycles;
2667 PyObject *result = po->result;
2668 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2669 Py_ssize_t r = po->r;
2670 Py_ssize_t i, j, k, index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002671
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002672 if (po->stopped)
2673 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002674
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002675 if (result == NULL) {
2676 /* On the first pass, initialize result tuple using the indices */
2677 result = PyTuple_New(r);
2678 if (result == NULL)
2679 goto empty;
2680 po->result = result;
2681 for (i=0; i<r ; i++) {
2682 index = indices[i];
2683 elem = PyTuple_GET_ITEM(pool, index);
2684 Py_INCREF(elem);
2685 PyTuple_SET_ITEM(result, i, elem);
2686 }
2687 } else {
2688 if (n == 0)
2689 goto empty;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002690
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002691 /* Copy the previous result tuple or re-use it if available */
2692 if (Py_REFCNT(result) > 1) {
2693 PyObject *old_result = result;
2694 result = PyTuple_New(r);
2695 if (result == NULL)
2696 goto empty;
2697 po->result = result;
2698 for (i=0; i<r ; i++) {
2699 elem = PyTuple_GET_ITEM(old_result, i);
2700 Py_INCREF(elem);
2701 PyTuple_SET_ITEM(result, i, elem);
2702 }
2703 Py_DECREF(old_result);
2704 }
2705 /* Now, we've got the only copy so we can update it in-place */
2706 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002707
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002708 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2709 for (i=r-1 ; i>=0 ; i--) {
2710 cycles[i] -= 1;
2711 if (cycles[i] == 0) {
2712 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2713 index = indices[i];
2714 for (j=i ; j<n-1 ; j++)
2715 indices[j] = indices[j+1];
2716 indices[n-1] = index;
2717 cycles[i] = n - i;
2718 } else {
2719 j = cycles[i];
2720 index = indices[i];
2721 indices[i] = indices[n-j];
2722 indices[n-j] = index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002723
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002724 for (k=i; k<r ; k++) {
2725 /* start with i, the leftmost element that changed */
2726 /* yield tuple(pool[k] for k in indices[:r]) */
2727 index = indices[k];
2728 elem = PyTuple_GET_ITEM(pool, index);
2729 Py_INCREF(elem);
2730 oldelem = PyTuple_GET_ITEM(result, k);
2731 PyTuple_SET_ITEM(result, k, elem);
2732 Py_DECREF(oldelem);
2733 }
2734 break;
2735 }
2736 }
2737 /* If i is negative, then the cycles have all
2738 rolled-over and we're done. */
2739 if (i < 0)
2740 goto empty;
2741 }
2742 Py_INCREF(result);
2743 return result;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002744
2745empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002746 po->stopped = 1;
2747 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002748}
2749
2750PyDoc_STRVAR(permutations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +00002751"permutations(iterable[, r]) --> permutations object\n\
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002752\n\
2753Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl8550d692008-06-10 12:46:39 +00002754permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002755
2756static PyTypeObject permutations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002757 PyVarObject_HEAD_INIT(NULL, 0)
2758 "itertools.permutations", /* tp_name */
2759 sizeof(permutationsobject), /* tp_basicsize */
2760 0, /* tp_itemsize */
2761 /* methods */
2762 (destructor)permutations_dealloc, /* tp_dealloc */
2763 0, /* tp_print */
2764 0, /* tp_getattr */
2765 0, /* tp_setattr */
2766 0, /* tp_compare */
2767 0, /* tp_repr */
2768 0, /* tp_as_number */
2769 0, /* tp_as_sequence */
2770 0, /* tp_as_mapping */
2771 0, /* tp_hash */
2772 0, /* tp_call */
2773 0, /* tp_str */
2774 PyObject_GenericGetAttr, /* tp_getattro */
2775 0, /* tp_setattro */
2776 0, /* tp_as_buffer */
2777 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2778 Py_TPFLAGS_BASETYPE, /* tp_flags */
2779 permutations_doc, /* tp_doc */
2780 (traverseproc)permutations_traverse, /* tp_traverse */
2781 0, /* tp_clear */
2782 0, /* tp_richcompare */
2783 0, /* tp_weaklistoffset */
2784 PyObject_SelfIter, /* tp_iter */
2785 (iternextfunc)permutations_next, /* tp_iternext */
2786 0, /* tp_methods */
2787 0, /* tp_members */
2788 0, /* tp_getset */
2789 0, /* tp_base */
2790 0, /* tp_dict */
2791 0, /* tp_descr_get */
2792 0, /* tp_descr_set */
2793 0, /* tp_dictoffset */
2794 0, /* tp_init */
2795 0, /* tp_alloc */
2796 permutations_new, /* tp_new */
2797 PyObject_GC_Del, /* tp_free */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002798};
2799
2800
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002801/* compress object ************************************************************/
2802
2803/* Equivalent to:
2804
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002805 def compress(data, selectors):
2806 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2807 return (d for d, s in izip(data, selectors) if s)
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002808*/
2809
2810typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002811 PyObject_HEAD
2812 PyObject *data;
2813 PyObject *selectors;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002814} compressobject;
2815
2816static PyTypeObject compress_type;
2817
2818static PyObject *
2819compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2820{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002821 PyObject *seq1, *seq2;
2822 PyObject *data=NULL, *selectors=NULL;
2823 compressobject *lz;
2824 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002825
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002826 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2827 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002828
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002829 data = PyObject_GetIter(seq1);
2830 if (data == NULL)
2831 goto fail;
2832 selectors = PyObject_GetIter(seq2);
2833 if (selectors == NULL)
2834 goto fail;
2835
2836 /* create compressobject structure */
2837 lz = (compressobject *)type->tp_alloc(type, 0);
2838 if (lz == NULL)
2839 goto fail;
2840 lz->data = data;
2841 lz->selectors = selectors;
2842 return (PyObject *)lz;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002843
2844fail:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002845 Py_XDECREF(data);
2846 Py_XDECREF(selectors);
2847 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002848}
2849
2850static void
2851compress_dealloc(compressobject *lz)
2852{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002853 PyObject_GC_UnTrack(lz);
2854 Py_XDECREF(lz->data);
2855 Py_XDECREF(lz->selectors);
2856 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002857}
2858
2859static int
2860compress_traverse(compressobject *lz, visitproc visit, void *arg)
2861{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002862 Py_VISIT(lz->data);
2863 Py_VISIT(lz->selectors);
2864 return 0;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002865}
2866
2867static PyObject *
2868compress_next(compressobject *lz)
2869{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002870 PyObject *data = lz->data, *selectors = lz->selectors;
2871 PyObject *datum, *selector;
2872 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2873 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2874 int ok;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002875
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002876 while (1) {
2877 /* Steps: get datum, get selector, evaluate selector.
2878 Order is important (to match the pure python version
2879 in terms of which input gets a chance to raise an
2880 exception first).
2881 */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002882
Serhiy Storchakabb845652013-04-06 22:04:10 +03002883 datum = datanext(data);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03002884 if (datum == NULL)
Serhiy Storchakabb845652013-04-06 22:04:10 +03002885 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002886
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002887 selector = selectornext(selectors);
2888 if (selector == NULL) {
2889 Py_DECREF(datum);
2890 return NULL;
2891 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002892
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002893 ok = PyObject_IsTrue(selector);
2894 Py_DECREF(selector);
2895 if (ok == 1)
2896 return datum;
2897 Py_DECREF(datum);
2898 if (ok == -1)
2899 return NULL;
2900 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002901}
2902
2903PyDoc_STRVAR(compress_doc,
Raymond Hettinger2e2909f2009-02-19 02:15:14 +00002904"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002905\n\
2906Return data elements corresponding to true selector elements.\n\
2907Forms a shorter iterator from selected data elements using the\n\
2908selectors to choose the data elements.");
2909
2910static PyTypeObject compress_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002911 PyVarObject_HEAD_INIT(NULL, 0)
2912 "itertools.compress", /* tp_name */
2913 sizeof(compressobject), /* tp_basicsize */
2914 0, /* tp_itemsize */
2915 /* methods */
2916 (destructor)compress_dealloc, /* tp_dealloc */
2917 0, /* tp_print */
2918 0, /* tp_getattr */
2919 0, /* tp_setattr */
2920 0, /* tp_compare */
2921 0, /* tp_repr */
2922 0, /* tp_as_number */
2923 0, /* tp_as_sequence */
2924 0, /* tp_as_mapping */
2925 0, /* tp_hash */
2926 0, /* tp_call */
2927 0, /* tp_str */
2928 PyObject_GenericGetAttr, /* tp_getattro */
2929 0, /* tp_setattro */
2930 0, /* tp_as_buffer */
2931 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2932 Py_TPFLAGS_BASETYPE, /* tp_flags */
2933 compress_doc, /* tp_doc */
2934 (traverseproc)compress_traverse, /* tp_traverse */
2935 0, /* tp_clear */
2936 0, /* tp_richcompare */
2937 0, /* tp_weaklistoffset */
2938 PyObject_SelfIter, /* tp_iter */
2939 (iternextfunc)compress_next, /* tp_iternext */
2940 0, /* tp_methods */
2941 0, /* tp_members */
2942 0, /* tp_getset */
2943 0, /* tp_base */
2944 0, /* tp_dict */
2945 0, /* tp_descr_get */
2946 0, /* tp_descr_set */
2947 0, /* tp_dictoffset */
2948 0, /* tp_init */
2949 0, /* tp_alloc */
2950 compress_new, /* tp_new */
2951 PyObject_GC_Del, /* tp_free */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002952};
2953
2954
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002955/* ifilter object ************************************************************/
2956
2957typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002958 PyObject_HEAD
2959 PyObject *func;
2960 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002961} ifilterobject;
2962
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002963static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002964
2965static PyObject *
2966ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2967{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002968 PyObject *func, *seq;
2969 PyObject *it;
2970 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002971
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002972 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
2973 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002974
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002975 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
2976 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002977
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002978 /* Get iterator. */
2979 it = PyObject_GetIter(seq);
2980 if (it == NULL)
2981 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002982
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002983 /* create ifilterobject structure */
2984 lz = (ifilterobject *)type->tp_alloc(type, 0);
2985 if (lz == NULL) {
2986 Py_DECREF(it);
2987 return NULL;
2988 }
2989 Py_INCREF(func);
2990 lz->func = func;
2991 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002992
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002993 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002994}
2995
2996static void
2997ifilter_dealloc(ifilterobject *lz)
2998{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002999 PyObject_GC_UnTrack(lz);
3000 Py_XDECREF(lz->func);
3001 Py_XDECREF(lz->it);
3002 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003003}
3004
3005static int
3006ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
3007{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003008 Py_VISIT(lz->it);
3009 Py_VISIT(lz->func);
3010 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003011}
3012
3013static PyObject *
3014ifilter_next(ifilterobject *lz)
3015{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003016 PyObject *item;
3017 PyObject *it = lz->it;
3018 long ok;
3019 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003020
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003021 iternext = *Py_TYPE(it)->tp_iternext;
3022 for (;;) {
3023 item = iternext(it);
3024 if (item == NULL)
3025 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003026
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003027 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3028 ok = PyObject_IsTrue(item);
3029 } else {
3030 PyObject *good;
3031 good = PyObject_CallFunctionObjArgs(lz->func,
3032 item, NULL);
3033 if (good == NULL) {
3034 Py_DECREF(item);
3035 return NULL;
3036 }
3037 ok = PyObject_IsTrue(good);
3038 Py_DECREF(good);
3039 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003040 if (ok > 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003041 return item;
3042 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003043 if (ok < 0)
3044 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003045 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003046}
3047
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003048PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003049"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003050\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003051Return those items of sequence for which function(item) is true.\n\
3052If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003053
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003054static PyTypeObject ifilter_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003055 PyVarObject_HEAD_INIT(NULL, 0)
3056 "itertools.ifilter", /* tp_name */
3057 sizeof(ifilterobject), /* tp_basicsize */
3058 0, /* tp_itemsize */
3059 /* methods */
3060 (destructor)ifilter_dealloc, /* tp_dealloc */
3061 0, /* tp_print */
3062 0, /* tp_getattr */
3063 0, /* tp_setattr */
3064 0, /* tp_compare */
3065 0, /* tp_repr */
3066 0, /* tp_as_number */
3067 0, /* tp_as_sequence */
3068 0, /* tp_as_mapping */
3069 0, /* tp_hash */
3070 0, /* tp_call */
3071 0, /* tp_str */
3072 PyObject_GenericGetAttr, /* tp_getattro */
3073 0, /* tp_setattro */
3074 0, /* tp_as_buffer */
3075 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3076 Py_TPFLAGS_BASETYPE, /* tp_flags */
3077 ifilter_doc, /* tp_doc */
3078 (traverseproc)ifilter_traverse, /* tp_traverse */
3079 0, /* tp_clear */
3080 0, /* tp_richcompare */
3081 0, /* tp_weaklistoffset */
3082 PyObject_SelfIter, /* tp_iter */
3083 (iternextfunc)ifilter_next, /* tp_iternext */
3084 0, /* tp_methods */
3085 0, /* tp_members */
3086 0, /* tp_getset */
3087 0, /* tp_base */
3088 0, /* tp_dict */
3089 0, /* tp_descr_get */
3090 0, /* tp_descr_set */
3091 0, /* tp_dictoffset */
3092 0, /* tp_init */
3093 0, /* tp_alloc */
3094 ifilter_new, /* tp_new */
3095 PyObject_GC_Del, /* tp_free */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003096};
3097
3098
3099/* ifilterfalse object ************************************************************/
3100
3101typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003102 PyObject_HEAD
3103 PyObject *func;
3104 PyObject *it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003105} ifilterfalseobject;
3106
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003107static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003108
3109static PyObject *
3110ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3111{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003112 PyObject *func, *seq;
3113 PyObject *it;
3114 ifilterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003115
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003116 if (type == &ifilterfalse_type &&
3117 !_PyArg_NoKeywords("ifilterfalse()", kwds))
3118 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003119
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003120 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
3121 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003122
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003123 /* Get iterator. */
3124 it = PyObject_GetIter(seq);
3125 if (it == NULL)
3126 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003127
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003128 /* create ifilterfalseobject structure */
3129 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
3130 if (lz == NULL) {
3131 Py_DECREF(it);
3132 return NULL;
3133 }
3134 Py_INCREF(func);
3135 lz->func = func;
3136 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003137
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003138 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003139}
3140
3141static void
3142ifilterfalse_dealloc(ifilterfalseobject *lz)
3143{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003144 PyObject_GC_UnTrack(lz);
3145 Py_XDECREF(lz->func);
3146 Py_XDECREF(lz->it);
3147 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003148}
3149
3150static int
3151ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
3152{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003153 Py_VISIT(lz->it);
3154 Py_VISIT(lz->func);
3155 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003156}
3157
3158static PyObject *
3159ifilterfalse_next(ifilterfalseobject *lz)
3160{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003161 PyObject *item;
3162 PyObject *it = lz->it;
3163 long ok;
3164 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003165
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003166 iternext = *Py_TYPE(it)->tp_iternext;
3167 for (;;) {
3168 item = iternext(it);
3169 if (item == NULL)
3170 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003171
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003172 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3173 ok = PyObject_IsTrue(item);
3174 } else {
3175 PyObject *good;
3176 good = PyObject_CallFunctionObjArgs(lz->func,
3177 item, NULL);
3178 if (good == NULL) {
3179 Py_DECREF(item);
3180 return NULL;
3181 }
3182 ok = PyObject_IsTrue(good);
3183 Py_DECREF(good);
3184 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003185 if (ok == 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003186 return item;
3187 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003188 if (ok < 0)
3189 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003190 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003191}
3192
Raymond Hettinger60eca932003-02-09 06:40:58 +00003193PyDoc_STRVAR(ifilterfalse_doc,
3194"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
3195\n\
3196Return those items of sequence for which function(item) is false.\n\
3197If function is None, return the items that are false.");
3198
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003199static PyTypeObject ifilterfalse_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003200 PyVarObject_HEAD_INIT(NULL, 0)
3201 "itertools.ifilterfalse", /* tp_name */
3202 sizeof(ifilterfalseobject), /* tp_basicsize */
3203 0, /* tp_itemsize */
3204 /* methods */
3205 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
3206 0, /* tp_print */
3207 0, /* tp_getattr */
3208 0, /* tp_setattr */
3209 0, /* tp_compare */
3210 0, /* tp_repr */
3211 0, /* tp_as_number */
3212 0, /* tp_as_sequence */
3213 0, /* tp_as_mapping */
3214 0, /* tp_hash */
3215 0, /* tp_call */
3216 0, /* tp_str */
3217 PyObject_GenericGetAttr, /* tp_getattro */
3218 0, /* tp_setattro */
3219 0, /* tp_as_buffer */
3220 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3221 Py_TPFLAGS_BASETYPE, /* tp_flags */
3222 ifilterfalse_doc, /* tp_doc */
3223 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
3224 0, /* tp_clear */
3225 0, /* tp_richcompare */
3226 0, /* tp_weaklistoffset */
3227 PyObject_SelfIter, /* tp_iter */
3228 (iternextfunc)ifilterfalse_next, /* tp_iternext */
3229 0, /* tp_methods */
3230 0, /* tp_members */
3231 0, /* tp_getset */
3232 0, /* tp_base */
3233 0, /* tp_dict */
3234 0, /* tp_descr_get */
3235 0, /* tp_descr_set */
3236 0, /* tp_dictoffset */
3237 0, /* tp_init */
3238 0, /* tp_alloc */
3239 ifilterfalse_new, /* tp_new */
3240 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003241};
3242
3243
3244/* count object ************************************************************/
3245
3246typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003247 PyObject_HEAD
3248 Py_ssize_t cnt;
3249 PyObject *long_cnt;
3250 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003251} countobject;
3252
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003253/* Counting logic and invariants:
3254
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003255fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003256
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003257 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3258 Advances with: cnt += 1
3259 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003260
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003261slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003262
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003263 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3264 All counting is done with python objects (no overflows or underflows).
3265 Advances with: long_cnt += long_step
3266 Step may be zero -- effectively a slow version of repeat(cnt).
3267 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003268*/
3269
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003270static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003271
3272static PyObject *
3273count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3274{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003275 countobject *lz;
3276 int slow_mode = 0;
3277 Py_ssize_t cnt = 0;
3278 PyObject *long_cnt = NULL;
3279 PyObject *long_step = NULL;
3280 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003281
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003282 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3283 kwlist, &long_cnt, &long_step))
3284 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003285
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003286 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3287 (long_step != NULL && !PyNumber_Check(long_step))) {
3288 PyErr_SetString(PyExc_TypeError, "a number is required");
3289 return NULL;
3290 }
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003291
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003292 if (long_cnt != NULL) {
3293 cnt = PyInt_AsSsize_t(long_cnt);
3294 if ((cnt == -1 && PyErr_Occurred()) || !PyInt_Check(long_cnt)) {
3295 PyErr_Clear();
3296 slow_mode = 1;
3297 }
3298 Py_INCREF(long_cnt);
3299 } else {
3300 cnt = 0;
3301 long_cnt = PyInt_FromLong(0);
3302 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003303
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003304 /* If not specified, step defaults to 1 */
3305 if (long_step == NULL) {
3306 long_step = PyInt_FromLong(1);
3307 if (long_step == NULL) {
3308 Py_DECREF(long_cnt);
3309 return NULL;
3310 }
3311 } else
3312 Py_INCREF(long_step);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003313
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003314 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003315
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003316 /* Fast mode only works when the step is 1 */
3317 if (!PyInt_Check(long_step) ||
3318 PyInt_AS_LONG(long_step) != 1) {
3319 slow_mode = 1;
3320 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003321
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003322 if (slow_mode)
3323 cnt = PY_SSIZE_T_MAX;
3324 else
3325 Py_CLEAR(long_cnt);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003326
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003327 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3328 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3329 assert(slow_mode ||
3330 (PyInt_Check(long_step) && PyInt_AS_LONG(long_step) == 1));
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003331
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003332 /* create countobject structure */
3333 lz = (countobject *)type->tp_alloc(type, 0);
3334 if (lz == NULL) {
3335 Py_XDECREF(long_cnt);
3336 return NULL;
3337 }
3338 lz->cnt = cnt;
3339 lz->long_cnt = long_cnt;
3340 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003341
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003342 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003343}
3344
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003345static void
3346count_dealloc(countobject *lz)
3347{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003348 PyObject_GC_UnTrack(lz);
3349 Py_XDECREF(lz->long_cnt);
3350 Py_XDECREF(lz->long_step);
3351 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003352}
3353
Benjamin Peterson062a7c32009-02-16 21:07:52 +00003354static int
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003355count_traverse(countobject *lz, visitproc visit, void *arg)
3356{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003357 Py_VISIT(lz->long_cnt);
3358 Py_VISIT(lz->long_step);
3359 return 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003360}
3361
3362static PyObject *
3363count_nextlong(countobject *lz)
3364{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003365 PyObject *long_cnt;
3366 PyObject *stepped_up;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003367
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003368 long_cnt = lz->long_cnt;
3369 if (long_cnt == NULL) {
3370 /* Switch to slow_mode */
3371 long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
3372 if (long_cnt == NULL)
3373 return NULL;
3374 }
3375 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003376
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003377 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3378 if (stepped_up == NULL)
3379 return NULL;
3380 lz->long_cnt = stepped_up;
3381 return long_cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003382}
3383
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003384static PyObject *
3385count_next(countobject *lz)
3386{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003387 if (lz->cnt == PY_SSIZE_T_MAX)
3388 return count_nextlong(lz);
3389 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003390}
3391
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003392static PyObject *
3393count_repr(countobject *lz)
3394{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003395 PyObject *cnt_repr, *step_repr = NULL;
3396 PyObject *result = NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003397
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003398 if (lz->cnt != PY_SSIZE_T_MAX)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003399 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003400
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003401 cnt_repr = PyObject_Repr(lz->long_cnt);
3402 if (cnt_repr == NULL)
3403 return NULL;
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003404
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003405 if (PyInt_Check(lz->long_step) && PyInt_AS_LONG(lz->long_step) == 1) {
3406 /* Don't display step when it is an integer equal to 1 */
3407 result = PyString_FromFormat("count(%s)",
3408 PyString_AS_STRING(cnt_repr));
3409 } else {
3410 step_repr = PyObject_Repr(lz->long_step);
3411 if (step_repr != NULL)
3412 result = PyString_FromFormat("count(%s, %s)",
3413 PyString_AS_STRING(cnt_repr),
3414 PyString_AS_STRING(step_repr));
3415 }
3416 Py_DECREF(cnt_repr);
3417 Py_XDECREF(step_repr);
3418 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003419}
3420
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003421static PyObject *
3422count_reduce(countobject *lz)
3423{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003424 if (lz->cnt == PY_SSIZE_T_MAX)
3425 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3426 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003427}
3428
3429PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3430
3431static PyMethodDef count_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003432 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3433 count_reduce_doc},
3434 {NULL, NULL} /* sentinel */
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003435};
3436
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003437PyDoc_STRVAR(count_doc,
Ezio Melottifdc1e0d2010-06-11 02:21:25 +00003438 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003439\n\
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003440Return a count object whose .next() method returns consecutive values.\n\
3441Equivalent to:\n\n\
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003442 def count(firstval=0, step=1):\n\
Eli Benderskye6802db2013-09-03 06:41:58 -07003443 x = firstval\n\
3444 while 1:\n\
3445 yield x\n\
3446 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003447
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003448static PyTypeObject count_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003449 PyVarObject_HEAD_INIT(NULL, 0)
3450 "itertools.count", /* tp_name */
3451 sizeof(countobject), /* tp_basicsize */
3452 0, /* tp_itemsize */
3453 /* methods */
3454 (destructor)count_dealloc, /* tp_dealloc */
3455 0, /* tp_print */
3456 0, /* tp_getattr */
3457 0, /* tp_setattr */
3458 0, /* tp_compare */
3459 (reprfunc)count_repr, /* tp_repr */
3460 0, /* tp_as_number */
3461 0, /* tp_as_sequence */
3462 0, /* tp_as_mapping */
3463 0, /* tp_hash */
3464 0, /* tp_call */
3465 0, /* tp_str */
3466 PyObject_GenericGetAttr, /* tp_getattro */
3467 0, /* tp_setattro */
3468 0, /* tp_as_buffer */
3469 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3470 Py_TPFLAGS_BASETYPE, /* tp_flags */
3471 count_doc, /* tp_doc */
3472 (traverseproc)count_traverse, /* tp_traverse */
3473 0, /* tp_clear */
3474 0, /* tp_richcompare */
3475 0, /* tp_weaklistoffset */
3476 PyObject_SelfIter, /* tp_iter */
3477 (iternextfunc)count_next, /* tp_iternext */
3478 count_methods, /* tp_methods */
3479 0, /* tp_members */
3480 0, /* tp_getset */
3481 0, /* tp_base */
3482 0, /* tp_dict */
3483 0, /* tp_descr_get */
3484 0, /* tp_descr_set */
3485 0, /* tp_dictoffset */
3486 0, /* tp_init */
3487 0, /* tp_alloc */
3488 count_new, /* tp_new */
3489 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003490};
3491
3492
3493/* izip object ************************************************************/
3494
3495#include "Python.h"
3496
3497typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003498 PyObject_HEAD
3499 Py_ssize_t tuplesize;
3500 PyObject *ittuple; /* tuple of iterators */
3501 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003502} izipobject;
3503
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003504static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003505
3506static PyObject *
3507izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3508{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003509 izipobject *lz;
3510 Py_ssize_t i;
3511 PyObject *ittuple; /* tuple of iterators */
3512 PyObject *result;
3513 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003514
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003515 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
3516 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003517
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003518 /* args must be a tuple */
3519 assert(PyTuple_Check(args));
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003520
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003521 /* obtain iterators */
3522 ittuple = PyTuple_New(tuplesize);
3523 if (ittuple == NULL)
3524 return NULL;
3525 for (i=0; i < tuplesize; ++i) {
3526 PyObject *item = PyTuple_GET_ITEM(args, i);
3527 PyObject *it = PyObject_GetIter(item);
3528 if (it == NULL) {
3529 if (PyErr_ExceptionMatches(PyExc_TypeError))
3530 PyErr_Format(PyExc_TypeError,
3531 "izip argument #%zd must support iteration",
3532 i+1);
3533 Py_DECREF(ittuple);
3534 return NULL;
3535 }
3536 PyTuple_SET_ITEM(ittuple, i, it);
3537 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003538
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003539 /* create a result holder */
3540 result = PyTuple_New(tuplesize);
3541 if (result == NULL) {
3542 Py_DECREF(ittuple);
3543 return NULL;
3544 }
3545 for (i=0 ; i < tuplesize ; i++) {
3546 Py_INCREF(Py_None);
3547 PyTuple_SET_ITEM(result, i, Py_None);
3548 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003549
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003550 /* create izipobject structure */
3551 lz = (izipobject *)type->tp_alloc(type, 0);
3552 if (lz == NULL) {
3553 Py_DECREF(ittuple);
3554 Py_DECREF(result);
3555 return NULL;
3556 }
3557 lz->ittuple = ittuple;
3558 lz->tuplesize = tuplesize;
3559 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003560
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003561 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003562}
3563
3564static void
3565izip_dealloc(izipobject *lz)
3566{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003567 PyObject_GC_UnTrack(lz);
3568 Py_XDECREF(lz->ittuple);
3569 Py_XDECREF(lz->result);
3570 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003571}
3572
3573static int
3574izip_traverse(izipobject *lz, visitproc visit, void *arg)
3575{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003576 Py_VISIT(lz->ittuple);
3577 Py_VISIT(lz->result);
3578 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003579}
3580
3581static PyObject *
3582izip_next(izipobject *lz)
3583{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003584 Py_ssize_t i;
3585 Py_ssize_t tuplesize = lz->tuplesize;
3586 PyObject *result = lz->result;
3587 PyObject *it;
3588 PyObject *item;
3589 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003590
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003591 if (tuplesize == 0)
3592 return NULL;
3593 if (Py_REFCNT(result) == 1) {
3594 Py_INCREF(result);
3595 for (i=0 ; i < tuplesize ; i++) {
3596 it = PyTuple_GET_ITEM(lz->ittuple, i);
3597 item = (*Py_TYPE(it)->tp_iternext)(it);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003598 if (item == NULL) {
3599 Py_DECREF(result);
3600 return NULL;
3601 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003602 olditem = PyTuple_GET_ITEM(result, i);
3603 PyTuple_SET_ITEM(result, i, item);
3604 Py_DECREF(olditem);
3605 }
3606 } else {
3607 result = PyTuple_New(tuplesize);
3608 if (result == NULL)
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003609 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003610 for (i=0 ; i < tuplesize ; i++) {
3611 it = PyTuple_GET_ITEM(lz->ittuple, i);
3612 item = (*Py_TYPE(it)->tp_iternext)(it);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003613 if (item == NULL) {
3614 Py_DECREF(result);
3615 return NULL;
3616 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003617 PyTuple_SET_ITEM(result, i, item);
3618 }
3619 }
3620 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003621}
3622
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003623PyDoc_STRVAR(izip_doc,
3624"izip(iter1 [,iter2 [...]]) --> izip object\n\
3625\n\
3626Return a izip object whose .next() method returns a tuple where\n\
3627the i-th element comes from the i-th iterable argument. The .next()\n\
3628method continues until the shortest iterable in the argument sequence\n\
3629is exhausted and then it raises StopIteration. Works like the zip()\n\
3630function but consumes less memory by returning an iterator instead of\n\
3631a list.");
3632
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003633static PyTypeObject izip_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003634 PyVarObject_HEAD_INIT(NULL, 0)
3635 "itertools.izip", /* tp_name */
3636 sizeof(izipobject), /* tp_basicsize */
3637 0, /* tp_itemsize */
3638 /* methods */
3639 (destructor)izip_dealloc, /* tp_dealloc */
3640 0, /* tp_print */
3641 0, /* tp_getattr */
3642 0, /* tp_setattr */
3643 0, /* tp_compare */
3644 0, /* tp_repr */
3645 0, /* tp_as_number */
3646 0, /* tp_as_sequence */
3647 0, /* tp_as_mapping */
3648 0, /* tp_hash */
3649 0, /* tp_call */
3650 0, /* tp_str */
3651 PyObject_GenericGetAttr, /* tp_getattro */
3652 0, /* tp_setattro */
3653 0, /* tp_as_buffer */
3654 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3655 Py_TPFLAGS_BASETYPE, /* tp_flags */
3656 izip_doc, /* tp_doc */
3657 (traverseproc)izip_traverse, /* tp_traverse */
3658 0, /* tp_clear */
3659 0, /* tp_richcompare */
3660 0, /* tp_weaklistoffset */
3661 PyObject_SelfIter, /* tp_iter */
3662 (iternextfunc)izip_next, /* tp_iternext */
3663 0, /* tp_methods */
3664 0, /* tp_members */
3665 0, /* tp_getset */
3666 0, /* tp_base */
3667 0, /* tp_dict */
3668 0, /* tp_descr_get */
3669 0, /* tp_descr_set */
3670 0, /* tp_dictoffset */
3671 0, /* tp_init */
3672 0, /* tp_alloc */
3673 izip_new, /* tp_new */
3674 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003675};
3676
3677
3678/* repeat object ************************************************************/
3679
3680typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003681 PyObject_HEAD
3682 PyObject *element;
3683 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003684} repeatobject;
3685
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003686static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003687
3688static PyObject *
3689repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3690{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003691 repeatobject *ro;
3692 PyObject *element;
Raymond Hettinger58ad2452014-06-24 21:53:45 -07003693 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003694 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003695
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003696 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3697 &element, &cnt))
3698 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003699
Raymond Hettinger58ad2452014-06-24 21:53:45 -07003700 if (kwds != NULL)
3701 n_kwds = PyDict_Size(kwds);
3702 /* Does user supply times argument? */
3703 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003704 cnt = 0;
3705
3706 ro = (repeatobject *)type->tp_alloc(type, 0);
3707 if (ro == NULL)
3708 return NULL;
3709 Py_INCREF(element);
3710 ro->element = element;
3711 ro->cnt = cnt;
3712 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003713}
3714
3715static void
3716repeat_dealloc(repeatobject *ro)
3717{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003718 PyObject_GC_UnTrack(ro);
3719 Py_XDECREF(ro->element);
3720 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003721}
3722
3723static int
3724repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3725{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003726 Py_VISIT(ro->element);
3727 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003728}
3729
3730static PyObject *
3731repeat_next(repeatobject *ro)
3732{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003733 if (ro->cnt == 0)
3734 return NULL;
3735 if (ro->cnt > 0)
3736 ro->cnt--;
3737 Py_INCREF(ro->element);
3738 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003739}
3740
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003741static PyObject *
3742repeat_repr(repeatobject *ro)
3743{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003744 PyObject *result, *objrepr;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003745
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003746 objrepr = PyObject_Repr(ro->element);
3747 if (objrepr == NULL)
3748 return NULL;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003749
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003750 if (ro->cnt == -1)
3751 result = PyString_FromFormat("repeat(%s)",
3752 PyString_AS_STRING(objrepr));
3753 else
3754 result = PyString_FromFormat("repeat(%s, %zd)",
3755 PyString_AS_STRING(objrepr), ro->cnt);
3756 Py_DECREF(objrepr);
3757 return result;
3758}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003759
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003760static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003761repeat_len(repeatobject *ro)
3762{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003763 if (ro->cnt == -1) {
3764 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3765 return NULL;
3766 }
3767 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003768}
3769
Armin Rigof5b3e362006-02-11 21:32:43 +00003770PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003771
3772static PyMethodDef repeat_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003773 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3774 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003775};
3776
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003777PyDoc_STRVAR(repeat_doc,
Raymond Hettinger182edae2009-02-19 02:38:25 +00003778"repeat(object [,times]) -> create an iterator which returns the object\n\
3779for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003780endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003781
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003782static PyTypeObject repeat_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003783 PyVarObject_HEAD_INIT(NULL, 0)
3784 "itertools.repeat", /* tp_name */
3785 sizeof(repeatobject), /* tp_basicsize */
3786 0, /* tp_itemsize */
3787 /* methods */
3788 (destructor)repeat_dealloc, /* tp_dealloc */
3789 0, /* tp_print */
3790 0, /* tp_getattr */
3791 0, /* tp_setattr */
3792 0, /* tp_compare */
3793 (reprfunc)repeat_repr, /* tp_repr */
3794 0, /* tp_as_number */
3795 0, /* tp_as_sequence */
3796 0, /* tp_as_mapping */
3797 0, /* tp_hash */
3798 0, /* tp_call */
3799 0, /* tp_str */
3800 PyObject_GenericGetAttr, /* tp_getattro */
3801 0, /* tp_setattro */
3802 0, /* tp_as_buffer */
3803 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3804 Py_TPFLAGS_BASETYPE, /* tp_flags */
3805 repeat_doc, /* tp_doc */
3806 (traverseproc)repeat_traverse, /* tp_traverse */
3807 0, /* tp_clear */
3808 0, /* tp_richcompare */
3809 0, /* tp_weaklistoffset */
3810 PyObject_SelfIter, /* tp_iter */
3811 (iternextfunc)repeat_next, /* tp_iternext */
3812 repeat_methods, /* tp_methods */
3813 0, /* tp_members */
3814 0, /* tp_getset */
3815 0, /* tp_base */
3816 0, /* tp_dict */
3817 0, /* tp_descr_get */
3818 0, /* tp_descr_set */
3819 0, /* tp_dictoffset */
3820 0, /* tp_init */
3821 0, /* tp_alloc */
3822 repeat_new, /* tp_new */
3823 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003824};
3825
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003826/* iziplongest object ************************************************************/
3827
3828#include "Python.h"
3829
3830typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003831 PyObject_HEAD
3832 Py_ssize_t tuplesize;
3833 Py_ssize_t numactive;
3834 PyObject *ittuple; /* tuple of iterators */
3835 PyObject *result;
3836 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003837} iziplongestobject;
3838
3839static PyTypeObject iziplongest_type;
3840
3841static PyObject *
3842izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3843{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003844 iziplongestobject *lz;
3845 Py_ssize_t i;
3846 PyObject *ittuple; /* tuple of iterators */
3847 PyObject *result;
3848 PyObject *fillvalue = Py_None;
3849 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003850
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003851 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3852 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3853 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3854 PyErr_SetString(PyExc_TypeError,
3855 "izip_longest() got an unexpected keyword argument");
3856 return NULL;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003857 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003858 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003859
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003860 /* args must be a tuple */
3861 assert(PyTuple_Check(args));
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003862
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003863 /* obtain iterators */
3864 ittuple = PyTuple_New(tuplesize);
3865 if (ittuple == NULL)
3866 return NULL;
3867 for (i=0; i < tuplesize; ++i) {
3868 PyObject *item = PyTuple_GET_ITEM(args, i);
3869 PyObject *it = PyObject_GetIter(item);
3870 if (it == NULL) {
3871 if (PyErr_ExceptionMatches(PyExc_TypeError))
3872 PyErr_Format(PyExc_TypeError,
3873 "izip_longest argument #%zd must support iteration",
3874 i+1);
3875 Py_DECREF(ittuple);
3876 return NULL;
3877 }
3878 PyTuple_SET_ITEM(ittuple, i, it);
3879 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003880
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003881 /* create a result holder */
3882 result = PyTuple_New(tuplesize);
3883 if (result == NULL) {
3884 Py_DECREF(ittuple);
3885 return NULL;
3886 }
3887 for (i=0 ; i < tuplesize ; i++) {
3888 Py_INCREF(Py_None);
3889 PyTuple_SET_ITEM(result, i, Py_None);
3890 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003891
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003892 /* create iziplongestobject structure */
3893 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3894 if (lz == NULL) {
3895 Py_DECREF(ittuple);
3896 Py_DECREF(result);
3897 return NULL;
3898 }
3899 lz->ittuple = ittuple;
3900 lz->tuplesize = tuplesize;
3901 lz->numactive = tuplesize;
3902 lz->result = result;
3903 Py_INCREF(fillvalue);
3904 lz->fillvalue = fillvalue;
3905 return (PyObject *)lz;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003906}
3907
3908static void
3909izip_longest_dealloc(iziplongestobject *lz)
3910{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003911 PyObject_GC_UnTrack(lz);
3912 Py_XDECREF(lz->ittuple);
3913 Py_XDECREF(lz->result);
3914 Py_XDECREF(lz->fillvalue);
3915 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003916}
3917
3918static int
3919izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3920{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003921 Py_VISIT(lz->ittuple);
3922 Py_VISIT(lz->result);
3923 Py_VISIT(lz->fillvalue);
3924 return 0;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003925}
3926
3927static PyObject *
3928izip_longest_next(iziplongestobject *lz)
3929{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003930 Py_ssize_t i;
3931 Py_ssize_t tuplesize = lz->tuplesize;
3932 PyObject *result = lz->result;
3933 PyObject *it;
3934 PyObject *item;
3935 PyObject *olditem;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003936
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003937 if (tuplesize == 0)
3938 return NULL;
3939 if (lz->numactive == 0)
3940 return NULL;
3941 if (Py_REFCNT(result) == 1) {
3942 Py_INCREF(result);
3943 for (i=0 ; i < tuplesize ; i++) {
3944 it = PyTuple_GET_ITEM(lz->ittuple, i);
3945 if (it == NULL) {
3946 Py_INCREF(lz->fillvalue);
3947 item = lz->fillvalue;
3948 } else {
3949 item = PyIter_Next(it);
3950 if (item == NULL) {
3951 lz->numactive -= 1;
3952 if (lz->numactive == 0 || PyErr_Occurred()) {
3953 lz->numactive = 0;
3954 Py_DECREF(result);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003955 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003956 } else {
3957 Py_INCREF(lz->fillvalue);
3958 item = lz->fillvalue;
3959 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3960 Py_DECREF(it);
3961 }
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003962 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003963 }
3964 olditem = PyTuple_GET_ITEM(result, i);
3965 PyTuple_SET_ITEM(result, i, item);
3966 Py_DECREF(olditem);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003967 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003968 } else {
3969 result = PyTuple_New(tuplesize);
3970 if (result == NULL)
3971 return NULL;
3972 for (i=0 ; i < tuplesize ; i++) {
3973 it = PyTuple_GET_ITEM(lz->ittuple, i);
3974 if (it == NULL) {
3975 Py_INCREF(lz->fillvalue);
3976 item = lz->fillvalue;
3977 } else {
3978 item = PyIter_Next(it);
3979 if (item == NULL) {
3980 lz->numactive -= 1;
3981 if (lz->numactive == 0 || PyErr_Occurred()) {
3982 lz->numactive = 0;
3983 Py_DECREF(result);
3984 return NULL;
3985 } else {
3986 Py_INCREF(lz->fillvalue);
3987 item = lz->fillvalue;
3988 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3989 Py_DECREF(it);
3990 }
3991 }
3992 }
3993 PyTuple_SET_ITEM(result, i, item);
3994 }
3995 }
3996 return result;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003997}
3998
3999PyDoc_STRVAR(izip_longest_doc,
4000"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
4001\n\
4002Return an izip_longest object whose .next() method returns a tuple where\n\
4003the i-th element comes from the i-th iterable argument. The .next()\n\
4004method continues until the longest iterable in the argument sequence\n\
4005is exhausted and then it raises StopIteration. When the shorter iterables\n\
4006are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4007defaults to None or can be specified by a keyword argument.\n\
4008");
4009
4010static PyTypeObject iziplongest_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004011 PyVarObject_HEAD_INIT(NULL, 0)
4012 "itertools.izip_longest", /* tp_name */
4013 sizeof(iziplongestobject), /* tp_basicsize */
4014 0, /* tp_itemsize */
4015 /* methods */
4016 (destructor)izip_longest_dealloc, /* tp_dealloc */
4017 0, /* tp_print */
4018 0, /* tp_getattr */
4019 0, /* tp_setattr */
4020 0, /* tp_compare */
4021 0, /* tp_repr */
4022 0, /* tp_as_number */
4023 0, /* tp_as_sequence */
4024 0, /* tp_as_mapping */
4025 0, /* tp_hash */
4026 0, /* tp_call */
4027 0, /* tp_str */
4028 PyObject_GenericGetAttr, /* tp_getattro */
4029 0, /* tp_setattro */
4030 0, /* tp_as_buffer */
4031 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4032 Py_TPFLAGS_BASETYPE, /* tp_flags */
4033 izip_longest_doc, /* tp_doc */
4034 (traverseproc)izip_longest_traverse, /* tp_traverse */
4035 0, /* tp_clear */
4036 0, /* tp_richcompare */
4037 0, /* tp_weaklistoffset */
4038 PyObject_SelfIter, /* tp_iter */
4039 (iternextfunc)izip_longest_next, /* tp_iternext */
4040 0, /* tp_methods */
4041 0, /* tp_members */
4042 0, /* tp_getset */
4043 0, /* tp_base */
4044 0, /* tp_dict */
4045 0, /* tp_descr_get */
4046 0, /* tp_descr_set */
4047 0, /* tp_dictoffset */
4048 0, /* tp_init */
4049 0, /* tp_alloc */
4050 izip_longest_new, /* tp_new */
4051 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd36862c2007-02-21 05:20:38 +00004052};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004053
4054/* module level code ********************************************************/
4055
4056PyDoc_STRVAR(module_doc,
4057"Functional tools for creating and using iterators.\n\
4058\n\
4059Infinite iterators:\n\
4060count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004061cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004062repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004063\n\
4064Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004065chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettingere76816b2009-01-29 03:43:44 +00004066compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004067dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4068groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00004069ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
4070ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004071islice(seq, [start,] stop [, step]) --> elements from\n\
4072 seq[start:stop:step]\n\
4073imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
4074starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004075tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004076takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004077izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4078izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4079\n\
4080Combinatoric generators:\n\
4081product(p, q, ... [repeat=1]) --> cartesian product\n\
4082permutations(p[, r])\n\
Raymond Hettinger9eac1192009-11-19 01:22:04 +00004083combinations(p, r)\n\
4084combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004085");
4086
4087
Raymond Hettingerad983e72003-11-12 14:32:26 +00004088static PyMethodDef module_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004089 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4090 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004091};
4092
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004093PyMODINIT_FUNC
4094inititertools(void)
4095{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004096 int i;
4097 PyObject *m;
4098 char *name;
4099 PyTypeObject *typelist[] = {
4100 &combinations_type,
4101 &cwr_type,
4102 &cycle_type,
4103 &dropwhile_type,
4104 &takewhile_type,
4105 &islice_type,
4106 &starmap_type,
4107 &imap_type,
4108 &chain_type,
4109 &compress_type,
4110 &ifilter_type,
4111 &ifilterfalse_type,
4112 &count_type,
4113 &izip_type,
4114 &iziplongest_type,
4115 &permutations_type,
4116 &product_type,
4117 &repeat_type,
4118 &groupby_type,
4119 NULL
4120 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004121
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004122 Py_TYPE(&teedataobject_type) = &PyType_Type;
4123 m = Py_InitModule3("itertools", module_methods, module_doc);
4124 if (m == NULL)
4125 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004126
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004127 for (i=0 ; typelist[i] != NULL ; i++) {
4128 if (PyType_Ready(typelist[i]) < 0)
4129 return;
4130 name = strchr(typelist[i]->tp_name, '.');
4131 assert (name != NULL);
4132 Py_INCREF(typelist[i]);
4133 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4134 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004135
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004136 if (PyType_Ready(&teedataobject_type) < 0)
4137 return;
4138 if (PyType_Ready(&tee_type) < 0)
4139 return;
4140 if (PyType_Ready(&_grouper_type) < 0)
4141 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004142}