blob: 68285b82f6c98c4184912f68159dee7232a4281a [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>
Raymond Hettinger96ef8112003-02-01 00:10:11 +00007*/
8
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00009
10/* groupby object ***********************************************************/
11
12typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +000013 PyObject_HEAD
14 PyObject *it;
15 PyObject *keyfunc;
16 PyObject *tgtkey;
17 PyObject *currkey;
18 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000019} groupbyobject;
20
21static PyTypeObject groupby_type;
22static PyObject *_grouper_create(groupbyobject *, PyObject *);
23
24static PyObject *
25groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
26{
Antoine Pitrouc83ea132010-05-09 14:46:46 +000027 static char *kwargs[] = {"iterable", "key", NULL};
28 groupbyobject *gbo;
29 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000030
Antoine Pitrouc83ea132010-05-09 14:46:46 +000031 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
32 &it, &keyfunc))
33 return NULL;
34
35 gbo = (groupbyobject *)type->tp_alloc(type, 0);
36 if (gbo == NULL)
37 return NULL;
38 gbo->tgtkey = NULL;
39 gbo->currkey = NULL;
40 gbo->currvalue = NULL;
41 gbo->keyfunc = keyfunc;
42 Py_INCREF(keyfunc);
43 gbo->it = PyObject_GetIter(it);
44 if (gbo->it == NULL) {
45 Py_DECREF(gbo);
46 return NULL;
47 }
48 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000049}
50
51static void
52groupby_dealloc(groupbyobject *gbo)
53{
Antoine Pitrouc83ea132010-05-09 14:46:46 +000054 PyObject_GC_UnTrack(gbo);
55 Py_XDECREF(gbo->it);
56 Py_XDECREF(gbo->keyfunc);
57 Py_XDECREF(gbo->tgtkey);
58 Py_XDECREF(gbo->currkey);
59 Py_XDECREF(gbo->currvalue);
60 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000061}
62
63static int
64groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
65{
Antoine Pitrouc83ea132010-05-09 14:46:46 +000066 Py_VISIT(gbo->it);
67 Py_VISIT(gbo->keyfunc);
68 Py_VISIT(gbo->tgtkey);
69 Py_VISIT(gbo->currkey);
70 Py_VISIT(gbo->currvalue);
71 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000072}
73
74static PyObject *
75groupby_next(groupbyobject *gbo)
76{
Antoine Pitrouc83ea132010-05-09 14:46:46 +000077 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000078
Antoine Pitrouc83ea132010-05-09 14:46:46 +000079 /* skip to next iteration group */
80 for (;;) {
81 if (gbo->currkey == NULL)
82 /* pass */;
83 else if (gbo->tgtkey == NULL)
84 break;
85 else {
86 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000087
Antoine Pitrouc83ea132010-05-09 14:46:46 +000088 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
89 gbo->currkey, Py_EQ);
90 if (rcmp == -1)
91 return NULL;
92 else if (rcmp == 0)
93 break;
94 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000095
Antoine Pitrouc83ea132010-05-09 14:46:46 +000096 newvalue = PyIter_Next(gbo->it);
97 if (newvalue == NULL)
98 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000099
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000100 if (gbo->keyfunc == Py_None) {
101 newkey = newvalue;
102 Py_INCREF(newvalue);
103 } else {
104 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
105 newvalue, NULL);
106 if (newkey == NULL) {
107 Py_DECREF(newvalue);
108 return NULL;
109 }
110 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000111
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000112 tmp = gbo->currkey;
113 gbo->currkey = newkey;
114 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000115
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000116 tmp = gbo->currvalue;
117 gbo->currvalue = newvalue;
118 Py_XDECREF(tmp);
119 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000120
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000121 Py_INCREF(gbo->currkey);
122 tmp = gbo->tgtkey;
123 gbo->tgtkey = gbo->currkey;
124 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000125
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000126 grouper = _grouper_create(gbo, gbo->tgtkey);
127 if (grouper == NULL)
128 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000129
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000130 r = PyTuple_Pack(2, gbo->currkey, grouper);
131 Py_DECREF(grouper);
132 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000133}
134
135PyDoc_STRVAR(groupby_doc,
136"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
137(key, sub-iterator) grouped by each value of key(value).\n");
138
139static PyTypeObject groupby_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000140 PyVarObject_HEAD_INIT(NULL, 0)
141 "itertools.groupby", /* tp_name */
142 sizeof(groupbyobject), /* tp_basicsize */
143 0, /* tp_itemsize */
144 /* methods */
145 (destructor)groupby_dealloc, /* tp_dealloc */
146 0, /* tp_print */
147 0, /* tp_getattr */
148 0, /* tp_setattr */
149 0, /* tp_compare */
150 0, /* tp_repr */
151 0, /* tp_as_number */
152 0, /* tp_as_sequence */
153 0, /* tp_as_mapping */
154 0, /* tp_hash */
155 0, /* tp_call */
156 0, /* tp_str */
157 PyObject_GenericGetAttr, /* tp_getattro */
158 0, /* tp_setattro */
159 0, /* tp_as_buffer */
160 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
161 Py_TPFLAGS_BASETYPE, /* tp_flags */
162 groupby_doc, /* tp_doc */
163 (traverseproc)groupby_traverse, /* tp_traverse */
164 0, /* tp_clear */
165 0, /* tp_richcompare */
166 0, /* tp_weaklistoffset */
167 PyObject_SelfIter, /* tp_iter */
168 (iternextfunc)groupby_next, /* tp_iternext */
169 0, /* tp_methods */
170 0, /* tp_members */
171 0, /* tp_getset */
172 0, /* tp_base */
173 0, /* tp_dict */
174 0, /* tp_descr_get */
175 0, /* tp_descr_set */
176 0, /* tp_dictoffset */
177 0, /* tp_init */
178 0, /* tp_alloc */
179 groupby_new, /* tp_new */
180 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000181};
182
183
184/* _grouper object (internal) ************************************************/
185
186typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000187 PyObject_HEAD
188 PyObject *parent;
189 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000190} _grouperobject;
191
192static PyTypeObject _grouper_type;
193
194static PyObject *
195_grouper_create(groupbyobject *parent, PyObject *tgtkey)
196{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000197 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000198
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000199 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
200 if (igo == NULL)
201 return NULL;
202 igo->parent = (PyObject *)parent;
203 Py_INCREF(parent);
204 igo->tgtkey = tgtkey;
205 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000206
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000207 PyObject_GC_Track(igo);
208 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000209}
210
211static void
212_grouper_dealloc(_grouperobject *igo)
213{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000214 PyObject_GC_UnTrack(igo);
215 Py_DECREF(igo->parent);
216 Py_DECREF(igo->tgtkey);
217 PyObject_GC_Del(igo);
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000218}
219
220static int
221_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
222{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000223 Py_VISIT(igo->parent);
224 Py_VISIT(igo->tgtkey);
225 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000226}
227
228static PyObject *
229_grouper_next(_grouperobject *igo)
230{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000231 groupbyobject *gbo = (groupbyobject *)igo->parent;
232 PyObject *newvalue, *newkey, *r;
233 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000234
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000235 if (gbo->currvalue == NULL) {
236 newvalue = PyIter_Next(gbo->it);
237 if (newvalue == NULL)
238 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000239
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000240 if (gbo->keyfunc == Py_None) {
241 newkey = newvalue;
242 Py_INCREF(newvalue);
243 } else {
244 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
245 newvalue, NULL);
246 if (newkey == NULL) {
247 Py_DECREF(newvalue);
248 return NULL;
249 }
250 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000251
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000252 assert(gbo->currkey == NULL);
253 gbo->currkey = newkey;
254 gbo->currvalue = newvalue;
255 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000256
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000257 assert(gbo->currkey != NULL);
258 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
259 if (rcmp <= 0)
260 /* got any error or current group is end */
261 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000262
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000263 r = gbo->currvalue;
264 gbo->currvalue = NULL;
265 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000266
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000267 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000268}
269
270static PyTypeObject _grouper_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000271 PyVarObject_HEAD_INIT(NULL, 0)
272 "itertools._grouper", /* tp_name */
273 sizeof(_grouperobject), /* tp_basicsize */
274 0, /* tp_itemsize */
275 /* methods */
276 (destructor)_grouper_dealloc, /* tp_dealloc */
277 0, /* tp_print */
278 0, /* tp_getattr */
279 0, /* tp_setattr */
280 0, /* tp_compare */
281 0, /* tp_repr */
282 0, /* tp_as_number */
283 0, /* tp_as_sequence */
284 0, /* tp_as_mapping */
285 0, /* tp_hash */
286 0, /* tp_call */
287 0, /* tp_str */
288 PyObject_GenericGetAttr, /* tp_getattro */
289 0, /* tp_setattro */
290 0, /* tp_as_buffer */
291 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
292 0, /* tp_doc */
293 (traverseproc)_grouper_traverse,/* tp_traverse */
294 0, /* tp_clear */
295 0, /* tp_richcompare */
296 0, /* tp_weaklistoffset */
297 PyObject_SelfIter, /* tp_iter */
298 (iternextfunc)_grouper_next, /* tp_iternext */
299 0, /* tp_methods */
300 0, /* tp_members */
301 0, /* tp_getset */
302 0, /* tp_base */
303 0, /* tp_dict */
304 0, /* tp_descr_get */
305 0, /* tp_descr_set */
306 0, /* tp_dictoffset */
307 0, /* tp_init */
308 0, /* tp_alloc */
309 0, /* tp_new */
310 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000311};
312
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000313
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000314
Raymond Hettingerad983e72003-11-12 14:32:26 +0000315/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000316
Raymond Hettingerad983e72003-11-12 14:32:26 +0000317/* The teedataobject pre-allocates space for LINKCELLS number of objects.
318 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000319 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000320 the other structure members including PyHEAD overhead. The larger the
321 value, the less memory overhead per object and the less time spent
322 allocating/deallocating new links. The smaller the number, the less
323 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000324*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000325#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000326
327typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000328 PyObject_HEAD
329 PyObject *it;
330 int numread;
331 PyObject *nextlink;
332 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000333} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000334
335typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000336 PyObject_HEAD
337 teedataobject *dataobj;
338 int index;
339 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000340} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000341
Raymond Hettingerad983e72003-11-12 14:32:26 +0000342static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000343
344static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000345teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000346{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000347 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000348
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000349 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
350 if (tdo == NULL)
351 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000352
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000353 tdo->numread = 0;
354 tdo->nextlink = NULL;
355 Py_INCREF(it);
356 tdo->it = it;
357 PyObject_GC_Track(tdo);
358 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000359}
360
361static PyObject *
362teedataobject_jumplink(teedataobject *tdo)
363{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000364 if (tdo->nextlink == NULL)
365 tdo->nextlink = teedataobject_new(tdo->it);
366 Py_XINCREF(tdo->nextlink);
367 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000368}
369
370static PyObject *
371teedataobject_getitem(teedataobject *tdo, int i)
372{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000373 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000374
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000375 assert(i < LINKCELLS);
376 if (i < tdo->numread)
377 value = tdo->values[i];
378 else {
379 /* this is the lead iterator, so fetch more data */
380 assert(i == tdo->numread);
381 value = PyIter_Next(tdo->it);
382 if (value == NULL)
383 return NULL;
384 tdo->numread++;
385 tdo->values[i] = value;
386 }
387 Py_INCREF(value);
388 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000389}
390
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000391static int
392teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
393{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000394 int i;
395 Py_VISIT(tdo->it);
396 for (i = 0; i < tdo->numread; i++)
397 Py_VISIT(tdo->values[i]);
398 Py_VISIT(tdo->nextlink);
399 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000400}
401
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200402static void
403teedataobject_safe_decref(PyObject *obj)
404{
405 while (obj && Py_TYPE(obj) == &teedataobject_type &&
406 Py_REFCNT(obj) == 1) {
407 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
408 ((teedataobject *)obj)->nextlink = NULL;
409 Py_DECREF(obj);
410 obj = nextlink;
411 }
412 Py_XDECREF(obj);
413}
414
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000415static int
416teedataobject_clear(teedataobject *tdo)
417{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000418 int i;
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200419 PyObject *tmp;
420
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000421 Py_CLEAR(tdo->it);
422 for (i=0 ; i<tdo->numread ; i++)
423 Py_CLEAR(tdo->values[i]);
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200424 tmp = tdo->nextlink;
425 tdo->nextlink = NULL;
426 teedataobject_safe_decref(tmp);
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000427 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000428}
429
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000430static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000431teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000432{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000433 PyObject_GC_UnTrack(tdo);
434 teedataobject_clear(tdo);
435 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000436}
437
Raymond Hettingerad983e72003-11-12 14:32:26 +0000438PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000439
Raymond Hettingerad983e72003-11-12 14:32:26 +0000440static PyTypeObject teedataobject_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000441 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
442 "itertools.tee_dataobject", /* tp_name */
443 sizeof(teedataobject), /* tp_basicsize */
444 0, /* tp_itemsize */
445 /* methods */
446 (destructor)teedataobject_dealloc, /* tp_dealloc */
447 0, /* tp_print */
448 0, /* tp_getattr */
449 0, /* tp_setattr */
450 0, /* tp_compare */
451 0, /* tp_repr */
452 0, /* tp_as_number */
453 0, /* tp_as_sequence */
454 0, /* tp_as_mapping */
455 0, /* tp_hash */
456 0, /* tp_call */
457 0, /* tp_str */
458 PyObject_GenericGetAttr, /* tp_getattro */
459 0, /* tp_setattro */
460 0, /* tp_as_buffer */
461 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
462 teedataobject_doc, /* tp_doc */
463 (traverseproc)teedataobject_traverse, /* tp_traverse */
464 (inquiry)teedataobject_clear, /* tp_clear */
465 0, /* tp_richcompare */
466 0, /* tp_weaklistoffset */
467 0, /* tp_iter */
468 0, /* tp_iternext */
469 0, /* tp_methods */
470 0, /* tp_members */
471 0, /* tp_getset */
472 0, /* tp_base */
473 0, /* tp_dict */
474 0, /* tp_descr_get */
475 0, /* tp_descr_set */
476 0, /* tp_dictoffset */
477 0, /* tp_init */
478 0, /* tp_alloc */
479 0, /* tp_new */
480 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000481};
482
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000483
484static PyTypeObject tee_type;
485
486static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000487tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000488{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000489 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000490
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000491 if (to->index >= LINKCELLS) {
492 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200493 if (link == NULL)
494 return NULL;
Serhiy Storchaka5951f232015-12-24 10:35:35 +0200495 Py_SETREF(to->dataobj, (teedataobject *)link);
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000496 to->index = 0;
497 }
498 value = teedataobject_getitem(to->dataobj, to->index);
499 if (value == NULL)
500 return NULL;
501 to->index++;
502 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000503}
504
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000505static int
506tee_traverse(teeobject *to, visitproc visit, void *arg)
507{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000508 Py_VISIT((PyObject *)to->dataobj);
509 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000510}
511
Raymond Hettingerad983e72003-11-12 14:32:26 +0000512static PyObject *
513tee_copy(teeobject *to)
514{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000515 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000516
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000517 newto = PyObject_GC_New(teeobject, &tee_type);
518 if (newto == NULL)
519 return NULL;
520 Py_INCREF(to->dataobj);
521 newto->dataobj = to->dataobj;
522 newto->index = to->index;
523 newto->weakreflist = NULL;
524 PyObject_GC_Track(newto);
525 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000526}
527
528PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
529
530static PyObject *
531tee_fromiterable(PyObject *iterable)
532{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000533 teeobject *to;
534 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000535
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000536 it = PyObject_GetIter(iterable);
537 if (it == NULL)
538 return NULL;
539 if (PyObject_TypeCheck(it, &tee_type)) {
540 to = (teeobject *)tee_copy((teeobject *)it);
541 goto done;
542 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000543
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000544 to = PyObject_GC_New(teeobject, &tee_type);
545 if (to == NULL)
546 goto done;
547 to->dataobj = (teedataobject *)teedataobject_new(it);
548 if (!to->dataobj) {
549 PyObject_GC_Del(to);
550 to = NULL;
551 goto done;
552 }
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000553
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000554 to->index = 0;
555 to->weakreflist = NULL;
556 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000557done:
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000558 Py_XDECREF(it);
559 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000560}
561
562static PyObject *
563tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
564{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000565 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000566
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000567 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
568 return NULL;
569 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000570}
571
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000572static int
573tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000574{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000575 if (to->weakreflist != NULL)
576 PyObject_ClearWeakRefs((PyObject *) to);
577 Py_CLEAR(to->dataobj);
578 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000579}
580
581static void
582tee_dealloc(teeobject *to)
583{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000584 PyObject_GC_UnTrack(to);
585 tee_clear(to);
586 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000587}
588
Raymond Hettingerad983e72003-11-12 14:32:26 +0000589PyDoc_STRVAR(teeobject_doc,
590"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000591
Raymond Hettingerad983e72003-11-12 14:32:26 +0000592static PyMethodDef tee_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000593 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
594 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000595};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000596
597static PyTypeObject tee_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000598 PyVarObject_HEAD_INIT(NULL, 0)
599 "itertools.tee", /* tp_name */
600 sizeof(teeobject), /* tp_basicsize */
601 0, /* tp_itemsize */
602 /* methods */
603 (destructor)tee_dealloc, /* tp_dealloc */
604 0, /* tp_print */
605 0, /* tp_getattr */
606 0, /* tp_setattr */
607 0, /* tp_compare */
608 0, /* tp_repr */
609 0, /* tp_as_number */
610 0, /* tp_as_sequence */
611 0, /* tp_as_mapping */
612 0, /* tp_hash */
613 0, /* tp_call */
614 0, /* tp_str */
615 0, /* tp_getattro */
616 0, /* tp_setattro */
617 0, /* tp_as_buffer */
618 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
619 teeobject_doc, /* tp_doc */
620 (traverseproc)tee_traverse, /* tp_traverse */
621 (inquiry)tee_clear, /* tp_clear */
622 0, /* tp_richcompare */
623 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
624 PyObject_SelfIter, /* tp_iter */
625 (iternextfunc)tee_next, /* tp_iternext */
626 tee_methods, /* tp_methods */
627 0, /* tp_members */
628 0, /* tp_getset */
629 0, /* tp_base */
630 0, /* tp_dict */
631 0, /* tp_descr_get */
632 0, /* tp_descr_set */
633 0, /* tp_dictoffset */
634 0, /* tp_init */
635 0, /* tp_alloc */
636 tee_new, /* tp_new */
637 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000638};
639
Raymond Hettingerad983e72003-11-12 14:32:26 +0000640static PyObject *
641tee(PyObject *self, PyObject *args)
642{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000643 Py_ssize_t i, n=2;
644 PyObject *it, *iterable, *copyable, *result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000645
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000646 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
647 return NULL;
648 if (n < 0) {
649 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
650 return NULL;
651 }
652 result = PyTuple_New(n);
653 if (result == NULL)
654 return NULL;
655 if (n == 0)
656 return result;
657 it = PyObject_GetIter(iterable);
658 if (it == NULL) {
659 Py_DECREF(result);
660 return NULL;
661 }
662 if (!PyObject_HasAttrString(it, "__copy__")) {
663 copyable = tee_fromiterable(it);
664 Py_DECREF(it);
665 if (copyable == NULL) {
666 Py_DECREF(result);
667 return NULL;
668 }
669 } else
670 copyable = it;
671 PyTuple_SET_ITEM(result, 0, copyable);
672 for (i=1 ; i<n ; i++) {
673 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
674 if (copyable == NULL) {
675 Py_DECREF(result);
676 return NULL;
677 }
678 PyTuple_SET_ITEM(result, i, copyable);
679 }
680 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000681}
682
683PyDoc_STRVAR(tee_doc,
684"tee(iterable, n=2) --> tuple of n independent iterators.");
685
686
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000687/* cycle object **********************************************************/
688
689typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000690 PyObject_HEAD
691 PyObject *it;
692 PyObject *saved;
693 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000694} cycleobject;
695
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000696static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000697
698static PyObject *
699cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
700{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000701 PyObject *it;
702 PyObject *iterable;
703 PyObject *saved;
704 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000705
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000706 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
707 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000708
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000709 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
710 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000711
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000712 /* Get iterator. */
713 it = PyObject_GetIter(iterable);
714 if (it == NULL)
715 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000716
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000717 saved = PyList_New(0);
718 if (saved == NULL) {
719 Py_DECREF(it);
720 return NULL;
721 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000722
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000723 /* create cycleobject structure */
724 lz = (cycleobject *)type->tp_alloc(type, 0);
725 if (lz == NULL) {
726 Py_DECREF(it);
727 Py_DECREF(saved);
728 return NULL;
729 }
730 lz->it = it;
731 lz->saved = saved;
732 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000733
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000734 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000735}
736
737static void
738cycle_dealloc(cycleobject *lz)
739{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000740 PyObject_GC_UnTrack(lz);
741 Py_XDECREF(lz->saved);
742 Py_XDECREF(lz->it);
743 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000744}
745
746static int
747cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
748{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000749 Py_VISIT(lz->it);
750 Py_VISIT(lz->saved);
751 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000752}
753
754static PyObject *
755cycle_next(cycleobject *lz)
756{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000757 PyObject *item;
758 PyObject *it;
759 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000760
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000761 while (1) {
762 item = PyIter_Next(lz->it);
763 if (item != NULL) {
764 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
765 Py_DECREF(item);
766 return NULL;
767 }
768 return item;
769 }
770 if (PyErr_Occurred()) {
771 if (PyErr_ExceptionMatches(PyExc_StopIteration))
772 PyErr_Clear();
773 else
774 return NULL;
775 }
776 if (PyList_Size(lz->saved) == 0)
777 return NULL;
778 it = PyObject_GetIter(lz->saved);
779 if (it == NULL)
780 return NULL;
781 tmp = lz->it;
782 lz->it = it;
783 lz->firstpass = 1;
784 Py_DECREF(tmp);
785 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000786}
787
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000788PyDoc_STRVAR(cycle_doc,
789"cycle(iterable) --> cycle object\n\
790\n\
791Return elements from the iterable until it is exhausted.\n\
792Then repeat the sequence indefinitely.");
793
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000794static PyTypeObject cycle_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000795 PyVarObject_HEAD_INIT(NULL, 0)
796 "itertools.cycle", /* tp_name */
797 sizeof(cycleobject), /* tp_basicsize */
798 0, /* tp_itemsize */
799 /* methods */
800 (destructor)cycle_dealloc, /* tp_dealloc */
801 0, /* tp_print */
802 0, /* tp_getattr */
803 0, /* tp_setattr */
804 0, /* tp_compare */
805 0, /* tp_repr */
806 0, /* tp_as_number */
807 0, /* tp_as_sequence */
808 0, /* tp_as_mapping */
809 0, /* tp_hash */
810 0, /* tp_call */
811 0, /* tp_str */
812 PyObject_GenericGetAttr, /* tp_getattro */
813 0, /* tp_setattro */
814 0, /* tp_as_buffer */
815 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
816 Py_TPFLAGS_BASETYPE, /* tp_flags */
817 cycle_doc, /* tp_doc */
818 (traverseproc)cycle_traverse, /* tp_traverse */
819 0, /* tp_clear */
820 0, /* tp_richcompare */
821 0, /* tp_weaklistoffset */
822 PyObject_SelfIter, /* tp_iter */
823 (iternextfunc)cycle_next, /* tp_iternext */
824 0, /* tp_methods */
825 0, /* tp_members */
826 0, /* tp_getset */
827 0, /* tp_base */
828 0, /* tp_dict */
829 0, /* tp_descr_get */
830 0, /* tp_descr_set */
831 0, /* tp_dictoffset */
832 0, /* tp_init */
833 0, /* tp_alloc */
834 cycle_new, /* tp_new */
835 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000836};
837
838
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000839/* dropwhile object **********************************************************/
840
841typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000842 PyObject_HEAD
843 PyObject *func;
844 PyObject *it;
845 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000846} dropwhileobject;
847
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000848static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000849
850static PyObject *
851dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
852{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000853 PyObject *func, *seq;
854 PyObject *it;
855 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000856
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000857 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
858 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000859
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000860 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
861 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000862
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000863 /* Get iterator. */
864 it = PyObject_GetIter(seq);
865 if (it == NULL)
866 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000867
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000868 /* create dropwhileobject structure */
869 lz = (dropwhileobject *)type->tp_alloc(type, 0);
870 if (lz == NULL) {
871 Py_DECREF(it);
872 return NULL;
873 }
874 Py_INCREF(func);
875 lz->func = func;
876 lz->it = it;
877 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000878
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000879 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000880}
881
882static void
883dropwhile_dealloc(dropwhileobject *lz)
884{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000885 PyObject_GC_UnTrack(lz);
886 Py_XDECREF(lz->func);
887 Py_XDECREF(lz->it);
888 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000889}
890
891static int
892dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
893{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000894 Py_VISIT(lz->it);
895 Py_VISIT(lz->func);
896 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000897}
898
899static PyObject *
900dropwhile_next(dropwhileobject *lz)
901{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000902 PyObject *item, *good;
903 PyObject *it = lz->it;
904 long ok;
905 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000906
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000907 iternext = *Py_TYPE(it)->tp_iternext;
908 for (;;) {
909 item = iternext(it);
910 if (item == NULL)
911 return NULL;
912 if (lz->start == 1)
913 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000914
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000915 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
916 if (good == NULL) {
917 Py_DECREF(item);
918 return NULL;
919 }
920 ok = PyObject_IsTrue(good);
921 Py_DECREF(good);
Antoine Pitrouc5bef752012-08-15 23:16:51 +0200922 if (ok == 0) {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000923 lz->start = 1;
924 return item;
925 }
926 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +0200927 if (ok < 0)
928 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000929 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000930}
931
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000932PyDoc_STRVAR(dropwhile_doc,
933"dropwhile(predicate, iterable) --> dropwhile object\n\
934\n\
935Drop items from the iterable while predicate(item) is true.\n\
936Afterwards, return every element until the iterable is exhausted.");
937
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000938static PyTypeObject dropwhile_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000939 PyVarObject_HEAD_INIT(NULL, 0)
940 "itertools.dropwhile", /* tp_name */
941 sizeof(dropwhileobject), /* tp_basicsize */
942 0, /* tp_itemsize */
943 /* methods */
944 (destructor)dropwhile_dealloc, /* tp_dealloc */
945 0, /* tp_print */
946 0, /* tp_getattr */
947 0, /* tp_setattr */
948 0, /* tp_compare */
949 0, /* tp_repr */
950 0, /* tp_as_number */
951 0, /* tp_as_sequence */
952 0, /* tp_as_mapping */
953 0, /* tp_hash */
954 0, /* tp_call */
955 0, /* tp_str */
956 PyObject_GenericGetAttr, /* tp_getattro */
957 0, /* tp_setattro */
958 0, /* tp_as_buffer */
959 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
960 Py_TPFLAGS_BASETYPE, /* tp_flags */
961 dropwhile_doc, /* tp_doc */
962 (traverseproc)dropwhile_traverse, /* tp_traverse */
963 0, /* tp_clear */
964 0, /* tp_richcompare */
965 0, /* tp_weaklistoffset */
966 PyObject_SelfIter, /* tp_iter */
967 (iternextfunc)dropwhile_next, /* tp_iternext */
968 0, /* tp_methods */
969 0, /* tp_members */
970 0, /* tp_getset */
971 0, /* tp_base */
972 0, /* tp_dict */
973 0, /* tp_descr_get */
974 0, /* tp_descr_set */
975 0, /* tp_dictoffset */
976 0, /* tp_init */
977 0, /* tp_alloc */
978 dropwhile_new, /* tp_new */
979 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000980};
981
982
983/* takewhile object **********************************************************/
984
985typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000986 PyObject_HEAD
987 PyObject *func;
988 PyObject *it;
989 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000990} takewhileobject;
991
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000992static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000993
994static PyObject *
995takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
996{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000997 PyObject *func, *seq;
998 PyObject *it;
999 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001000
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001001 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1002 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001003
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001004 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1005 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001006
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001007 /* Get iterator. */
1008 it = PyObject_GetIter(seq);
1009 if (it == NULL)
1010 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001011
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001012 /* create takewhileobject structure */
1013 lz = (takewhileobject *)type->tp_alloc(type, 0);
1014 if (lz == NULL) {
1015 Py_DECREF(it);
1016 return NULL;
1017 }
1018 Py_INCREF(func);
1019 lz->func = func;
1020 lz->it = it;
1021 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001022
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001023 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001024}
1025
1026static void
1027takewhile_dealloc(takewhileobject *lz)
1028{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001029 PyObject_GC_UnTrack(lz);
1030 Py_XDECREF(lz->func);
1031 Py_XDECREF(lz->it);
1032 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001033}
1034
1035static int
1036takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1037{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001038 Py_VISIT(lz->it);
1039 Py_VISIT(lz->func);
1040 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001041}
1042
1043static PyObject *
1044takewhile_next(takewhileobject *lz)
1045{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001046 PyObject *item, *good;
1047 PyObject *it = lz->it;
1048 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001049
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001050 if (lz->stop == 1)
1051 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001052
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001053 item = (*Py_TYPE(it)->tp_iternext)(it);
1054 if (item == NULL)
1055 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001056
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001057 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1058 if (good == NULL) {
1059 Py_DECREF(item);
1060 return NULL;
1061 }
1062 ok = PyObject_IsTrue(good);
1063 Py_DECREF(good);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02001064 if (ok > 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001065 return item;
1066 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02001067 if (ok == 0)
1068 lz->stop = 1;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001069 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001070}
1071
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001072PyDoc_STRVAR(takewhile_doc,
1073"takewhile(predicate, iterable) --> takewhile object\n\
1074\n\
1075Return successive entries from an iterable as long as the \n\
1076predicate evaluates to true for each entry.");
1077
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001078static PyTypeObject takewhile_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001079 PyVarObject_HEAD_INIT(NULL, 0)
1080 "itertools.takewhile", /* tp_name */
1081 sizeof(takewhileobject), /* tp_basicsize */
1082 0, /* tp_itemsize */
1083 /* methods */
1084 (destructor)takewhile_dealloc, /* tp_dealloc */
1085 0, /* tp_print */
1086 0, /* tp_getattr */
1087 0, /* tp_setattr */
1088 0, /* tp_compare */
1089 0, /* tp_repr */
1090 0, /* tp_as_number */
1091 0, /* tp_as_sequence */
1092 0, /* tp_as_mapping */
1093 0, /* tp_hash */
1094 0, /* tp_call */
1095 0, /* tp_str */
1096 PyObject_GenericGetAttr, /* tp_getattro */
1097 0, /* tp_setattro */
1098 0, /* tp_as_buffer */
1099 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1100 Py_TPFLAGS_BASETYPE, /* tp_flags */
1101 takewhile_doc, /* tp_doc */
1102 (traverseproc)takewhile_traverse, /* tp_traverse */
1103 0, /* tp_clear */
1104 0, /* tp_richcompare */
1105 0, /* tp_weaklistoffset */
1106 PyObject_SelfIter, /* tp_iter */
1107 (iternextfunc)takewhile_next, /* tp_iternext */
1108 0, /* tp_methods */
1109 0, /* tp_members */
1110 0, /* tp_getset */
1111 0, /* tp_base */
1112 0, /* tp_dict */
1113 0, /* tp_descr_get */
1114 0, /* tp_descr_set */
1115 0, /* tp_dictoffset */
1116 0, /* tp_init */
1117 0, /* tp_alloc */
1118 takewhile_new, /* tp_new */
1119 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001120};
1121
1122
1123/* islice object ************************************************************/
1124
1125typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001126 PyObject_HEAD
1127 PyObject *it;
1128 Py_ssize_t next;
1129 Py_ssize_t stop;
1130 Py_ssize_t step;
1131 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001132} isliceobject;
1133
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001134static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001135
1136static PyObject *
1137islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1138{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001139 PyObject *seq;
1140 Py_ssize_t start=0, stop=-1, step=1;
1141 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1142 Py_ssize_t numargs;
1143 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001144
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001145 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1146 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001147
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001148 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1149 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001150
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001151 numargs = PyTuple_Size(args);
1152 if (numargs == 2) {
1153 if (a1 != Py_None) {
1154 stop = PyInt_AsSsize_t(a1);
1155 if (stop == -1) {
1156 if (PyErr_Occurred())
1157 PyErr_Clear();
1158 PyErr_SetString(PyExc_ValueError,
1159 "Stop argument for islice() must be None or an integer: 0 <= x <= maxint.");
1160 return NULL;
1161 }
1162 }
1163 } else {
1164 if (a1 != Py_None)
1165 start = PyInt_AsSsize_t(a1);
1166 if (start == -1 && PyErr_Occurred())
1167 PyErr_Clear();
1168 if (a2 != Py_None) {
1169 stop = PyInt_AsSsize_t(a2);
1170 if (stop == -1) {
1171 if (PyErr_Occurred())
1172 PyErr_Clear();
1173 PyErr_SetString(PyExc_ValueError,
1174 "Stop argument for islice() must be None or an integer: 0 <= x <= maxint.");
1175 return NULL;
1176 }
1177 }
1178 }
1179 if (start<0 || stop<-1) {
1180 PyErr_SetString(PyExc_ValueError,
1181 "Indices for islice() must be None or an integer: 0 <= x <= maxint.");
1182 return NULL;
1183 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001184
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001185 if (a3 != NULL) {
1186 if (a3 != Py_None)
1187 step = PyInt_AsSsize_t(a3);
1188 if (step == -1 && PyErr_Occurred())
1189 PyErr_Clear();
1190 }
1191 if (step<1) {
1192 PyErr_SetString(PyExc_ValueError,
1193 "Step for islice() must be a positive integer or None.");
1194 return NULL;
1195 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001197 /* Get iterator. */
1198 it = PyObject_GetIter(seq);
1199 if (it == NULL)
1200 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001201
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001202 /* create isliceobject structure */
1203 lz = (isliceobject *)type->tp_alloc(type, 0);
1204 if (lz == NULL) {
1205 Py_DECREF(it);
1206 return NULL;
1207 }
1208 lz->it = it;
1209 lz->next = start;
1210 lz->stop = stop;
1211 lz->step = step;
1212 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001213
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001214 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001215}
1216
1217static void
1218islice_dealloc(isliceobject *lz)
1219{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001220 PyObject_GC_UnTrack(lz);
1221 Py_XDECREF(lz->it);
1222 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001223}
1224
1225static int
1226islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1227{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001228 Py_VISIT(lz->it);
1229 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001230}
1231
1232static PyObject *
1233islice_next(isliceobject *lz)
1234{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001235 PyObject *item;
1236 PyObject *it = lz->it;
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001237 Py_ssize_t stop = lz->stop;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001238 Py_ssize_t oldnext;
1239 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001240
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001241 if (it == NULL)
1242 return NULL;
1243
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001244 iternext = *Py_TYPE(it)->tp_iternext;
1245 while (lz->cnt < lz->next) {
1246 item = iternext(it);
1247 if (item == NULL)
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001248 goto empty;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001249 Py_DECREF(item);
1250 lz->cnt++;
1251 }
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001252 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001253 goto empty;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001254 item = iternext(it);
1255 if (item == NULL)
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001256 goto empty;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001257 lz->cnt++;
1258 oldnext = lz->next;
Mark Dickinsona96b0d12011-09-24 09:01:16 +01001259 /* The (size_t) cast below avoids the danger of undefined
1260 behaviour from signed integer overflow. */
1261 lz->next += (size_t)lz->step;
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001262 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1263 lz->next = stop;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001264 return item;
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001265
1266empty:
1267 Py_CLEAR(lz->it);
1268 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001269}
1270
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001271PyDoc_STRVAR(islice_doc,
1272"islice(iterable, [start,] stop [, step]) --> islice object\n\
1273\n\
1274Return an iterator whose next() method returns selected values from an\n\
1275iterable. If start is specified, will skip all preceding elements;\n\
1276otherwise, start defaults to zero. Step defaults to one. If\n\
1277specified as another value, step determines how many values are \n\
1278skipped between successive calls. Works like a slice() on a list\n\
1279but returns an iterator.");
1280
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001281static PyTypeObject islice_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001282 PyVarObject_HEAD_INIT(NULL, 0)
1283 "itertools.islice", /* tp_name */
1284 sizeof(isliceobject), /* tp_basicsize */
1285 0, /* tp_itemsize */
1286 /* methods */
1287 (destructor)islice_dealloc, /* tp_dealloc */
1288 0, /* tp_print */
1289 0, /* tp_getattr */
1290 0, /* tp_setattr */
1291 0, /* tp_compare */
1292 0, /* tp_repr */
1293 0, /* tp_as_number */
1294 0, /* tp_as_sequence */
1295 0, /* tp_as_mapping */
1296 0, /* tp_hash */
1297 0, /* tp_call */
1298 0, /* tp_str */
1299 PyObject_GenericGetAttr, /* tp_getattro */
1300 0, /* tp_setattro */
1301 0, /* tp_as_buffer */
1302 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1303 Py_TPFLAGS_BASETYPE, /* tp_flags */
1304 islice_doc, /* tp_doc */
1305 (traverseproc)islice_traverse, /* tp_traverse */
1306 0, /* tp_clear */
1307 0, /* tp_richcompare */
1308 0, /* tp_weaklistoffset */
1309 PyObject_SelfIter, /* tp_iter */
1310 (iternextfunc)islice_next, /* tp_iternext */
1311 0, /* tp_methods */
1312 0, /* tp_members */
1313 0, /* tp_getset */
1314 0, /* tp_base */
1315 0, /* tp_dict */
1316 0, /* tp_descr_get */
1317 0, /* tp_descr_set */
1318 0, /* tp_dictoffset */
1319 0, /* tp_init */
1320 0, /* tp_alloc */
1321 islice_new, /* tp_new */
1322 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001323};
1324
1325
1326/* starmap object ************************************************************/
1327
1328typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001329 PyObject_HEAD
1330 PyObject *func;
1331 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001332} starmapobject;
1333
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001334static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001335
1336static PyObject *
1337starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1338{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001339 PyObject *func, *seq;
1340 PyObject *it;
1341 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001342
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001343 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1344 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001345
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001346 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1347 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001348
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001349 /* Get iterator. */
1350 it = PyObject_GetIter(seq);
1351 if (it == NULL)
1352 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001353
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001354 /* create starmapobject structure */
1355 lz = (starmapobject *)type->tp_alloc(type, 0);
1356 if (lz == NULL) {
1357 Py_DECREF(it);
1358 return NULL;
1359 }
1360 Py_INCREF(func);
1361 lz->func = func;
1362 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001363
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001364 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001365}
1366
1367static void
1368starmap_dealloc(starmapobject *lz)
1369{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001370 PyObject_GC_UnTrack(lz);
1371 Py_XDECREF(lz->func);
1372 Py_XDECREF(lz->it);
1373 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001374}
1375
1376static int
1377starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1378{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001379 Py_VISIT(lz->it);
1380 Py_VISIT(lz->func);
1381 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001382}
1383
1384static PyObject *
1385starmap_next(starmapobject *lz)
1386{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001387 PyObject *args;
1388 PyObject *result;
1389 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001390
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001391 args = (*Py_TYPE(it)->tp_iternext)(it);
1392 if (args == NULL)
1393 return NULL;
1394 if (!PyTuple_CheckExact(args)) {
1395 PyObject *newargs = PySequence_Tuple(args);
1396 Py_DECREF(args);
1397 if (newargs == NULL)
1398 return NULL;
1399 args = newargs;
1400 }
1401 result = PyObject_Call(lz->func, args, NULL);
1402 Py_DECREF(args);
1403 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001404}
1405
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001406PyDoc_STRVAR(starmap_doc,
1407"starmap(function, sequence) --> starmap object\n\
1408\n\
1409Return an iterator whose values are returned from the function evaluated\n\
Serhiy Storchakac72e66a2015-11-02 15:06:09 +02001410with an argument tuple taken from the given sequence.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001411
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001412static PyTypeObject starmap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001413 PyVarObject_HEAD_INIT(NULL, 0)
1414 "itertools.starmap", /* tp_name */
1415 sizeof(starmapobject), /* tp_basicsize */
1416 0, /* tp_itemsize */
1417 /* methods */
1418 (destructor)starmap_dealloc, /* tp_dealloc */
1419 0, /* tp_print */
1420 0, /* tp_getattr */
1421 0, /* tp_setattr */
1422 0, /* tp_compare */
1423 0, /* tp_repr */
1424 0, /* tp_as_number */
1425 0, /* tp_as_sequence */
1426 0, /* tp_as_mapping */
1427 0, /* tp_hash */
1428 0, /* tp_call */
1429 0, /* tp_str */
1430 PyObject_GenericGetAttr, /* tp_getattro */
1431 0, /* tp_setattro */
1432 0, /* tp_as_buffer */
1433 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1434 Py_TPFLAGS_BASETYPE, /* tp_flags */
1435 starmap_doc, /* tp_doc */
1436 (traverseproc)starmap_traverse, /* tp_traverse */
1437 0, /* tp_clear */
1438 0, /* tp_richcompare */
1439 0, /* tp_weaklistoffset */
1440 PyObject_SelfIter, /* tp_iter */
1441 (iternextfunc)starmap_next, /* tp_iternext */
1442 0, /* tp_methods */
1443 0, /* tp_members */
1444 0, /* tp_getset */
1445 0, /* tp_base */
1446 0, /* tp_dict */
1447 0, /* tp_descr_get */
1448 0, /* tp_descr_set */
1449 0, /* tp_dictoffset */
1450 0, /* tp_init */
1451 0, /* tp_alloc */
1452 starmap_new, /* tp_new */
1453 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001454};
1455
1456
1457/* imap object ************************************************************/
1458
1459typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001460 PyObject_HEAD
1461 PyObject *iters;
1462 PyObject *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001463} imapobject;
1464
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001465static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001466
1467static PyObject *
1468imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1469{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001470 PyObject *it, *iters, *func;
1471 imapobject *lz;
1472 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001473
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001474 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
1475 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001476
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001477 numargs = PyTuple_Size(args);
1478 if (numargs < 2) {
1479 PyErr_SetString(PyExc_TypeError,
1480 "imap() must have at least two arguments.");
1481 return NULL;
1482 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001483
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001484 iters = PyTuple_New(numargs-1);
1485 if (iters == NULL)
1486 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001487
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001488 for (i=1 ; i<numargs ; i++) {
1489 /* Get iterator. */
1490 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1491 if (it == NULL) {
1492 Py_DECREF(iters);
1493 return NULL;
1494 }
1495 PyTuple_SET_ITEM(iters, i-1, it);
1496 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001497
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001498 /* create imapobject structure */
1499 lz = (imapobject *)type->tp_alloc(type, 0);
1500 if (lz == NULL) {
1501 Py_DECREF(iters);
1502 return NULL;
1503 }
1504 lz->iters = iters;
1505 func = PyTuple_GET_ITEM(args, 0);
1506 Py_INCREF(func);
1507 lz->func = func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001508
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001509 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001510}
1511
1512static void
1513imap_dealloc(imapobject *lz)
1514{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001515 PyObject_GC_UnTrack(lz);
1516 Py_XDECREF(lz->iters);
1517 Py_XDECREF(lz->func);
1518 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001519}
1520
1521static int
1522imap_traverse(imapobject *lz, visitproc visit, void *arg)
1523{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001524 Py_VISIT(lz->iters);
1525 Py_VISIT(lz->func);
1526 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001527}
1528
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001529/*
Raymond Hettinger2012f172003-02-07 05:32:58 +00001530imap() is an iterator version of __builtins__.map() except that it does
1531not have the None fill-in feature. That was intentionally left out for
1532the following reasons:
1533
1534 1) Itertools are designed to be easily combined and chained together.
1535 Having all tools stop with the shortest input is a unifying principle
1536 that makes it easier to combine finite iterators (supplying data) with
1537 infinite iterators like count() and repeat() (for supplying sequential
1538 or constant arguments to a function).
1539
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001540 2) In typical use cases for combining itertools, having one finite data
Raymond Hettinger2012f172003-02-07 05:32:58 +00001541 supplier run out before another is likely to be an error condition which
1542 should not pass silently by automatically supplying None.
1543
1544 3) The use cases for automatic None fill-in are rare -- not many functions
1545 do something useful when a parameter suddenly switches type and becomes
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001546 None.
Raymond Hettinger2012f172003-02-07 05:32:58 +00001547
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001548 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001549 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001550
1551 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1552*/
1553
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001554static PyObject *
1555imap_next(imapobject *lz)
1556{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001557 PyObject *val;
1558 PyObject *argtuple;
1559 PyObject *result;
1560 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001561
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001562 numargs = PyTuple_Size(lz->iters);
1563 argtuple = PyTuple_New(numargs);
1564 if (argtuple == NULL)
1565 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001566
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001567 for (i=0 ; i<numargs ; i++) {
1568 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1569 if (val == NULL) {
1570 Py_DECREF(argtuple);
1571 return NULL;
1572 }
1573 PyTuple_SET_ITEM(argtuple, i, val);
1574 }
1575 if (lz->func == Py_None)
1576 return argtuple;
1577 result = PyObject_Call(lz->func, argtuple, NULL);
1578 Py_DECREF(argtuple);
1579 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001580}
1581
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001582PyDoc_STRVAR(imap_doc,
1583"imap(func, *iterables) --> imap object\n\
1584\n\
1585Make an iterator that computes the function using arguments from\n\
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001586each of the iterables. Like map() except that it returns\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001587an iterator instead of a list and that it stops when the shortest\n\
1588iterable is exhausted instead of filling in None for shorter\n\
1589iterables.");
1590
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001591static PyTypeObject imap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001592 PyVarObject_HEAD_INIT(NULL, 0)
1593 "itertools.imap", /* tp_name */
1594 sizeof(imapobject), /* tp_basicsize */
1595 0, /* tp_itemsize */
1596 /* methods */
1597 (destructor)imap_dealloc, /* tp_dealloc */
1598 0, /* tp_print */
1599 0, /* tp_getattr */
1600 0, /* tp_setattr */
1601 0, /* tp_compare */
1602 0, /* tp_repr */
1603 0, /* tp_as_number */
1604 0, /* tp_as_sequence */
1605 0, /* tp_as_mapping */
1606 0, /* tp_hash */
1607 0, /* tp_call */
1608 0, /* tp_str */
1609 PyObject_GenericGetAttr, /* tp_getattro */
1610 0, /* tp_setattro */
1611 0, /* tp_as_buffer */
1612 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1613 Py_TPFLAGS_BASETYPE, /* tp_flags */
1614 imap_doc, /* tp_doc */
1615 (traverseproc)imap_traverse, /* tp_traverse */
1616 0, /* tp_clear */
1617 0, /* tp_richcompare */
1618 0, /* tp_weaklistoffset */
1619 PyObject_SelfIter, /* tp_iter */
1620 (iternextfunc)imap_next, /* tp_iternext */
1621 0, /* tp_methods */
1622 0, /* tp_members */
1623 0, /* tp_getset */
1624 0, /* tp_base */
1625 0, /* tp_dict */
1626 0, /* tp_descr_get */
1627 0, /* tp_descr_set */
1628 0, /* tp_dictoffset */
1629 0, /* tp_init */
1630 0, /* tp_alloc */
1631 imap_new, /* tp_new */
1632 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001633};
1634
1635
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001636/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001637
1638typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001639 PyObject_HEAD
1640 PyObject *source; /* Iterator over input iterables */
1641 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001642} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001643
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001644static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001645
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001646static PyObject *
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001647chain_new_internal(PyTypeObject *type, PyObject *source)
1648{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001649 chainobject *lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001650
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001651 lz = (chainobject *)type->tp_alloc(type, 0);
1652 if (lz == NULL) {
1653 Py_DECREF(source);
1654 return NULL;
1655 }
1656
1657 lz->source = source;
1658 lz->active = NULL;
1659 return (PyObject *)lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001660}
1661
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001662static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001663chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001664{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001665 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001666
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001667 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1668 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001669
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001670 source = PyObject_GetIter(args);
1671 if (source == NULL)
1672 return NULL;
1673
1674 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001675}
1676
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001677static PyObject *
1678chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1679{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001680 PyObject *source;
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001681
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001682 source = PyObject_GetIter(arg);
1683 if (source == NULL)
1684 return NULL;
1685
1686 return chain_new_internal(type, source);
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001687}
1688
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001689static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001690chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001691{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001692 PyObject_GC_UnTrack(lz);
1693 Py_XDECREF(lz->active);
1694 Py_XDECREF(lz->source);
1695 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001696}
1697
Raymond Hettinger2012f172003-02-07 05:32:58 +00001698static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001699chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001700{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001701 Py_VISIT(lz->source);
1702 Py_VISIT(lz->active);
1703 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001704}
1705
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001706static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001707chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001708{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001709 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001710
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001711 if (lz->source == NULL)
1712 return NULL; /* already stopped */
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001713
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001714 if (lz->active == NULL) {
1715 PyObject *iterable = PyIter_Next(lz->source);
1716 if (iterable == NULL) {
1717 Py_CLEAR(lz->source);
1718 return NULL; /* no more input sources */
1719 }
1720 lz->active = PyObject_GetIter(iterable);
1721 Py_DECREF(iterable);
1722 if (lz->active == NULL) {
1723 Py_CLEAR(lz->source);
1724 return NULL; /* input not iterable */
1725 }
1726 }
1727 item = PyIter_Next(lz->active);
1728 if (item != NULL)
1729 return item;
1730 if (PyErr_Occurred()) {
1731 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1732 PyErr_Clear();
1733 else
1734 return NULL; /* input raised an exception */
1735 }
1736 Py_CLEAR(lz->active);
1737 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001738}
1739
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001740PyDoc_STRVAR(chain_doc,
1741"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001742\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001743Return a chain object whose .next() method returns elements from the\n\
1744first iterable until it is exhausted, then elements from the next\n\
1745iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001746
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001747PyDoc_STRVAR(chain_from_iterable_doc,
1748"chain.from_iterable(iterable) --> chain object\n\
1749\n\
1750Alternate chain() contructor taking a single iterable argument\n\
1751that evaluates lazily.");
1752
1753static PyMethodDef chain_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001754 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1755 chain_from_iterable_doc},
1756 {NULL, NULL} /* sentinel */
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001757};
1758
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001759static PyTypeObject chain_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001760 PyVarObject_HEAD_INIT(NULL, 0)
1761 "itertools.chain", /* tp_name */
1762 sizeof(chainobject), /* tp_basicsize */
1763 0, /* tp_itemsize */
1764 /* methods */
1765 (destructor)chain_dealloc, /* tp_dealloc */
1766 0, /* tp_print */
1767 0, /* tp_getattr */
1768 0, /* tp_setattr */
1769 0, /* tp_compare */
1770 0, /* tp_repr */
1771 0, /* tp_as_number */
1772 0, /* tp_as_sequence */
1773 0, /* tp_as_mapping */
1774 0, /* tp_hash */
1775 0, /* tp_call */
1776 0, /* tp_str */
1777 PyObject_GenericGetAttr, /* tp_getattro */
1778 0, /* tp_setattro */
1779 0, /* tp_as_buffer */
1780 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1781 Py_TPFLAGS_BASETYPE, /* tp_flags */
1782 chain_doc, /* tp_doc */
1783 (traverseproc)chain_traverse, /* tp_traverse */
1784 0, /* tp_clear */
1785 0, /* tp_richcompare */
1786 0, /* tp_weaklistoffset */
1787 PyObject_SelfIter, /* tp_iter */
1788 (iternextfunc)chain_next, /* tp_iternext */
1789 chain_methods, /* tp_methods */
1790 0, /* tp_members */
1791 0, /* tp_getset */
1792 0, /* tp_base */
1793 0, /* tp_dict */
1794 0, /* tp_descr_get */
1795 0, /* tp_descr_set */
1796 0, /* tp_dictoffset */
1797 0, /* tp_init */
1798 0, /* tp_alloc */
1799 chain_new, /* tp_new */
1800 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001801};
1802
1803
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001804/* product object ************************************************************/
1805
1806typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001807 PyObject_HEAD
1808 PyObject *pools; /* tuple of pool tuples */
1809 Py_ssize_t *indices; /* one index per pool */
1810 PyObject *result; /* most recently returned result tuple */
1811 int stopped; /* set to 1 when the product iterator is exhausted */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001812} productobject;
1813
1814static PyTypeObject product_type;
1815
1816static PyObject *
1817product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1818{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001819 productobject *lz;
1820 Py_ssize_t nargs, npools, repeat=1;
1821 PyObject *pools = NULL;
1822 Py_ssize_t *indices = NULL;
1823 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001824
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001825 if (kwds != NULL) {
1826 char *kwlist[] = {"repeat", 0};
1827 PyObject *tmpargs = PyTuple_New(0);
1828 if (tmpargs == NULL)
1829 return NULL;
1830 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1831 Py_DECREF(tmpargs);
1832 return NULL;
1833 }
1834 Py_DECREF(tmpargs);
1835 if (repeat < 0) {
1836 PyErr_SetString(PyExc_ValueError,
1837 "repeat argument cannot be negative");
1838 return NULL;
1839 }
1840 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001841
Benjamin Petersondda91212015-02-01 21:34:07 -05001842 assert(PyTuple_CheckExact(args));
1843 if (repeat == 0) {
1844 nargs = 0;
1845 } else {
1846 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02001847 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Petersondda91212015-02-01 21:34:07 -05001848 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
1849 return NULL;
1850 }
1851 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001852 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001853
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02001854 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001855 if (indices == NULL) {
1856 PyErr_NoMemory();
1857 goto error;
1858 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001859
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001860 pools = PyTuple_New(npools);
1861 if (pools == NULL)
1862 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001863
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001864 for (i=0; i < nargs ; ++i) {
1865 PyObject *item = PyTuple_GET_ITEM(args, i);
1866 PyObject *pool = PySequence_Tuple(item);
1867 if (pool == NULL)
1868 goto error;
1869 PyTuple_SET_ITEM(pools, i, pool);
1870 indices[i] = 0;
1871 }
1872 for ( ; i < npools; ++i) {
1873 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1874 Py_INCREF(pool);
1875 PyTuple_SET_ITEM(pools, i, pool);
1876 indices[i] = 0;
1877 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001878
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001879 /* create productobject structure */
1880 lz = (productobject *)type->tp_alloc(type, 0);
1881 if (lz == NULL)
1882 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001883
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001884 lz->pools = pools;
1885 lz->indices = indices;
1886 lz->result = NULL;
1887 lz->stopped = 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001888
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001889 return (PyObject *)lz;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001890
1891error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001892 if (indices != NULL)
1893 PyMem_Free(indices);
1894 Py_XDECREF(pools);
1895 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001896}
1897
1898static void
1899product_dealloc(productobject *lz)
1900{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001901 PyObject_GC_UnTrack(lz);
1902 Py_XDECREF(lz->pools);
1903 Py_XDECREF(lz->result);
1904 if (lz->indices != NULL)
1905 PyMem_Free(lz->indices);
1906 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001907}
1908
1909static int
1910product_traverse(productobject *lz, visitproc visit, void *arg)
1911{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001912 Py_VISIT(lz->pools);
1913 Py_VISIT(lz->result);
1914 return 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001915}
1916
1917static PyObject *
1918product_next(productobject *lz)
1919{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001920 PyObject *pool;
1921 PyObject *elem;
1922 PyObject *oldelem;
1923 PyObject *pools = lz->pools;
1924 PyObject *result = lz->result;
1925 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1926 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001927
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001928 if (lz->stopped)
1929 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001930
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001931 if (result == NULL) {
1932 /* On the first pass, return an initial tuple filled with the
1933 first element from each pool. */
1934 result = PyTuple_New(npools);
1935 if (result == NULL)
1936 goto empty;
1937 lz->result = result;
1938 for (i=0; i < npools; i++) {
1939 pool = PyTuple_GET_ITEM(pools, i);
1940 if (PyTuple_GET_SIZE(pool) == 0)
1941 goto empty;
1942 elem = PyTuple_GET_ITEM(pool, 0);
1943 Py_INCREF(elem);
1944 PyTuple_SET_ITEM(result, i, elem);
1945 }
1946 } else {
1947 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001948
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001949 /* Copy the previous result tuple or re-use it if available */
1950 if (Py_REFCNT(result) > 1) {
1951 PyObject *old_result = result;
1952 result = PyTuple_New(npools);
1953 if (result == NULL)
1954 goto empty;
1955 lz->result = result;
1956 for (i=0; i < npools; i++) {
1957 elem = PyTuple_GET_ITEM(old_result, i);
1958 Py_INCREF(elem);
1959 PyTuple_SET_ITEM(result, i, elem);
1960 }
1961 Py_DECREF(old_result);
1962 }
1963 /* Now, we've got the only copy so we can update it in-place */
1964 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001965
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001966 /* Update the pool indices right-to-left. Only advance to the
1967 next pool when the previous one rolls-over */
1968 for (i=npools-1 ; i >= 0 ; i--) {
1969 pool = PyTuple_GET_ITEM(pools, i);
1970 indices[i]++;
1971 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1972 /* Roll-over and advance to next pool */
1973 indices[i] = 0;
1974 elem = PyTuple_GET_ITEM(pool, 0);
1975 Py_INCREF(elem);
1976 oldelem = PyTuple_GET_ITEM(result, i);
1977 PyTuple_SET_ITEM(result, i, elem);
1978 Py_DECREF(oldelem);
1979 } else {
1980 /* No rollover. Just increment and stop here. */
1981 elem = PyTuple_GET_ITEM(pool, indices[i]);
1982 Py_INCREF(elem);
1983 oldelem = PyTuple_GET_ITEM(result, i);
1984 PyTuple_SET_ITEM(result, i, elem);
1985 Py_DECREF(oldelem);
1986 break;
1987 }
1988 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00001989
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001990 /* If i is negative, then the indices have all rolled-over
1991 and we're done. */
1992 if (i < 0)
1993 goto empty;
1994 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001995
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001996 Py_INCREF(result);
1997 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001998
1999empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002000 lz->stopped = 1;
2001 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002002}
2003
2004PyDoc_STRVAR(product_doc,
2005"product(*iterables) --> product object\n\
2006\n\
Raymond Hettinger73d79632008-02-23 02:20:41 +00002007Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002008For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2009The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2010cycle in a manner similar to an odometer (with the rightmost element changing\n\
2011on every iteration).\n\n\
Raymond Hettinger94034ea2009-02-09 18:39:41 +00002012To compute the product of an iterable with itself, specify the number\n\
2013of repetitions with the optional repeat keyword argument. For example,\n\
2014product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002015product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2016product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2017
2018static PyTypeObject product_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002019 PyVarObject_HEAD_INIT(NULL, 0)
2020 "itertools.product", /* tp_name */
2021 sizeof(productobject), /* tp_basicsize */
2022 0, /* tp_itemsize */
2023 /* methods */
2024 (destructor)product_dealloc, /* tp_dealloc */
2025 0, /* tp_print */
2026 0, /* tp_getattr */
2027 0, /* tp_setattr */
2028 0, /* tp_compare */
2029 0, /* tp_repr */
2030 0, /* tp_as_number */
2031 0, /* tp_as_sequence */
2032 0, /* tp_as_mapping */
2033 0, /* tp_hash */
2034 0, /* tp_call */
2035 0, /* tp_str */
2036 PyObject_GenericGetAttr, /* tp_getattro */
2037 0, /* tp_setattro */
2038 0, /* tp_as_buffer */
2039 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2040 Py_TPFLAGS_BASETYPE, /* tp_flags */
2041 product_doc, /* tp_doc */
2042 (traverseproc)product_traverse, /* tp_traverse */
2043 0, /* tp_clear */
2044 0, /* tp_richcompare */
2045 0, /* tp_weaklistoffset */
2046 PyObject_SelfIter, /* tp_iter */
2047 (iternextfunc)product_next, /* tp_iternext */
2048 0, /* tp_methods */
2049 0, /* tp_members */
2050 0, /* tp_getset */
2051 0, /* tp_base */
2052 0, /* tp_dict */
2053 0, /* tp_descr_get */
2054 0, /* tp_descr_set */
2055 0, /* tp_dictoffset */
2056 0, /* tp_init */
2057 0, /* tp_alloc */
2058 product_new, /* tp_new */
2059 PyObject_GC_Del, /* tp_free */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002060};
2061
2062
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002063/* combinations object ************************************************************/
2064
2065typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002066 PyObject_HEAD
2067 PyObject *pool; /* input converted to a tuple */
2068 Py_ssize_t *indices; /* one index per result element */
2069 PyObject *result; /* most recently returned result tuple */
2070 Py_ssize_t r; /* size of result tuple */
2071 int stopped; /* set to 1 when the combinations iterator is exhausted */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002072} combinationsobject;
2073
2074static PyTypeObject combinations_type;
2075
2076static PyObject *
2077combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2078{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002079 combinationsobject *co;
2080 Py_ssize_t n;
2081 Py_ssize_t r;
2082 PyObject *pool = NULL;
2083 PyObject *iterable = NULL;
2084 Py_ssize_t *indices = NULL;
2085 Py_ssize_t i;
2086 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002087
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002088 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2089 &iterable, &r))
2090 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002091
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002092 pool = PySequence_Tuple(iterable);
2093 if (pool == NULL)
2094 goto error;
2095 n = PyTuple_GET_SIZE(pool);
2096 if (r < 0) {
2097 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2098 goto error;
2099 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002100
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02002101 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002102 if (indices == NULL) {
2103 PyErr_NoMemory();
2104 goto error;
2105 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002106
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002107 for (i=0 ; i<r ; i++)
2108 indices[i] = i;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002109
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002110 /* create combinationsobject structure */
2111 co = (combinationsobject *)type->tp_alloc(type, 0);
2112 if (co == NULL)
2113 goto error;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002114
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002115 co->pool = pool;
2116 co->indices = indices;
2117 co->result = NULL;
2118 co->r = r;
2119 co->stopped = r > n ? 1 : 0;
2120
2121 return (PyObject *)co;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002122
2123error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002124 if (indices != NULL)
2125 PyMem_Free(indices);
2126 Py_XDECREF(pool);
2127 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002128}
2129
2130static void
2131combinations_dealloc(combinationsobject *co)
2132{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002133 PyObject_GC_UnTrack(co);
2134 Py_XDECREF(co->pool);
2135 Py_XDECREF(co->result);
2136 if (co->indices != NULL)
2137 PyMem_Free(co->indices);
2138 Py_TYPE(co)->tp_free(co);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002139}
2140
2141static int
2142combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2143{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002144 Py_VISIT(co->pool);
2145 Py_VISIT(co->result);
2146 return 0;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002147}
2148
2149static PyObject *
2150combinations_next(combinationsobject *co)
2151{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002152 PyObject *elem;
2153 PyObject *oldelem;
2154 PyObject *pool = co->pool;
2155 Py_ssize_t *indices = co->indices;
2156 PyObject *result = co->result;
2157 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2158 Py_ssize_t r = co->r;
2159 Py_ssize_t i, j, index;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002160
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002161 if (co->stopped)
2162 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002163
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002164 if (result == NULL) {
2165 /* On the first pass, initialize result tuple using the indices */
2166 result = PyTuple_New(r);
2167 if (result == NULL)
2168 goto empty;
2169 co->result = result;
2170 for (i=0; i<r ; i++) {
2171 index = indices[i];
2172 elem = PyTuple_GET_ITEM(pool, index);
2173 Py_INCREF(elem);
2174 PyTuple_SET_ITEM(result, i, elem);
2175 }
2176 } else {
2177 /* Copy the previous result tuple or re-use it if available */
2178 if (Py_REFCNT(result) > 1) {
2179 PyObject *old_result = result;
2180 result = PyTuple_New(r);
2181 if (result == NULL)
2182 goto empty;
2183 co->result = result;
2184 for (i=0; i<r ; i++) {
2185 elem = PyTuple_GET_ITEM(old_result, i);
2186 Py_INCREF(elem);
2187 PyTuple_SET_ITEM(result, i, elem);
2188 }
2189 Py_DECREF(old_result);
2190 }
2191 /* Now, we've got the only copy so we can update it in-place
2192 * CPython's empty tuple is a singleton and cached in
2193 * PyTuple's freelist.
2194 */
2195 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002196
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002197 /* Scan indices right-to-left until finding one that is not
2198 at its maximum (i + n - r). */
2199 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2200 ;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002201
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002202 /* If i is negative, then the indices are all at
2203 their maximum value and we're done. */
2204 if (i < 0)
2205 goto empty;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002206
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002207 /* Increment the current index which we know is not at its
2208 maximum. Then move back to the right setting each index
2209 to its lowest possible value (one higher than the index
2210 to its left -- this maintains the sort order invariant). */
2211 indices[i]++;
2212 for (j=i+1 ; j<r ; j++)
2213 indices[j] = indices[j-1] + 1;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002214
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002215 /* Update the result tuple for the new indices
2216 starting with i, the leftmost index that changed */
2217 for ( ; i<r ; i++) {
2218 index = indices[i];
2219 elem = PyTuple_GET_ITEM(pool, index);
2220 Py_INCREF(elem);
2221 oldelem = PyTuple_GET_ITEM(result, i);
2222 PyTuple_SET_ITEM(result, i, elem);
2223 Py_DECREF(oldelem);
2224 }
2225 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002226
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002227 Py_INCREF(result);
2228 return result;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002229
2230empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002231 co->stopped = 1;
2232 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002233}
2234
2235PyDoc_STRVAR(combinations_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002236"combinations(iterable, r) --> combinations object\n\
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002237\n\
2238Return successive r-length combinations of elements in the iterable.\n\n\
2239combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2240
2241static PyTypeObject combinations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002242 PyVarObject_HEAD_INIT(NULL, 0)
2243 "itertools.combinations", /* tp_name */
2244 sizeof(combinationsobject), /* tp_basicsize */
2245 0, /* tp_itemsize */
2246 /* methods */
2247 (destructor)combinations_dealloc, /* tp_dealloc */
2248 0, /* tp_print */
2249 0, /* tp_getattr */
2250 0, /* tp_setattr */
2251 0, /* tp_compare */
2252 0, /* tp_repr */
2253 0, /* tp_as_number */
2254 0, /* tp_as_sequence */
2255 0, /* tp_as_mapping */
2256 0, /* tp_hash */
2257 0, /* tp_call */
2258 0, /* tp_str */
2259 PyObject_GenericGetAttr, /* tp_getattro */
2260 0, /* tp_setattro */
2261 0, /* tp_as_buffer */
2262 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2263 Py_TPFLAGS_BASETYPE, /* tp_flags */
2264 combinations_doc, /* tp_doc */
2265 (traverseproc)combinations_traverse, /* tp_traverse */
2266 0, /* tp_clear */
2267 0, /* tp_richcompare */
2268 0, /* tp_weaklistoffset */
2269 PyObject_SelfIter, /* tp_iter */
2270 (iternextfunc)combinations_next, /* tp_iternext */
2271 0, /* tp_methods */
2272 0, /* tp_members */
2273 0, /* tp_getset */
2274 0, /* tp_base */
2275 0, /* tp_dict */
2276 0, /* tp_descr_get */
2277 0, /* tp_descr_set */
2278 0, /* tp_dictoffset */
2279 0, /* tp_init */
2280 0, /* tp_alloc */
2281 combinations_new, /* tp_new */
2282 PyObject_GC_Del, /* tp_free */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002283};
2284
2285
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002286/* combinations with replacement object *******************************************/
2287
2288/* Equivalent to:
2289
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002290 def combinations_with_replacement(iterable, r):
2291 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2292 # number items returned: (n+r-1)! / r! / (n-1)!
2293 pool = tuple(iterable)
2294 n = len(pool)
2295 indices = [0] * r
2296 yield tuple(pool[i] for i in indices)
2297 while 1:
2298 for i in reversed(range(r)):
2299 if indices[i] != n - 1:
2300 break
2301 else:
2302 return
2303 indices[i:] = [indices[i] + 1] * (r - i)
2304 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002305
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002306 def combinations_with_replacement2(iterable, r):
2307 'Alternate version that filters from product()'
2308 pool = tuple(iterable)
2309 n = len(pool)
2310 for indices in product(range(n), repeat=r):
2311 if sorted(indices) == list(indices):
2312 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002313*/
2314typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002315 PyObject_HEAD
2316 PyObject *pool; /* input converted to a tuple */
2317 Py_ssize_t *indices; /* one index per result element */
2318 PyObject *result; /* most recently returned result tuple */
2319 Py_ssize_t r; /* size of result tuple */
2320 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002321} cwrobject;
2322
2323static PyTypeObject cwr_type;
2324
2325static PyObject *
2326cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2327{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002328 cwrobject *co;
2329 Py_ssize_t n;
2330 Py_ssize_t r;
2331 PyObject *pool = NULL;
2332 PyObject *iterable = NULL;
2333 Py_ssize_t *indices = NULL;
2334 Py_ssize_t i;
2335 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002336
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002337 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2338 &iterable, &r))
2339 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002340
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002341 pool = PySequence_Tuple(iterable);
2342 if (pool == NULL)
2343 goto error;
2344 n = PyTuple_GET_SIZE(pool);
2345 if (r < 0) {
2346 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2347 goto error;
2348 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002349
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02002350 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002351 if (indices == NULL) {
2352 PyErr_NoMemory();
2353 goto error;
2354 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002355
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002356 for (i=0 ; i<r ; i++)
2357 indices[i] = 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002358
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002359 /* create cwrobject structure */
2360 co = (cwrobject *)type->tp_alloc(type, 0);
2361 if (co == NULL)
2362 goto error;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002363
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002364 co->pool = pool;
2365 co->indices = indices;
2366 co->result = NULL;
2367 co->r = r;
2368 co->stopped = !n && r;
2369
2370 return (PyObject *)co;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002371
2372error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002373 if (indices != NULL)
2374 PyMem_Free(indices);
2375 Py_XDECREF(pool);
2376 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002377}
2378
2379static void
2380cwr_dealloc(cwrobject *co)
2381{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002382 PyObject_GC_UnTrack(co);
2383 Py_XDECREF(co->pool);
2384 Py_XDECREF(co->result);
2385 if (co->indices != NULL)
2386 PyMem_Free(co->indices);
2387 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002388}
2389
2390static int
2391cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2392{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002393 Py_VISIT(co->pool);
2394 Py_VISIT(co->result);
2395 return 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002396}
2397
2398static PyObject *
2399cwr_next(cwrobject *co)
2400{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002401 PyObject *elem;
2402 PyObject *oldelem;
2403 PyObject *pool = co->pool;
2404 Py_ssize_t *indices = co->indices;
2405 PyObject *result = co->result;
2406 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2407 Py_ssize_t r = co->r;
2408 Py_ssize_t i, j, index;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002409
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002410 if (co->stopped)
2411 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002412
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002413 if (result == NULL) {
2414 /* On the first pass, initialize result tuple using the indices */
2415 result = PyTuple_New(r);
2416 if (result == NULL)
2417 goto empty;
2418 co->result = result;
2419 for (i=0; i<r ; i++) {
2420 index = indices[i];
2421 elem = PyTuple_GET_ITEM(pool, index);
2422 Py_INCREF(elem);
2423 PyTuple_SET_ITEM(result, i, elem);
2424 }
2425 } else {
2426 /* Copy the previous result tuple or re-use it if available */
2427 if (Py_REFCNT(result) > 1) {
2428 PyObject *old_result = result;
2429 result = PyTuple_New(r);
2430 if (result == NULL)
2431 goto empty;
2432 co->result = result;
2433 for (i=0; i<r ; i++) {
2434 elem = PyTuple_GET_ITEM(old_result, i);
2435 Py_INCREF(elem);
2436 PyTuple_SET_ITEM(result, i, elem);
2437 }
2438 Py_DECREF(old_result);
2439 }
2440 /* Now, we've got the only copy so we can update it in-place CPython's
2441 empty tuple is a singleton and cached in PyTuple's freelist. */
2442 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002443
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002444 /* Scan indices right-to-left until finding one that is not
2445 * at its maximum (n-1). */
2446 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2447 ;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002448
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002449 /* If i is negative, then the indices are all at
2450 their maximum value and we're done. */
2451 if (i < 0)
2452 goto empty;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002453
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002454 /* Increment the current index which we know is not at its
2455 maximum. Then set all to the right to the same value. */
2456 indices[i]++;
2457 for (j=i+1 ; j<r ; j++)
2458 indices[j] = indices[j-1];
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002459
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002460 /* Update the result tuple for the new indices
2461 starting with i, the leftmost index that changed */
2462 for ( ; i<r ; i++) {
2463 index = indices[i];
2464 elem = PyTuple_GET_ITEM(pool, index);
2465 Py_INCREF(elem);
2466 oldelem = PyTuple_GET_ITEM(result, i);
2467 PyTuple_SET_ITEM(result, i, elem);
2468 Py_DECREF(oldelem);
2469 }
2470 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002471
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002472 Py_INCREF(result);
2473 return result;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002474
2475empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002476 co->stopped = 1;
2477 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002478}
2479
2480PyDoc_STRVAR(cwr_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002481"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002482\n\
2483Return successive r-length combinations of elements in the iterable\n\
2484allowing individual elements to have successive repeats.\n\
2485combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2486
2487static PyTypeObject cwr_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002488 PyVarObject_HEAD_INIT(NULL, 0)
2489 "itertools.combinations_with_replacement", /* tp_name */
2490 sizeof(cwrobject), /* tp_basicsize */
2491 0, /* tp_itemsize */
2492 /* methods */
2493 (destructor)cwr_dealloc, /* tp_dealloc */
2494 0, /* tp_print */
2495 0, /* tp_getattr */
2496 0, /* tp_setattr */
2497 0, /* tp_compare */
2498 0, /* tp_repr */
2499 0, /* tp_as_number */
2500 0, /* tp_as_sequence */
2501 0, /* tp_as_mapping */
2502 0, /* tp_hash */
2503 0, /* tp_call */
2504 0, /* tp_str */
2505 PyObject_GenericGetAttr, /* tp_getattro */
2506 0, /* tp_setattro */
2507 0, /* tp_as_buffer */
2508 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2509 Py_TPFLAGS_BASETYPE, /* tp_flags */
2510 cwr_doc, /* tp_doc */
2511 (traverseproc)cwr_traverse, /* tp_traverse */
2512 0, /* tp_clear */
2513 0, /* tp_richcompare */
2514 0, /* tp_weaklistoffset */
2515 PyObject_SelfIter, /* tp_iter */
2516 (iternextfunc)cwr_next, /* tp_iternext */
2517 0, /* tp_methods */
2518 0, /* tp_members */
2519 0, /* tp_getset */
2520 0, /* tp_base */
2521 0, /* tp_dict */
2522 0, /* tp_descr_get */
2523 0, /* tp_descr_set */
2524 0, /* tp_dictoffset */
2525 0, /* tp_init */
2526 0, /* tp_alloc */
2527 cwr_new, /* tp_new */
2528 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002529};
2530
2531
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002532/* permutations object ************************************************************
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002533
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002534def permutations(iterable, r=None):
2535 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2536 pool = tuple(iterable)
2537 n = len(pool)
2538 r = n if r is None else r
2539 indices = range(n)
2540 cycles = range(n-r+1, n+1)[::-1]
2541 yield tuple(pool[i] for i in indices[:r])
2542 while n:
Serhiy Storchakaf0aa88f2015-06-11 00:06:27 +03002543 for i in reversed(range(r)):
2544 cycles[i] -= 1
2545 if cycles[i] == 0:
2546 indices[i:] = indices[i+1:] + indices[i:i+1]
2547 cycles[i] = n - i
2548 else:
2549 j = cycles[i]
2550 indices[i], indices[-j] = indices[-j], indices[i]
2551 yield tuple(pool[i] for i in indices[:r])
2552 break
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002553 else:
Serhiy Storchakaf0aa88f2015-06-11 00:06:27 +03002554 return
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002555*/
2556
2557typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002558 PyObject_HEAD
2559 PyObject *pool; /* input converted to a tuple */
2560 Py_ssize_t *indices; /* one index per element in the pool */
2561 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2562 PyObject *result; /* most recently returned result tuple */
2563 Py_ssize_t r; /* size of result tuple */
2564 int stopped; /* set to 1 when the permutations iterator is exhausted */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002565} permutationsobject;
2566
2567static PyTypeObject permutations_type;
2568
2569static PyObject *
2570permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2571{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002572 permutationsobject *po;
2573 Py_ssize_t n;
2574 Py_ssize_t r;
2575 PyObject *robj = Py_None;
2576 PyObject *pool = NULL;
2577 PyObject *iterable = NULL;
2578 Py_ssize_t *indices = NULL;
2579 Py_ssize_t *cycles = NULL;
2580 Py_ssize_t i;
2581 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002582
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002583 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2584 &iterable, &robj))
2585 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002586
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002587 pool = PySequence_Tuple(iterable);
2588 if (pool == NULL)
2589 goto error;
2590 n = PyTuple_GET_SIZE(pool);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002591
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002592 r = n;
2593 if (robj != Py_None) {
2594 r = PyInt_AsSsize_t(robj);
2595 if (r == -1 && PyErr_Occurred())
2596 goto error;
2597 }
2598 if (r < 0) {
2599 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2600 goto error;
2601 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002602
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02002603 indices = PyMem_New(Py_ssize_t, n);
2604 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002605 if (indices == NULL || cycles == NULL) {
2606 PyErr_NoMemory();
2607 goto error;
2608 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002609
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002610 for (i=0 ; i<n ; i++)
2611 indices[i] = i;
2612 for (i=0 ; i<r ; i++)
2613 cycles[i] = n - i;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002614
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002615 /* create permutationsobject structure */
2616 po = (permutationsobject *)type->tp_alloc(type, 0);
2617 if (po == NULL)
2618 goto error;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002619
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002620 po->pool = pool;
2621 po->indices = indices;
2622 po->cycles = cycles;
2623 po->result = NULL;
2624 po->r = r;
2625 po->stopped = r > n ? 1 : 0;
2626
2627 return (PyObject *)po;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002628
2629error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002630 if (indices != NULL)
2631 PyMem_Free(indices);
2632 if (cycles != NULL)
2633 PyMem_Free(cycles);
2634 Py_XDECREF(pool);
2635 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002636}
2637
2638static void
2639permutations_dealloc(permutationsobject *po)
2640{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002641 PyObject_GC_UnTrack(po);
2642 Py_XDECREF(po->pool);
2643 Py_XDECREF(po->result);
2644 PyMem_Free(po->indices);
2645 PyMem_Free(po->cycles);
2646 Py_TYPE(po)->tp_free(po);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002647}
2648
2649static int
2650permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2651{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002652 Py_VISIT(po->pool);
2653 Py_VISIT(po->result);
2654 return 0;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002655}
2656
2657static PyObject *
2658permutations_next(permutationsobject *po)
2659{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002660 PyObject *elem;
2661 PyObject *oldelem;
2662 PyObject *pool = po->pool;
2663 Py_ssize_t *indices = po->indices;
2664 Py_ssize_t *cycles = po->cycles;
2665 PyObject *result = po->result;
2666 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2667 Py_ssize_t r = po->r;
2668 Py_ssize_t i, j, k, index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002669
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002670 if (po->stopped)
2671 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002672
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002673 if (result == NULL) {
2674 /* On the first pass, initialize result tuple using the indices */
2675 result = PyTuple_New(r);
2676 if (result == NULL)
2677 goto empty;
2678 po->result = result;
2679 for (i=0; i<r ; i++) {
2680 index = indices[i];
2681 elem = PyTuple_GET_ITEM(pool, index);
2682 Py_INCREF(elem);
2683 PyTuple_SET_ITEM(result, i, elem);
2684 }
2685 } else {
2686 if (n == 0)
2687 goto empty;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002688
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002689 /* Copy the previous result tuple or re-use it if available */
2690 if (Py_REFCNT(result) > 1) {
2691 PyObject *old_result = result;
2692 result = PyTuple_New(r);
2693 if (result == NULL)
2694 goto empty;
2695 po->result = result;
2696 for (i=0; i<r ; i++) {
2697 elem = PyTuple_GET_ITEM(old_result, i);
2698 Py_INCREF(elem);
2699 PyTuple_SET_ITEM(result, i, elem);
2700 }
2701 Py_DECREF(old_result);
2702 }
2703 /* Now, we've got the only copy so we can update it in-place */
2704 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002705
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002706 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2707 for (i=r-1 ; i>=0 ; i--) {
2708 cycles[i] -= 1;
2709 if (cycles[i] == 0) {
2710 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2711 index = indices[i];
2712 for (j=i ; j<n-1 ; j++)
2713 indices[j] = indices[j+1];
2714 indices[n-1] = index;
2715 cycles[i] = n - i;
2716 } else {
2717 j = cycles[i];
2718 index = indices[i];
2719 indices[i] = indices[n-j];
2720 indices[n-j] = index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002721
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002722 for (k=i; k<r ; k++) {
2723 /* start with i, the leftmost element that changed */
2724 /* yield tuple(pool[k] for k in indices[:r]) */
2725 index = indices[k];
2726 elem = PyTuple_GET_ITEM(pool, index);
2727 Py_INCREF(elem);
2728 oldelem = PyTuple_GET_ITEM(result, k);
2729 PyTuple_SET_ITEM(result, k, elem);
2730 Py_DECREF(oldelem);
2731 }
2732 break;
2733 }
2734 }
2735 /* If i is negative, then the cycles have all
2736 rolled-over and we're done. */
2737 if (i < 0)
2738 goto empty;
2739 }
2740 Py_INCREF(result);
2741 return result;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002742
2743empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002744 po->stopped = 1;
2745 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002746}
2747
2748PyDoc_STRVAR(permutations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +00002749"permutations(iterable[, r]) --> permutations object\n\
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002750\n\
2751Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl8550d692008-06-10 12:46:39 +00002752permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002753
2754static PyTypeObject permutations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002755 PyVarObject_HEAD_INIT(NULL, 0)
2756 "itertools.permutations", /* tp_name */
2757 sizeof(permutationsobject), /* tp_basicsize */
2758 0, /* tp_itemsize */
2759 /* methods */
2760 (destructor)permutations_dealloc, /* tp_dealloc */
2761 0, /* tp_print */
2762 0, /* tp_getattr */
2763 0, /* tp_setattr */
2764 0, /* tp_compare */
2765 0, /* tp_repr */
2766 0, /* tp_as_number */
2767 0, /* tp_as_sequence */
2768 0, /* tp_as_mapping */
2769 0, /* tp_hash */
2770 0, /* tp_call */
2771 0, /* tp_str */
2772 PyObject_GenericGetAttr, /* tp_getattro */
2773 0, /* tp_setattro */
2774 0, /* tp_as_buffer */
2775 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2776 Py_TPFLAGS_BASETYPE, /* tp_flags */
2777 permutations_doc, /* tp_doc */
2778 (traverseproc)permutations_traverse, /* tp_traverse */
2779 0, /* tp_clear */
2780 0, /* tp_richcompare */
2781 0, /* tp_weaklistoffset */
2782 PyObject_SelfIter, /* tp_iter */
2783 (iternextfunc)permutations_next, /* tp_iternext */
2784 0, /* tp_methods */
2785 0, /* tp_members */
2786 0, /* tp_getset */
2787 0, /* tp_base */
2788 0, /* tp_dict */
2789 0, /* tp_descr_get */
2790 0, /* tp_descr_set */
2791 0, /* tp_dictoffset */
2792 0, /* tp_init */
2793 0, /* tp_alloc */
2794 permutations_new, /* tp_new */
2795 PyObject_GC_Del, /* tp_free */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002796};
2797
2798
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002799/* compress object ************************************************************/
2800
2801/* Equivalent to:
2802
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002803 def compress(data, selectors):
2804 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2805 return (d for d, s in izip(data, selectors) if s)
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002806*/
2807
2808typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002809 PyObject_HEAD
2810 PyObject *data;
2811 PyObject *selectors;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002812} compressobject;
2813
2814static PyTypeObject compress_type;
2815
2816static PyObject *
2817compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2818{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002819 PyObject *seq1, *seq2;
2820 PyObject *data=NULL, *selectors=NULL;
2821 compressobject *lz;
2822 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002823
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002824 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2825 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002826
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002827 data = PyObject_GetIter(seq1);
2828 if (data == NULL)
2829 goto fail;
2830 selectors = PyObject_GetIter(seq2);
2831 if (selectors == NULL)
2832 goto fail;
2833
2834 /* create compressobject structure */
2835 lz = (compressobject *)type->tp_alloc(type, 0);
2836 if (lz == NULL)
2837 goto fail;
2838 lz->data = data;
2839 lz->selectors = selectors;
2840 return (PyObject *)lz;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002841
2842fail:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002843 Py_XDECREF(data);
2844 Py_XDECREF(selectors);
2845 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002846}
2847
2848static void
2849compress_dealloc(compressobject *lz)
2850{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002851 PyObject_GC_UnTrack(lz);
2852 Py_XDECREF(lz->data);
2853 Py_XDECREF(lz->selectors);
2854 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002855}
2856
2857static int
2858compress_traverse(compressobject *lz, visitproc visit, void *arg)
2859{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002860 Py_VISIT(lz->data);
2861 Py_VISIT(lz->selectors);
2862 return 0;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002863}
2864
2865static PyObject *
2866compress_next(compressobject *lz)
2867{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002868 PyObject *data = lz->data, *selectors = lz->selectors;
2869 PyObject *datum, *selector;
2870 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2871 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2872 int ok;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002873
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002874 while (1) {
2875 /* Steps: get datum, get selector, evaluate selector.
2876 Order is important (to match the pure python version
2877 in terms of which input gets a chance to raise an
2878 exception first).
2879 */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002880
Serhiy Storchakabb845652013-04-06 22:04:10 +03002881 datum = datanext(data);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03002882 if (datum == NULL)
Serhiy Storchakabb845652013-04-06 22:04:10 +03002883 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002884
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002885 selector = selectornext(selectors);
2886 if (selector == NULL) {
2887 Py_DECREF(datum);
2888 return NULL;
2889 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002890
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002891 ok = PyObject_IsTrue(selector);
2892 Py_DECREF(selector);
2893 if (ok == 1)
2894 return datum;
2895 Py_DECREF(datum);
2896 if (ok == -1)
2897 return NULL;
2898 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002899}
2900
2901PyDoc_STRVAR(compress_doc,
Raymond Hettinger2e2909f2009-02-19 02:15:14 +00002902"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002903\n\
2904Return data elements corresponding to true selector elements.\n\
2905Forms a shorter iterator from selected data elements using the\n\
2906selectors to choose the data elements.");
2907
2908static PyTypeObject compress_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002909 PyVarObject_HEAD_INIT(NULL, 0)
2910 "itertools.compress", /* tp_name */
2911 sizeof(compressobject), /* tp_basicsize */
2912 0, /* tp_itemsize */
2913 /* methods */
2914 (destructor)compress_dealloc, /* tp_dealloc */
2915 0, /* tp_print */
2916 0, /* tp_getattr */
2917 0, /* tp_setattr */
2918 0, /* tp_compare */
2919 0, /* tp_repr */
2920 0, /* tp_as_number */
2921 0, /* tp_as_sequence */
2922 0, /* tp_as_mapping */
2923 0, /* tp_hash */
2924 0, /* tp_call */
2925 0, /* tp_str */
2926 PyObject_GenericGetAttr, /* tp_getattro */
2927 0, /* tp_setattro */
2928 0, /* tp_as_buffer */
2929 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2930 Py_TPFLAGS_BASETYPE, /* tp_flags */
2931 compress_doc, /* tp_doc */
2932 (traverseproc)compress_traverse, /* tp_traverse */
2933 0, /* tp_clear */
2934 0, /* tp_richcompare */
2935 0, /* tp_weaklistoffset */
2936 PyObject_SelfIter, /* tp_iter */
2937 (iternextfunc)compress_next, /* tp_iternext */
2938 0, /* tp_methods */
2939 0, /* tp_members */
2940 0, /* tp_getset */
2941 0, /* tp_base */
2942 0, /* tp_dict */
2943 0, /* tp_descr_get */
2944 0, /* tp_descr_set */
2945 0, /* tp_dictoffset */
2946 0, /* tp_init */
2947 0, /* tp_alloc */
2948 compress_new, /* tp_new */
2949 PyObject_GC_Del, /* tp_free */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002950};
2951
2952
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002953/* ifilter object ************************************************************/
2954
2955typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002956 PyObject_HEAD
2957 PyObject *func;
2958 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002959} ifilterobject;
2960
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002961static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002962
2963static PyObject *
2964ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2965{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002966 PyObject *func, *seq;
2967 PyObject *it;
2968 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002969
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002970 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
2971 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002972
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002973 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
2974 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002975
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002976 /* Get iterator. */
2977 it = PyObject_GetIter(seq);
2978 if (it == NULL)
2979 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002980
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002981 /* create ifilterobject structure */
2982 lz = (ifilterobject *)type->tp_alloc(type, 0);
2983 if (lz == NULL) {
2984 Py_DECREF(it);
2985 return NULL;
2986 }
2987 Py_INCREF(func);
2988 lz->func = func;
2989 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002990
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002991 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002992}
2993
2994static void
2995ifilter_dealloc(ifilterobject *lz)
2996{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002997 PyObject_GC_UnTrack(lz);
2998 Py_XDECREF(lz->func);
2999 Py_XDECREF(lz->it);
3000 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003001}
3002
3003static int
3004ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
3005{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003006 Py_VISIT(lz->it);
3007 Py_VISIT(lz->func);
3008 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003009}
3010
3011static PyObject *
3012ifilter_next(ifilterobject *lz)
3013{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003014 PyObject *item;
3015 PyObject *it = lz->it;
3016 long ok;
3017 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003018
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003019 iternext = *Py_TYPE(it)->tp_iternext;
3020 for (;;) {
3021 item = iternext(it);
3022 if (item == NULL)
3023 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003024
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003025 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3026 ok = PyObject_IsTrue(item);
3027 } else {
3028 PyObject *good;
3029 good = PyObject_CallFunctionObjArgs(lz->func,
3030 item, NULL);
3031 if (good == NULL) {
3032 Py_DECREF(item);
3033 return NULL;
3034 }
3035 ok = PyObject_IsTrue(good);
3036 Py_DECREF(good);
3037 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003038 if (ok > 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003039 return item;
3040 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003041 if (ok < 0)
3042 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003043 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003044}
3045
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003046PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003047"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003048\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003049Return those items of sequence for which function(item) is true.\n\
3050If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003051
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003052static PyTypeObject ifilter_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003053 PyVarObject_HEAD_INIT(NULL, 0)
3054 "itertools.ifilter", /* tp_name */
3055 sizeof(ifilterobject), /* tp_basicsize */
3056 0, /* tp_itemsize */
3057 /* methods */
3058 (destructor)ifilter_dealloc, /* tp_dealloc */
3059 0, /* tp_print */
3060 0, /* tp_getattr */
3061 0, /* tp_setattr */
3062 0, /* tp_compare */
3063 0, /* tp_repr */
3064 0, /* tp_as_number */
3065 0, /* tp_as_sequence */
3066 0, /* tp_as_mapping */
3067 0, /* tp_hash */
3068 0, /* tp_call */
3069 0, /* tp_str */
3070 PyObject_GenericGetAttr, /* tp_getattro */
3071 0, /* tp_setattro */
3072 0, /* tp_as_buffer */
3073 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3074 Py_TPFLAGS_BASETYPE, /* tp_flags */
3075 ifilter_doc, /* tp_doc */
3076 (traverseproc)ifilter_traverse, /* tp_traverse */
3077 0, /* tp_clear */
3078 0, /* tp_richcompare */
3079 0, /* tp_weaklistoffset */
3080 PyObject_SelfIter, /* tp_iter */
3081 (iternextfunc)ifilter_next, /* tp_iternext */
3082 0, /* tp_methods */
3083 0, /* tp_members */
3084 0, /* tp_getset */
3085 0, /* tp_base */
3086 0, /* tp_dict */
3087 0, /* tp_descr_get */
3088 0, /* tp_descr_set */
3089 0, /* tp_dictoffset */
3090 0, /* tp_init */
3091 0, /* tp_alloc */
3092 ifilter_new, /* tp_new */
3093 PyObject_GC_Del, /* tp_free */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003094};
3095
3096
3097/* ifilterfalse object ************************************************************/
3098
3099typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003100 PyObject_HEAD
3101 PyObject *func;
3102 PyObject *it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003103} ifilterfalseobject;
3104
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003105static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003106
3107static PyObject *
3108ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3109{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003110 PyObject *func, *seq;
3111 PyObject *it;
3112 ifilterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003113
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003114 if (type == &ifilterfalse_type &&
3115 !_PyArg_NoKeywords("ifilterfalse()", kwds))
3116 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003117
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003118 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
3119 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003120
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003121 /* Get iterator. */
3122 it = PyObject_GetIter(seq);
3123 if (it == NULL)
3124 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003125
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003126 /* create ifilterfalseobject structure */
3127 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
3128 if (lz == NULL) {
3129 Py_DECREF(it);
3130 return NULL;
3131 }
3132 Py_INCREF(func);
3133 lz->func = func;
3134 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003135
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003136 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003137}
3138
3139static void
3140ifilterfalse_dealloc(ifilterfalseobject *lz)
3141{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003142 PyObject_GC_UnTrack(lz);
3143 Py_XDECREF(lz->func);
3144 Py_XDECREF(lz->it);
3145 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003146}
3147
3148static int
3149ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
3150{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003151 Py_VISIT(lz->it);
3152 Py_VISIT(lz->func);
3153 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003154}
3155
3156static PyObject *
3157ifilterfalse_next(ifilterfalseobject *lz)
3158{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003159 PyObject *item;
3160 PyObject *it = lz->it;
3161 long ok;
3162 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003163
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003164 iternext = *Py_TYPE(it)->tp_iternext;
3165 for (;;) {
3166 item = iternext(it);
3167 if (item == NULL)
3168 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003169
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003170 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3171 ok = PyObject_IsTrue(item);
3172 } else {
3173 PyObject *good;
3174 good = PyObject_CallFunctionObjArgs(lz->func,
3175 item, NULL);
3176 if (good == NULL) {
3177 Py_DECREF(item);
3178 return NULL;
3179 }
3180 ok = PyObject_IsTrue(good);
3181 Py_DECREF(good);
3182 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003183 if (ok == 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003184 return item;
3185 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003186 if (ok < 0)
3187 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003188 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003189}
3190
Raymond Hettinger60eca932003-02-09 06:40:58 +00003191PyDoc_STRVAR(ifilterfalse_doc,
3192"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
3193\n\
3194Return those items of sequence for which function(item) is false.\n\
3195If function is None, return the items that are false.");
3196
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003197static PyTypeObject ifilterfalse_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003198 PyVarObject_HEAD_INIT(NULL, 0)
3199 "itertools.ifilterfalse", /* tp_name */
3200 sizeof(ifilterfalseobject), /* tp_basicsize */
3201 0, /* tp_itemsize */
3202 /* methods */
3203 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
3204 0, /* tp_print */
3205 0, /* tp_getattr */
3206 0, /* tp_setattr */
3207 0, /* tp_compare */
3208 0, /* tp_repr */
3209 0, /* tp_as_number */
3210 0, /* tp_as_sequence */
3211 0, /* tp_as_mapping */
3212 0, /* tp_hash */
3213 0, /* tp_call */
3214 0, /* tp_str */
3215 PyObject_GenericGetAttr, /* tp_getattro */
3216 0, /* tp_setattro */
3217 0, /* tp_as_buffer */
3218 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3219 Py_TPFLAGS_BASETYPE, /* tp_flags */
3220 ifilterfalse_doc, /* tp_doc */
3221 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
3222 0, /* tp_clear */
3223 0, /* tp_richcompare */
3224 0, /* tp_weaklistoffset */
3225 PyObject_SelfIter, /* tp_iter */
3226 (iternextfunc)ifilterfalse_next, /* tp_iternext */
3227 0, /* tp_methods */
3228 0, /* tp_members */
3229 0, /* tp_getset */
3230 0, /* tp_base */
3231 0, /* tp_dict */
3232 0, /* tp_descr_get */
3233 0, /* tp_descr_set */
3234 0, /* tp_dictoffset */
3235 0, /* tp_init */
3236 0, /* tp_alloc */
3237 ifilterfalse_new, /* tp_new */
3238 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003239};
3240
3241
3242/* count object ************************************************************/
3243
3244typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003245 PyObject_HEAD
3246 Py_ssize_t cnt;
3247 PyObject *long_cnt;
3248 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003249} countobject;
3250
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003251/* Counting logic and invariants:
3252
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003253fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003254
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003255 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3256 Advances with: cnt += 1
3257 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003258
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003259slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003260
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003261 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3262 All counting is done with python objects (no overflows or underflows).
3263 Advances with: long_cnt += long_step
3264 Step may be zero -- effectively a slow version of repeat(cnt).
3265 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003266*/
3267
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003268static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003269
3270static PyObject *
3271count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3272{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003273 countobject *lz;
3274 int slow_mode = 0;
3275 Py_ssize_t cnt = 0;
3276 PyObject *long_cnt = NULL;
3277 PyObject *long_step = NULL;
3278 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003279
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003280 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3281 kwlist, &long_cnt, &long_step))
3282 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003283
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003284 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3285 (long_step != NULL && !PyNumber_Check(long_step))) {
3286 PyErr_SetString(PyExc_TypeError, "a number is required");
3287 return NULL;
3288 }
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003289
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003290 if (long_cnt != NULL) {
3291 cnt = PyInt_AsSsize_t(long_cnt);
3292 if ((cnt == -1 && PyErr_Occurred()) || !PyInt_Check(long_cnt)) {
3293 PyErr_Clear();
3294 slow_mode = 1;
3295 }
3296 Py_INCREF(long_cnt);
3297 } else {
3298 cnt = 0;
3299 long_cnt = PyInt_FromLong(0);
3300 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003301
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003302 /* If not specified, step defaults to 1 */
3303 if (long_step == NULL) {
3304 long_step = PyInt_FromLong(1);
3305 if (long_step == NULL) {
3306 Py_DECREF(long_cnt);
3307 return NULL;
3308 }
3309 } else
3310 Py_INCREF(long_step);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003311
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003312 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003313
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003314 /* Fast mode only works when the step is 1 */
3315 if (!PyInt_Check(long_step) ||
3316 PyInt_AS_LONG(long_step) != 1) {
3317 slow_mode = 1;
3318 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003319
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003320 if (slow_mode)
3321 cnt = PY_SSIZE_T_MAX;
3322 else
3323 Py_CLEAR(long_cnt);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003324
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003325 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3326 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3327 assert(slow_mode ||
3328 (PyInt_Check(long_step) && PyInt_AS_LONG(long_step) == 1));
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003329
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003330 /* create countobject structure */
3331 lz = (countobject *)type->tp_alloc(type, 0);
3332 if (lz == NULL) {
3333 Py_XDECREF(long_cnt);
3334 return NULL;
3335 }
3336 lz->cnt = cnt;
3337 lz->long_cnt = long_cnt;
3338 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003339
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003340 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003341}
3342
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003343static void
3344count_dealloc(countobject *lz)
3345{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003346 PyObject_GC_UnTrack(lz);
3347 Py_XDECREF(lz->long_cnt);
3348 Py_XDECREF(lz->long_step);
3349 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003350}
3351
Benjamin Peterson062a7c32009-02-16 21:07:52 +00003352static int
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003353count_traverse(countobject *lz, visitproc visit, void *arg)
3354{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003355 Py_VISIT(lz->long_cnt);
3356 Py_VISIT(lz->long_step);
3357 return 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003358}
3359
3360static PyObject *
3361count_nextlong(countobject *lz)
3362{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003363 PyObject *long_cnt;
3364 PyObject *stepped_up;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003365
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003366 long_cnt = lz->long_cnt;
3367 if (long_cnt == NULL) {
3368 /* Switch to slow_mode */
3369 long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
3370 if (long_cnt == NULL)
3371 return NULL;
3372 }
3373 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003374
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003375 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3376 if (stepped_up == NULL)
3377 return NULL;
3378 lz->long_cnt = stepped_up;
3379 return long_cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003380}
3381
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003382static PyObject *
3383count_next(countobject *lz)
3384{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003385 if (lz->cnt == PY_SSIZE_T_MAX)
3386 return count_nextlong(lz);
3387 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003388}
3389
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003390static PyObject *
3391count_repr(countobject *lz)
3392{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003393 PyObject *cnt_repr, *step_repr = NULL;
3394 PyObject *result = NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003395
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003396 if (lz->cnt != PY_SSIZE_T_MAX)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003397 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003398
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003399 cnt_repr = PyObject_Repr(lz->long_cnt);
3400 if (cnt_repr == NULL)
3401 return NULL;
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003402
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003403 if (PyInt_Check(lz->long_step) && PyInt_AS_LONG(lz->long_step) == 1) {
3404 /* Don't display step when it is an integer equal to 1 */
3405 result = PyString_FromFormat("count(%s)",
3406 PyString_AS_STRING(cnt_repr));
3407 } else {
3408 step_repr = PyObject_Repr(lz->long_step);
3409 if (step_repr != NULL)
3410 result = PyString_FromFormat("count(%s, %s)",
3411 PyString_AS_STRING(cnt_repr),
3412 PyString_AS_STRING(step_repr));
3413 }
3414 Py_DECREF(cnt_repr);
3415 Py_XDECREF(step_repr);
3416 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003417}
3418
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003419static PyObject *
3420count_reduce(countobject *lz)
3421{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003422 if (lz->cnt == PY_SSIZE_T_MAX)
3423 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3424 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003425}
3426
3427PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3428
3429static PyMethodDef count_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003430 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3431 count_reduce_doc},
3432 {NULL, NULL} /* sentinel */
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003433};
3434
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003435PyDoc_STRVAR(count_doc,
Ezio Melottifdc1e0d2010-06-11 02:21:25 +00003436 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003437\n\
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003438Return a count object whose .next() method returns consecutive values.\n\
3439Equivalent to:\n\n\
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003440 def count(firstval=0, step=1):\n\
Eli Benderskye6802db2013-09-03 06:41:58 -07003441 x = firstval\n\
3442 while 1:\n\
3443 yield x\n\
3444 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003445
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003446static PyTypeObject count_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003447 PyVarObject_HEAD_INIT(NULL, 0)
3448 "itertools.count", /* tp_name */
3449 sizeof(countobject), /* tp_basicsize */
3450 0, /* tp_itemsize */
3451 /* methods */
3452 (destructor)count_dealloc, /* tp_dealloc */
3453 0, /* tp_print */
3454 0, /* tp_getattr */
3455 0, /* tp_setattr */
3456 0, /* tp_compare */
3457 (reprfunc)count_repr, /* tp_repr */
3458 0, /* tp_as_number */
3459 0, /* tp_as_sequence */
3460 0, /* tp_as_mapping */
3461 0, /* tp_hash */
3462 0, /* tp_call */
3463 0, /* tp_str */
3464 PyObject_GenericGetAttr, /* tp_getattro */
3465 0, /* tp_setattro */
3466 0, /* tp_as_buffer */
3467 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3468 Py_TPFLAGS_BASETYPE, /* tp_flags */
3469 count_doc, /* tp_doc */
3470 (traverseproc)count_traverse, /* tp_traverse */
3471 0, /* tp_clear */
3472 0, /* tp_richcompare */
3473 0, /* tp_weaklistoffset */
3474 PyObject_SelfIter, /* tp_iter */
3475 (iternextfunc)count_next, /* tp_iternext */
3476 count_methods, /* tp_methods */
3477 0, /* tp_members */
3478 0, /* tp_getset */
3479 0, /* tp_base */
3480 0, /* tp_dict */
3481 0, /* tp_descr_get */
3482 0, /* tp_descr_set */
3483 0, /* tp_dictoffset */
3484 0, /* tp_init */
3485 0, /* tp_alloc */
3486 count_new, /* tp_new */
3487 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003488};
3489
3490
3491/* izip object ************************************************************/
3492
3493#include "Python.h"
3494
3495typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003496 PyObject_HEAD
3497 Py_ssize_t tuplesize;
3498 PyObject *ittuple; /* tuple of iterators */
3499 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003500} izipobject;
3501
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003502static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003503
3504static PyObject *
3505izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3506{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003507 izipobject *lz;
3508 Py_ssize_t i;
3509 PyObject *ittuple; /* tuple of iterators */
3510 PyObject *result;
3511 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003512
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003513 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
3514 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003515
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003516 /* args must be a tuple */
3517 assert(PyTuple_Check(args));
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003518
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003519 /* obtain iterators */
3520 ittuple = PyTuple_New(tuplesize);
3521 if (ittuple == NULL)
3522 return NULL;
3523 for (i=0; i < tuplesize; ++i) {
3524 PyObject *item = PyTuple_GET_ITEM(args, i);
3525 PyObject *it = PyObject_GetIter(item);
3526 if (it == NULL) {
3527 if (PyErr_ExceptionMatches(PyExc_TypeError))
3528 PyErr_Format(PyExc_TypeError,
3529 "izip argument #%zd must support iteration",
3530 i+1);
3531 Py_DECREF(ittuple);
3532 return NULL;
3533 }
3534 PyTuple_SET_ITEM(ittuple, i, it);
3535 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003536
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003537 /* create a result holder */
3538 result = PyTuple_New(tuplesize);
3539 if (result == NULL) {
3540 Py_DECREF(ittuple);
3541 return NULL;
3542 }
3543 for (i=0 ; i < tuplesize ; i++) {
3544 Py_INCREF(Py_None);
3545 PyTuple_SET_ITEM(result, i, Py_None);
3546 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003547
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003548 /* create izipobject structure */
3549 lz = (izipobject *)type->tp_alloc(type, 0);
3550 if (lz == NULL) {
3551 Py_DECREF(ittuple);
3552 Py_DECREF(result);
3553 return NULL;
3554 }
3555 lz->ittuple = ittuple;
3556 lz->tuplesize = tuplesize;
3557 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003558
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003559 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003560}
3561
3562static void
3563izip_dealloc(izipobject *lz)
3564{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003565 PyObject_GC_UnTrack(lz);
3566 Py_XDECREF(lz->ittuple);
3567 Py_XDECREF(lz->result);
3568 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003569}
3570
3571static int
3572izip_traverse(izipobject *lz, visitproc visit, void *arg)
3573{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003574 Py_VISIT(lz->ittuple);
3575 Py_VISIT(lz->result);
3576 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003577}
3578
3579static PyObject *
3580izip_next(izipobject *lz)
3581{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003582 Py_ssize_t i;
3583 Py_ssize_t tuplesize = lz->tuplesize;
3584 PyObject *result = lz->result;
3585 PyObject *it;
3586 PyObject *item;
3587 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003588
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003589 if (tuplesize == 0)
3590 return NULL;
3591 if (Py_REFCNT(result) == 1) {
3592 Py_INCREF(result);
3593 for (i=0 ; i < tuplesize ; i++) {
3594 it = PyTuple_GET_ITEM(lz->ittuple, i);
3595 item = (*Py_TYPE(it)->tp_iternext)(it);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003596 if (item == NULL) {
3597 Py_DECREF(result);
3598 return NULL;
3599 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003600 olditem = PyTuple_GET_ITEM(result, i);
3601 PyTuple_SET_ITEM(result, i, item);
3602 Py_DECREF(olditem);
3603 }
3604 } else {
3605 result = PyTuple_New(tuplesize);
3606 if (result == NULL)
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003607 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003608 for (i=0 ; i < tuplesize ; i++) {
3609 it = PyTuple_GET_ITEM(lz->ittuple, i);
3610 item = (*Py_TYPE(it)->tp_iternext)(it);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003611 if (item == NULL) {
3612 Py_DECREF(result);
3613 return NULL;
3614 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003615 PyTuple_SET_ITEM(result, i, item);
3616 }
3617 }
3618 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003619}
3620
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003621PyDoc_STRVAR(izip_doc,
3622"izip(iter1 [,iter2 [...]]) --> izip object\n\
3623\n\
3624Return a izip object whose .next() method returns a tuple where\n\
3625the i-th element comes from the i-th iterable argument. The .next()\n\
3626method continues until the shortest iterable in the argument sequence\n\
3627is exhausted and then it raises StopIteration. Works like the zip()\n\
3628function but consumes less memory by returning an iterator instead of\n\
3629a list.");
3630
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003631static PyTypeObject izip_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003632 PyVarObject_HEAD_INIT(NULL, 0)
3633 "itertools.izip", /* tp_name */
3634 sizeof(izipobject), /* tp_basicsize */
3635 0, /* tp_itemsize */
3636 /* methods */
3637 (destructor)izip_dealloc, /* tp_dealloc */
3638 0, /* tp_print */
3639 0, /* tp_getattr */
3640 0, /* tp_setattr */
3641 0, /* tp_compare */
3642 0, /* tp_repr */
3643 0, /* tp_as_number */
3644 0, /* tp_as_sequence */
3645 0, /* tp_as_mapping */
3646 0, /* tp_hash */
3647 0, /* tp_call */
3648 0, /* tp_str */
3649 PyObject_GenericGetAttr, /* tp_getattro */
3650 0, /* tp_setattro */
3651 0, /* tp_as_buffer */
3652 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3653 Py_TPFLAGS_BASETYPE, /* tp_flags */
3654 izip_doc, /* tp_doc */
3655 (traverseproc)izip_traverse, /* tp_traverse */
3656 0, /* tp_clear */
3657 0, /* tp_richcompare */
3658 0, /* tp_weaklistoffset */
3659 PyObject_SelfIter, /* tp_iter */
3660 (iternextfunc)izip_next, /* tp_iternext */
3661 0, /* tp_methods */
3662 0, /* tp_members */
3663 0, /* tp_getset */
3664 0, /* tp_base */
3665 0, /* tp_dict */
3666 0, /* tp_descr_get */
3667 0, /* tp_descr_set */
3668 0, /* tp_dictoffset */
3669 0, /* tp_init */
3670 0, /* tp_alloc */
3671 izip_new, /* tp_new */
3672 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003673};
3674
3675
3676/* repeat object ************************************************************/
3677
3678typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003679 PyObject_HEAD
3680 PyObject *element;
3681 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003682} repeatobject;
3683
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003684static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003685
3686static PyObject *
3687repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3688{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003689 repeatobject *ro;
3690 PyObject *element;
Raymond Hettinger58ad2452014-06-24 21:53:45 -07003691 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003692 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003693
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003694 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3695 &element, &cnt))
3696 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003697
Raymond Hettinger58ad2452014-06-24 21:53:45 -07003698 if (kwds != NULL)
3699 n_kwds = PyDict_Size(kwds);
3700 /* Does user supply times argument? */
3701 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003702 cnt = 0;
3703
3704 ro = (repeatobject *)type->tp_alloc(type, 0);
3705 if (ro == NULL)
3706 return NULL;
3707 Py_INCREF(element);
3708 ro->element = element;
3709 ro->cnt = cnt;
3710 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003711}
3712
3713static void
3714repeat_dealloc(repeatobject *ro)
3715{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003716 PyObject_GC_UnTrack(ro);
3717 Py_XDECREF(ro->element);
3718 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003719}
3720
3721static int
3722repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3723{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003724 Py_VISIT(ro->element);
3725 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003726}
3727
3728static PyObject *
3729repeat_next(repeatobject *ro)
3730{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003731 if (ro->cnt == 0)
3732 return NULL;
3733 if (ro->cnt > 0)
3734 ro->cnt--;
3735 Py_INCREF(ro->element);
3736 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003737}
3738
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003739static PyObject *
3740repeat_repr(repeatobject *ro)
3741{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003742 PyObject *result, *objrepr;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003743
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003744 objrepr = PyObject_Repr(ro->element);
3745 if (objrepr == NULL)
3746 return NULL;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003747
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003748 if (ro->cnt == -1)
3749 result = PyString_FromFormat("repeat(%s)",
3750 PyString_AS_STRING(objrepr));
3751 else
3752 result = PyString_FromFormat("repeat(%s, %zd)",
3753 PyString_AS_STRING(objrepr), ro->cnt);
3754 Py_DECREF(objrepr);
3755 return result;
3756}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003757
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003758static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003759repeat_len(repeatobject *ro)
3760{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003761 if (ro->cnt == -1) {
3762 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3763 return NULL;
3764 }
3765 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003766}
3767
Armin Rigof5b3e362006-02-11 21:32:43 +00003768PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003769
3770static PyMethodDef repeat_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003771 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3772 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003773};
3774
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003775PyDoc_STRVAR(repeat_doc,
Raymond Hettinger182edae2009-02-19 02:38:25 +00003776"repeat(object [,times]) -> create an iterator which returns the object\n\
3777for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003778endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003779
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003780static PyTypeObject repeat_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003781 PyVarObject_HEAD_INIT(NULL, 0)
3782 "itertools.repeat", /* tp_name */
3783 sizeof(repeatobject), /* tp_basicsize */
3784 0, /* tp_itemsize */
3785 /* methods */
3786 (destructor)repeat_dealloc, /* tp_dealloc */
3787 0, /* tp_print */
3788 0, /* tp_getattr */
3789 0, /* tp_setattr */
3790 0, /* tp_compare */
3791 (reprfunc)repeat_repr, /* tp_repr */
3792 0, /* tp_as_number */
3793 0, /* tp_as_sequence */
3794 0, /* tp_as_mapping */
3795 0, /* tp_hash */
3796 0, /* tp_call */
3797 0, /* tp_str */
3798 PyObject_GenericGetAttr, /* tp_getattro */
3799 0, /* tp_setattro */
3800 0, /* tp_as_buffer */
3801 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3802 Py_TPFLAGS_BASETYPE, /* tp_flags */
3803 repeat_doc, /* tp_doc */
3804 (traverseproc)repeat_traverse, /* tp_traverse */
3805 0, /* tp_clear */
3806 0, /* tp_richcompare */
3807 0, /* tp_weaklistoffset */
3808 PyObject_SelfIter, /* tp_iter */
3809 (iternextfunc)repeat_next, /* tp_iternext */
3810 repeat_methods, /* tp_methods */
3811 0, /* tp_members */
3812 0, /* tp_getset */
3813 0, /* tp_base */
3814 0, /* tp_dict */
3815 0, /* tp_descr_get */
3816 0, /* tp_descr_set */
3817 0, /* tp_dictoffset */
3818 0, /* tp_init */
3819 0, /* tp_alloc */
3820 repeat_new, /* tp_new */
3821 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003822};
3823
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003824/* iziplongest object ************************************************************/
3825
3826#include "Python.h"
3827
3828typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003829 PyObject_HEAD
3830 Py_ssize_t tuplesize;
3831 Py_ssize_t numactive;
3832 PyObject *ittuple; /* tuple of iterators */
3833 PyObject *result;
3834 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003835} iziplongestobject;
3836
3837static PyTypeObject iziplongest_type;
3838
3839static PyObject *
3840izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3841{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003842 iziplongestobject *lz;
3843 Py_ssize_t i;
3844 PyObject *ittuple; /* tuple of iterators */
3845 PyObject *result;
3846 PyObject *fillvalue = Py_None;
3847 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003848
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003849 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3850 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3851 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3852 PyErr_SetString(PyExc_TypeError,
3853 "izip_longest() got an unexpected keyword argument");
3854 return NULL;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003855 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003856 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003857
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003858 /* args must be a tuple */
3859 assert(PyTuple_Check(args));
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003860
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003861 /* obtain iterators */
3862 ittuple = PyTuple_New(tuplesize);
3863 if (ittuple == NULL)
3864 return NULL;
3865 for (i=0; i < tuplesize; ++i) {
3866 PyObject *item = PyTuple_GET_ITEM(args, i);
3867 PyObject *it = PyObject_GetIter(item);
3868 if (it == NULL) {
3869 if (PyErr_ExceptionMatches(PyExc_TypeError))
3870 PyErr_Format(PyExc_TypeError,
3871 "izip_longest argument #%zd must support iteration",
3872 i+1);
3873 Py_DECREF(ittuple);
3874 return NULL;
3875 }
3876 PyTuple_SET_ITEM(ittuple, i, it);
3877 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003878
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003879 /* create a result holder */
3880 result = PyTuple_New(tuplesize);
3881 if (result == NULL) {
3882 Py_DECREF(ittuple);
3883 return NULL;
3884 }
3885 for (i=0 ; i < tuplesize ; i++) {
3886 Py_INCREF(Py_None);
3887 PyTuple_SET_ITEM(result, i, Py_None);
3888 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003889
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003890 /* create iziplongestobject structure */
3891 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3892 if (lz == NULL) {
3893 Py_DECREF(ittuple);
3894 Py_DECREF(result);
3895 return NULL;
3896 }
3897 lz->ittuple = ittuple;
3898 lz->tuplesize = tuplesize;
3899 lz->numactive = tuplesize;
3900 lz->result = result;
3901 Py_INCREF(fillvalue);
3902 lz->fillvalue = fillvalue;
3903 return (PyObject *)lz;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003904}
3905
3906static void
3907izip_longest_dealloc(iziplongestobject *lz)
3908{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003909 PyObject_GC_UnTrack(lz);
3910 Py_XDECREF(lz->ittuple);
3911 Py_XDECREF(lz->result);
3912 Py_XDECREF(lz->fillvalue);
3913 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003914}
3915
3916static int
3917izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3918{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003919 Py_VISIT(lz->ittuple);
3920 Py_VISIT(lz->result);
3921 Py_VISIT(lz->fillvalue);
3922 return 0;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003923}
3924
3925static PyObject *
3926izip_longest_next(iziplongestobject *lz)
3927{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003928 Py_ssize_t i;
3929 Py_ssize_t tuplesize = lz->tuplesize;
3930 PyObject *result = lz->result;
3931 PyObject *it;
3932 PyObject *item;
3933 PyObject *olditem;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003934
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003935 if (tuplesize == 0)
3936 return NULL;
3937 if (lz->numactive == 0)
3938 return NULL;
3939 if (Py_REFCNT(result) == 1) {
3940 Py_INCREF(result);
3941 for (i=0 ; i < tuplesize ; i++) {
3942 it = PyTuple_GET_ITEM(lz->ittuple, i);
3943 if (it == NULL) {
3944 Py_INCREF(lz->fillvalue);
3945 item = lz->fillvalue;
3946 } else {
3947 item = PyIter_Next(it);
3948 if (item == NULL) {
3949 lz->numactive -= 1;
3950 if (lz->numactive == 0 || PyErr_Occurred()) {
3951 lz->numactive = 0;
3952 Py_DECREF(result);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003953 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003954 } else {
3955 Py_INCREF(lz->fillvalue);
3956 item = lz->fillvalue;
3957 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3958 Py_DECREF(it);
3959 }
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003960 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003961 }
3962 olditem = PyTuple_GET_ITEM(result, i);
3963 PyTuple_SET_ITEM(result, i, item);
3964 Py_DECREF(olditem);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003965 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003966 } else {
3967 result = PyTuple_New(tuplesize);
3968 if (result == NULL)
3969 return NULL;
3970 for (i=0 ; i < tuplesize ; i++) {
3971 it = PyTuple_GET_ITEM(lz->ittuple, i);
3972 if (it == NULL) {
3973 Py_INCREF(lz->fillvalue);
3974 item = lz->fillvalue;
3975 } else {
3976 item = PyIter_Next(it);
3977 if (item == NULL) {
3978 lz->numactive -= 1;
3979 if (lz->numactive == 0 || PyErr_Occurred()) {
3980 lz->numactive = 0;
3981 Py_DECREF(result);
3982 return NULL;
3983 } else {
3984 Py_INCREF(lz->fillvalue);
3985 item = lz->fillvalue;
3986 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3987 Py_DECREF(it);
3988 }
3989 }
3990 }
3991 PyTuple_SET_ITEM(result, i, item);
3992 }
3993 }
3994 return result;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003995}
3996
3997PyDoc_STRVAR(izip_longest_doc,
3998"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3999\n\
4000Return an izip_longest object whose .next() method returns a tuple where\n\
4001the i-th element comes from the i-th iterable argument. The .next()\n\
4002method continues until the longest iterable in the argument sequence\n\
4003is exhausted and then it raises StopIteration. When the shorter iterables\n\
4004are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4005defaults to None or can be specified by a keyword argument.\n\
4006");
4007
4008static PyTypeObject iziplongest_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004009 PyVarObject_HEAD_INIT(NULL, 0)
4010 "itertools.izip_longest", /* tp_name */
4011 sizeof(iziplongestobject), /* tp_basicsize */
4012 0, /* tp_itemsize */
4013 /* methods */
4014 (destructor)izip_longest_dealloc, /* tp_dealloc */
4015 0, /* tp_print */
4016 0, /* tp_getattr */
4017 0, /* tp_setattr */
4018 0, /* tp_compare */
4019 0, /* tp_repr */
4020 0, /* tp_as_number */
4021 0, /* tp_as_sequence */
4022 0, /* tp_as_mapping */
4023 0, /* tp_hash */
4024 0, /* tp_call */
4025 0, /* tp_str */
4026 PyObject_GenericGetAttr, /* tp_getattro */
4027 0, /* tp_setattro */
4028 0, /* tp_as_buffer */
4029 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4030 Py_TPFLAGS_BASETYPE, /* tp_flags */
4031 izip_longest_doc, /* tp_doc */
4032 (traverseproc)izip_longest_traverse, /* tp_traverse */
4033 0, /* tp_clear */
4034 0, /* tp_richcompare */
4035 0, /* tp_weaklistoffset */
4036 PyObject_SelfIter, /* tp_iter */
4037 (iternextfunc)izip_longest_next, /* tp_iternext */
4038 0, /* tp_methods */
4039 0, /* tp_members */
4040 0, /* tp_getset */
4041 0, /* tp_base */
4042 0, /* tp_dict */
4043 0, /* tp_descr_get */
4044 0, /* tp_descr_set */
4045 0, /* tp_dictoffset */
4046 0, /* tp_init */
4047 0, /* tp_alloc */
4048 izip_longest_new, /* tp_new */
4049 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd36862c2007-02-21 05:20:38 +00004050};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004051
4052/* module level code ********************************************************/
4053
4054PyDoc_STRVAR(module_doc,
4055"Functional tools for creating and using iterators.\n\
4056\n\
4057Infinite iterators:\n\
4058count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004059cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004060repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004061\n\
4062Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004063chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettingere76816b2009-01-29 03:43:44 +00004064compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004065dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4066groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00004067ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
4068ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004069islice(seq, [start,] stop [, step]) --> elements from\n\
4070 seq[start:stop:step]\n\
4071imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
4072starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004073tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004074takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004075izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4076izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4077\n\
4078Combinatoric generators:\n\
4079product(p, q, ... [repeat=1]) --> cartesian product\n\
4080permutations(p[, r])\n\
Raymond Hettinger9eac1192009-11-19 01:22:04 +00004081combinations(p, r)\n\
4082combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004083");
4084
4085
Raymond Hettingerad983e72003-11-12 14:32:26 +00004086static PyMethodDef module_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004087 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4088 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004089};
4090
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004091PyMODINIT_FUNC
4092inititertools(void)
4093{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004094 int i;
4095 PyObject *m;
4096 char *name;
4097 PyTypeObject *typelist[] = {
4098 &combinations_type,
4099 &cwr_type,
4100 &cycle_type,
4101 &dropwhile_type,
4102 &takewhile_type,
4103 &islice_type,
4104 &starmap_type,
4105 &imap_type,
4106 &chain_type,
4107 &compress_type,
4108 &ifilter_type,
4109 &ifilterfalse_type,
4110 &count_type,
4111 &izip_type,
4112 &iziplongest_type,
4113 &permutations_type,
4114 &product_type,
4115 &repeat_type,
4116 &groupby_type,
4117 NULL
4118 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004119
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004120 Py_TYPE(&teedataobject_type) = &PyType_Type;
4121 m = Py_InitModule3("itertools", module_methods, module_doc);
4122 if (m == NULL)
4123 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004124
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004125 for (i=0 ; typelist[i] != NULL ; i++) {
4126 if (PyType_Ready(typelist[i]) < 0)
4127 return;
4128 name = strchr(typelist[i]->tp_name, '.');
4129 assert (name != NULL);
4130 Py_INCREF(typelist[i]);
4131 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4132 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004133
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004134 if (PyType_Ready(&teedataobject_type) < 0)
4135 return;
4136 if (PyType_Ready(&tee_type) < 0)
4137 return;
4138 if (PyType_Ready(&_grouper_type) < 0)
4139 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004140}