blob: cf3aadfd3e07128e2afecc9cdbe1b10e9752f2d0 [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 Storchaka763a61c2016-04-10 18:05:12 +0300495 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
T. Woutersd694a062017-03-30 12:49:22 -07001711 /* lz->source is the iterator of iterables. If it's NULL, we've already
1712 * consumed them all. lz->active is the current iterator. If it's NULL,
1713 * we should grab a new one from lz->source. */
1714 while (lz->source != NULL) {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001715 if (lz->active == NULL) {
T. Woutersd694a062017-03-30 12:49:22 -07001716 PyObject *iterable = PyIter_Next(lz->source);
1717 if (iterable == NULL) {
1718 Py_CLEAR(lz->source);
1719 return NULL; /* no more input sources */
1720 }
1721 lz->active = PyObject_GetIter(iterable);
1722 Py_DECREF(iterable);
1723 if (lz->active == NULL) {
1724 Py_CLEAR(lz->source);
1725 return NULL; /* input not iterable */
1726 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001727 }
T. Woutersd694a062017-03-30 12:49:22 -07001728 item = PyIter_Next(lz->active);
1729 if (item != NULL)
1730 return item;
1731 if (PyErr_Occurred()) {
1732 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1733 PyErr_Clear();
1734 else
1735 return NULL; /* input raised an exception */
1736 }
1737 /* lz->active is consumed, try with the next iterable. */
1738 Py_CLEAR(lz->active);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001739 }
T. Woutersd694a062017-03-30 12:49:22 -07001740 /* Everything had been consumed already. */
1741 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001742}
1743
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001744PyDoc_STRVAR(chain_doc,
1745"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001746\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001747Return a chain object whose .next() method returns elements from the\n\
1748first iterable until it is exhausted, then elements from the next\n\
1749iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001750
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001751PyDoc_STRVAR(chain_from_iterable_doc,
1752"chain.from_iterable(iterable) --> chain object\n\
1753\n\
Martin Pantera850ef62016-07-28 01:11:04 +00001754Alternate chain() constructor taking a single iterable argument\n\
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001755that evaluates lazily.");
1756
1757static PyMethodDef chain_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001758 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1759 chain_from_iterable_doc},
1760 {NULL, NULL} /* sentinel */
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001761};
1762
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001763static PyTypeObject chain_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001764 PyVarObject_HEAD_INIT(NULL, 0)
1765 "itertools.chain", /* tp_name */
1766 sizeof(chainobject), /* tp_basicsize */
1767 0, /* tp_itemsize */
1768 /* methods */
1769 (destructor)chain_dealloc, /* tp_dealloc */
1770 0, /* tp_print */
1771 0, /* tp_getattr */
1772 0, /* tp_setattr */
1773 0, /* tp_compare */
1774 0, /* tp_repr */
1775 0, /* tp_as_number */
1776 0, /* tp_as_sequence */
1777 0, /* tp_as_mapping */
1778 0, /* tp_hash */
1779 0, /* tp_call */
1780 0, /* tp_str */
1781 PyObject_GenericGetAttr, /* tp_getattro */
1782 0, /* tp_setattro */
1783 0, /* tp_as_buffer */
1784 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1785 Py_TPFLAGS_BASETYPE, /* tp_flags */
1786 chain_doc, /* tp_doc */
1787 (traverseproc)chain_traverse, /* tp_traverse */
1788 0, /* tp_clear */
1789 0, /* tp_richcompare */
1790 0, /* tp_weaklistoffset */
1791 PyObject_SelfIter, /* tp_iter */
1792 (iternextfunc)chain_next, /* tp_iternext */
1793 chain_methods, /* tp_methods */
1794 0, /* tp_members */
1795 0, /* tp_getset */
1796 0, /* tp_base */
1797 0, /* tp_dict */
1798 0, /* tp_descr_get */
1799 0, /* tp_descr_set */
1800 0, /* tp_dictoffset */
1801 0, /* tp_init */
1802 0, /* tp_alloc */
1803 chain_new, /* tp_new */
1804 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001805};
1806
1807
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001808/* product object ************************************************************/
1809
1810typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001811 PyObject_HEAD
1812 PyObject *pools; /* tuple of pool tuples */
1813 Py_ssize_t *indices; /* one index per pool */
1814 PyObject *result; /* most recently returned result tuple */
1815 int stopped; /* set to 1 when the product iterator is exhausted */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001816} productobject;
1817
1818static PyTypeObject product_type;
1819
1820static PyObject *
1821product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1822{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001823 productobject *lz;
1824 Py_ssize_t nargs, npools, repeat=1;
1825 PyObject *pools = NULL;
1826 Py_ssize_t *indices = NULL;
1827 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001828
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001829 if (kwds != NULL) {
1830 char *kwlist[] = {"repeat", 0};
1831 PyObject *tmpargs = PyTuple_New(0);
1832 if (tmpargs == NULL)
1833 return NULL;
1834 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1835 Py_DECREF(tmpargs);
1836 return NULL;
1837 }
1838 Py_DECREF(tmpargs);
1839 if (repeat < 0) {
1840 PyErr_SetString(PyExc_ValueError,
1841 "repeat argument cannot be negative");
1842 return NULL;
1843 }
1844 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001845
Benjamin Petersondda91212015-02-01 21:34:07 -05001846 assert(PyTuple_CheckExact(args));
1847 if (repeat == 0) {
1848 nargs = 0;
1849 } else {
1850 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02001851 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Petersondda91212015-02-01 21:34:07 -05001852 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
1853 return NULL;
1854 }
1855 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001856 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001857
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02001858 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001859 if (indices == NULL) {
1860 PyErr_NoMemory();
1861 goto error;
1862 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001863
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001864 pools = PyTuple_New(npools);
1865 if (pools == NULL)
1866 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001867
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001868 for (i=0; i < nargs ; ++i) {
1869 PyObject *item = PyTuple_GET_ITEM(args, i);
1870 PyObject *pool = PySequence_Tuple(item);
1871 if (pool == NULL)
1872 goto error;
1873 PyTuple_SET_ITEM(pools, i, pool);
1874 indices[i] = 0;
1875 }
1876 for ( ; i < npools; ++i) {
1877 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1878 Py_INCREF(pool);
1879 PyTuple_SET_ITEM(pools, i, pool);
1880 indices[i] = 0;
1881 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001882
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001883 /* create productobject structure */
1884 lz = (productobject *)type->tp_alloc(type, 0);
1885 if (lz == NULL)
1886 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001887
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001888 lz->pools = pools;
1889 lz->indices = indices;
1890 lz->result = NULL;
1891 lz->stopped = 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001892
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001893 return (PyObject *)lz;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001894
1895error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001896 if (indices != NULL)
1897 PyMem_Free(indices);
1898 Py_XDECREF(pools);
1899 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001900}
1901
1902static void
1903product_dealloc(productobject *lz)
1904{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001905 PyObject_GC_UnTrack(lz);
1906 Py_XDECREF(lz->pools);
1907 Py_XDECREF(lz->result);
1908 if (lz->indices != NULL)
1909 PyMem_Free(lz->indices);
1910 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001911}
1912
1913static int
1914product_traverse(productobject *lz, visitproc visit, void *arg)
1915{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001916 Py_VISIT(lz->pools);
1917 Py_VISIT(lz->result);
1918 return 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001919}
1920
1921static PyObject *
1922product_next(productobject *lz)
1923{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001924 PyObject *pool;
1925 PyObject *elem;
1926 PyObject *oldelem;
1927 PyObject *pools = lz->pools;
1928 PyObject *result = lz->result;
1929 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1930 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001931
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001932 if (lz->stopped)
1933 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001934
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001935 if (result == NULL) {
1936 /* On the first pass, return an initial tuple filled with the
1937 first element from each pool. */
1938 result = PyTuple_New(npools);
1939 if (result == NULL)
1940 goto empty;
1941 lz->result = result;
1942 for (i=0; i < npools; i++) {
1943 pool = PyTuple_GET_ITEM(pools, i);
1944 if (PyTuple_GET_SIZE(pool) == 0)
1945 goto empty;
1946 elem = PyTuple_GET_ITEM(pool, 0);
1947 Py_INCREF(elem);
1948 PyTuple_SET_ITEM(result, i, elem);
1949 }
1950 } else {
1951 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001952
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001953 /* Copy the previous result tuple or re-use it if available */
1954 if (Py_REFCNT(result) > 1) {
1955 PyObject *old_result = result;
1956 result = PyTuple_New(npools);
1957 if (result == NULL)
1958 goto empty;
1959 lz->result = result;
1960 for (i=0; i < npools; i++) {
1961 elem = PyTuple_GET_ITEM(old_result, i);
1962 Py_INCREF(elem);
1963 PyTuple_SET_ITEM(result, i, elem);
1964 }
1965 Py_DECREF(old_result);
1966 }
1967 /* Now, we've got the only copy so we can update it in-place */
1968 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001969
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001970 /* Update the pool indices right-to-left. Only advance to the
1971 next pool when the previous one rolls-over */
1972 for (i=npools-1 ; i >= 0 ; i--) {
1973 pool = PyTuple_GET_ITEM(pools, i);
1974 indices[i]++;
1975 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1976 /* Roll-over and advance to next pool */
1977 indices[i] = 0;
1978 elem = PyTuple_GET_ITEM(pool, 0);
1979 Py_INCREF(elem);
1980 oldelem = PyTuple_GET_ITEM(result, i);
1981 PyTuple_SET_ITEM(result, i, elem);
1982 Py_DECREF(oldelem);
1983 } else {
1984 /* No rollover. Just increment and stop here. */
1985 elem = PyTuple_GET_ITEM(pool, indices[i]);
1986 Py_INCREF(elem);
1987 oldelem = PyTuple_GET_ITEM(result, i);
1988 PyTuple_SET_ITEM(result, i, elem);
1989 Py_DECREF(oldelem);
1990 break;
1991 }
1992 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00001993
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001994 /* If i is negative, then the indices have all rolled-over
1995 and we're done. */
1996 if (i < 0)
1997 goto empty;
1998 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001999
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002000 Py_INCREF(result);
2001 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002002
2003empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002004 lz->stopped = 1;
2005 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002006}
2007
2008PyDoc_STRVAR(product_doc,
2009"product(*iterables) --> product object\n\
2010\n\
Raymond Hettinger73d79632008-02-23 02:20:41 +00002011Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002012For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2013The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2014cycle in a manner similar to an odometer (with the rightmost element changing\n\
2015on every iteration).\n\n\
Raymond Hettinger94034ea2009-02-09 18:39:41 +00002016To compute the product of an iterable with itself, specify the number\n\
2017of repetitions with the optional repeat keyword argument. For example,\n\
2018product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002019product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2020product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2021
2022static PyTypeObject product_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002023 PyVarObject_HEAD_INIT(NULL, 0)
2024 "itertools.product", /* tp_name */
2025 sizeof(productobject), /* tp_basicsize */
2026 0, /* tp_itemsize */
2027 /* methods */
2028 (destructor)product_dealloc, /* tp_dealloc */
2029 0, /* tp_print */
2030 0, /* tp_getattr */
2031 0, /* tp_setattr */
2032 0, /* tp_compare */
2033 0, /* tp_repr */
2034 0, /* tp_as_number */
2035 0, /* tp_as_sequence */
2036 0, /* tp_as_mapping */
2037 0, /* tp_hash */
2038 0, /* tp_call */
2039 0, /* tp_str */
2040 PyObject_GenericGetAttr, /* tp_getattro */
2041 0, /* tp_setattro */
2042 0, /* tp_as_buffer */
2043 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2044 Py_TPFLAGS_BASETYPE, /* tp_flags */
2045 product_doc, /* tp_doc */
2046 (traverseproc)product_traverse, /* tp_traverse */
2047 0, /* tp_clear */
2048 0, /* tp_richcompare */
2049 0, /* tp_weaklistoffset */
2050 PyObject_SelfIter, /* tp_iter */
2051 (iternextfunc)product_next, /* tp_iternext */
2052 0, /* tp_methods */
2053 0, /* tp_members */
2054 0, /* tp_getset */
2055 0, /* tp_base */
2056 0, /* tp_dict */
2057 0, /* tp_descr_get */
2058 0, /* tp_descr_set */
2059 0, /* tp_dictoffset */
2060 0, /* tp_init */
2061 0, /* tp_alloc */
2062 product_new, /* tp_new */
2063 PyObject_GC_Del, /* tp_free */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002064};
2065
2066
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002067/* combinations object ************************************************************/
2068
2069typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002070 PyObject_HEAD
2071 PyObject *pool; /* input converted to a tuple */
2072 Py_ssize_t *indices; /* one index per result element */
2073 PyObject *result; /* most recently returned result tuple */
2074 Py_ssize_t r; /* size of result tuple */
2075 int stopped; /* set to 1 when the combinations iterator is exhausted */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002076} combinationsobject;
2077
2078static PyTypeObject combinations_type;
2079
2080static PyObject *
2081combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2082{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002083 combinationsobject *co;
2084 Py_ssize_t n;
2085 Py_ssize_t r;
2086 PyObject *pool = NULL;
2087 PyObject *iterable = NULL;
2088 Py_ssize_t *indices = NULL;
2089 Py_ssize_t i;
2090 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002091
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002092 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2093 &iterable, &r))
2094 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002095
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002096 pool = PySequence_Tuple(iterable);
2097 if (pool == NULL)
2098 goto error;
2099 n = PyTuple_GET_SIZE(pool);
2100 if (r < 0) {
2101 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2102 goto error;
2103 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002104
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02002105 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002106 if (indices == NULL) {
2107 PyErr_NoMemory();
2108 goto error;
2109 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002110
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002111 for (i=0 ; i<r ; i++)
2112 indices[i] = i;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002113
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002114 /* create combinationsobject structure */
2115 co = (combinationsobject *)type->tp_alloc(type, 0);
2116 if (co == NULL)
2117 goto error;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002118
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002119 co->pool = pool;
2120 co->indices = indices;
2121 co->result = NULL;
2122 co->r = r;
2123 co->stopped = r > n ? 1 : 0;
2124
2125 return (PyObject *)co;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002126
2127error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002128 if (indices != NULL)
2129 PyMem_Free(indices);
2130 Py_XDECREF(pool);
2131 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002132}
2133
2134static void
2135combinations_dealloc(combinationsobject *co)
2136{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002137 PyObject_GC_UnTrack(co);
2138 Py_XDECREF(co->pool);
2139 Py_XDECREF(co->result);
2140 if (co->indices != NULL)
2141 PyMem_Free(co->indices);
2142 Py_TYPE(co)->tp_free(co);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002143}
2144
2145static int
2146combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2147{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002148 Py_VISIT(co->pool);
2149 Py_VISIT(co->result);
2150 return 0;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002151}
2152
2153static PyObject *
2154combinations_next(combinationsobject *co)
2155{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002156 PyObject *elem;
2157 PyObject *oldelem;
2158 PyObject *pool = co->pool;
2159 Py_ssize_t *indices = co->indices;
2160 PyObject *result = co->result;
2161 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2162 Py_ssize_t r = co->r;
2163 Py_ssize_t i, j, index;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002164
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002165 if (co->stopped)
2166 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002167
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002168 if (result == NULL) {
2169 /* On the first pass, initialize result tuple using the indices */
2170 result = PyTuple_New(r);
2171 if (result == NULL)
2172 goto empty;
2173 co->result = result;
2174 for (i=0; i<r ; i++) {
2175 index = indices[i];
2176 elem = PyTuple_GET_ITEM(pool, index);
2177 Py_INCREF(elem);
2178 PyTuple_SET_ITEM(result, i, elem);
2179 }
2180 } else {
2181 /* Copy the previous result tuple or re-use it if available */
2182 if (Py_REFCNT(result) > 1) {
2183 PyObject *old_result = result;
2184 result = PyTuple_New(r);
2185 if (result == NULL)
2186 goto empty;
2187 co->result = result;
2188 for (i=0; i<r ; i++) {
2189 elem = PyTuple_GET_ITEM(old_result, i);
2190 Py_INCREF(elem);
2191 PyTuple_SET_ITEM(result, i, elem);
2192 }
2193 Py_DECREF(old_result);
2194 }
2195 /* Now, we've got the only copy so we can update it in-place
2196 * CPython's empty tuple is a singleton and cached in
2197 * PyTuple's freelist.
2198 */
2199 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002200
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002201 /* Scan indices right-to-left until finding one that is not
2202 at its maximum (i + n - r). */
2203 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2204 ;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002205
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002206 /* If i is negative, then the indices are all at
2207 their maximum value and we're done. */
2208 if (i < 0)
2209 goto empty;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002210
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002211 /* Increment the current index which we know is not at its
2212 maximum. Then move back to the right setting each index
2213 to its lowest possible value (one higher than the index
2214 to its left -- this maintains the sort order invariant). */
2215 indices[i]++;
2216 for (j=i+1 ; j<r ; j++)
2217 indices[j] = indices[j-1] + 1;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002218
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002219 /* Update the result tuple for the new indices
2220 starting with i, the leftmost index that changed */
2221 for ( ; i<r ; i++) {
2222 index = indices[i];
2223 elem = PyTuple_GET_ITEM(pool, index);
2224 Py_INCREF(elem);
2225 oldelem = PyTuple_GET_ITEM(result, i);
2226 PyTuple_SET_ITEM(result, i, elem);
2227 Py_DECREF(oldelem);
2228 }
2229 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002230
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002231 Py_INCREF(result);
2232 return result;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002233
2234empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002235 co->stopped = 1;
2236 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002237}
2238
2239PyDoc_STRVAR(combinations_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002240"combinations(iterable, r) --> combinations object\n\
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002241\n\
2242Return successive r-length combinations of elements in the iterable.\n\n\
2243combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2244
2245static PyTypeObject combinations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002246 PyVarObject_HEAD_INIT(NULL, 0)
2247 "itertools.combinations", /* tp_name */
2248 sizeof(combinationsobject), /* tp_basicsize */
2249 0, /* tp_itemsize */
2250 /* methods */
2251 (destructor)combinations_dealloc, /* tp_dealloc */
2252 0, /* tp_print */
2253 0, /* tp_getattr */
2254 0, /* tp_setattr */
2255 0, /* tp_compare */
2256 0, /* tp_repr */
2257 0, /* tp_as_number */
2258 0, /* tp_as_sequence */
2259 0, /* tp_as_mapping */
2260 0, /* tp_hash */
2261 0, /* tp_call */
2262 0, /* tp_str */
2263 PyObject_GenericGetAttr, /* tp_getattro */
2264 0, /* tp_setattro */
2265 0, /* tp_as_buffer */
2266 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2267 Py_TPFLAGS_BASETYPE, /* tp_flags */
2268 combinations_doc, /* tp_doc */
2269 (traverseproc)combinations_traverse, /* tp_traverse */
2270 0, /* tp_clear */
2271 0, /* tp_richcompare */
2272 0, /* tp_weaklistoffset */
2273 PyObject_SelfIter, /* tp_iter */
2274 (iternextfunc)combinations_next, /* tp_iternext */
2275 0, /* tp_methods */
2276 0, /* tp_members */
2277 0, /* tp_getset */
2278 0, /* tp_base */
2279 0, /* tp_dict */
2280 0, /* tp_descr_get */
2281 0, /* tp_descr_set */
2282 0, /* tp_dictoffset */
2283 0, /* tp_init */
2284 0, /* tp_alloc */
2285 combinations_new, /* tp_new */
2286 PyObject_GC_Del, /* tp_free */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002287};
2288
2289
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002290/* combinations with replacement object *******************************************/
2291
2292/* Equivalent to:
2293
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002294 def combinations_with_replacement(iterable, r):
2295 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2296 # number items returned: (n+r-1)! / r! / (n-1)!
2297 pool = tuple(iterable)
2298 n = len(pool)
2299 indices = [0] * r
2300 yield tuple(pool[i] for i in indices)
2301 while 1:
2302 for i in reversed(range(r)):
2303 if indices[i] != n - 1:
2304 break
2305 else:
2306 return
2307 indices[i:] = [indices[i] + 1] * (r - i)
2308 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002309
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002310 def combinations_with_replacement2(iterable, r):
2311 'Alternate version that filters from product()'
2312 pool = tuple(iterable)
2313 n = len(pool)
2314 for indices in product(range(n), repeat=r):
2315 if sorted(indices) == list(indices):
2316 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002317*/
2318typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002319 PyObject_HEAD
2320 PyObject *pool; /* input converted to a tuple */
2321 Py_ssize_t *indices; /* one index per result element */
2322 PyObject *result; /* most recently returned result tuple */
2323 Py_ssize_t r; /* size of result tuple */
2324 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002325} cwrobject;
2326
2327static PyTypeObject cwr_type;
2328
2329static PyObject *
2330cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2331{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002332 cwrobject *co;
2333 Py_ssize_t n;
2334 Py_ssize_t r;
2335 PyObject *pool = NULL;
2336 PyObject *iterable = NULL;
2337 Py_ssize_t *indices = NULL;
2338 Py_ssize_t i;
2339 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002340
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002341 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2342 &iterable, &r))
2343 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002344
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002345 pool = PySequence_Tuple(iterable);
2346 if (pool == NULL)
2347 goto error;
2348 n = PyTuple_GET_SIZE(pool);
2349 if (r < 0) {
2350 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2351 goto error;
2352 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002353
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02002354 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002355 if (indices == NULL) {
2356 PyErr_NoMemory();
2357 goto error;
2358 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002359
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002360 for (i=0 ; i<r ; i++)
2361 indices[i] = 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002362
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002363 /* create cwrobject structure */
2364 co = (cwrobject *)type->tp_alloc(type, 0);
2365 if (co == NULL)
2366 goto error;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002367
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002368 co->pool = pool;
2369 co->indices = indices;
2370 co->result = NULL;
2371 co->r = r;
2372 co->stopped = !n && r;
2373
2374 return (PyObject *)co;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002375
2376error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002377 if (indices != NULL)
2378 PyMem_Free(indices);
2379 Py_XDECREF(pool);
2380 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002381}
2382
2383static void
2384cwr_dealloc(cwrobject *co)
2385{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002386 PyObject_GC_UnTrack(co);
2387 Py_XDECREF(co->pool);
2388 Py_XDECREF(co->result);
2389 if (co->indices != NULL)
2390 PyMem_Free(co->indices);
2391 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002392}
2393
2394static int
2395cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2396{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002397 Py_VISIT(co->pool);
2398 Py_VISIT(co->result);
2399 return 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002400}
2401
2402static PyObject *
2403cwr_next(cwrobject *co)
2404{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002405 PyObject *elem;
2406 PyObject *oldelem;
2407 PyObject *pool = co->pool;
2408 Py_ssize_t *indices = co->indices;
2409 PyObject *result = co->result;
2410 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2411 Py_ssize_t r = co->r;
2412 Py_ssize_t i, j, index;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002413
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002414 if (co->stopped)
2415 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002416
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002417 if (result == NULL) {
2418 /* On the first pass, initialize result tuple using the indices */
2419 result = PyTuple_New(r);
2420 if (result == NULL)
2421 goto empty;
2422 co->result = result;
2423 for (i=0; i<r ; i++) {
2424 index = indices[i];
2425 elem = PyTuple_GET_ITEM(pool, index);
2426 Py_INCREF(elem);
2427 PyTuple_SET_ITEM(result, i, elem);
2428 }
2429 } else {
2430 /* Copy the previous result tuple or re-use it if available */
2431 if (Py_REFCNT(result) > 1) {
2432 PyObject *old_result = result;
2433 result = PyTuple_New(r);
2434 if (result == NULL)
2435 goto empty;
2436 co->result = result;
2437 for (i=0; i<r ; i++) {
2438 elem = PyTuple_GET_ITEM(old_result, i);
2439 Py_INCREF(elem);
2440 PyTuple_SET_ITEM(result, i, elem);
2441 }
2442 Py_DECREF(old_result);
2443 }
2444 /* Now, we've got the only copy so we can update it in-place CPython's
2445 empty tuple is a singleton and cached in PyTuple's freelist. */
2446 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002447
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002448 /* Scan indices right-to-left until finding one that is not
2449 * at its maximum (n-1). */
2450 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2451 ;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002452
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002453 /* If i is negative, then the indices are all at
2454 their maximum value and we're done. */
2455 if (i < 0)
2456 goto empty;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002457
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002458 /* Increment the current index which we know is not at its
2459 maximum. Then set all to the right to the same value. */
2460 indices[i]++;
2461 for (j=i+1 ; j<r ; j++)
2462 indices[j] = indices[j-1];
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002463
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002464 /* Update the result tuple for the new indices
2465 starting with i, the leftmost index that changed */
2466 for ( ; i<r ; i++) {
2467 index = indices[i];
2468 elem = PyTuple_GET_ITEM(pool, index);
2469 Py_INCREF(elem);
2470 oldelem = PyTuple_GET_ITEM(result, i);
2471 PyTuple_SET_ITEM(result, i, elem);
2472 Py_DECREF(oldelem);
2473 }
2474 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002475
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002476 Py_INCREF(result);
2477 return result;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002478
2479empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002480 co->stopped = 1;
2481 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002482}
2483
2484PyDoc_STRVAR(cwr_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002485"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002486\n\
2487Return successive r-length combinations of elements in the iterable\n\
2488allowing individual elements to have successive repeats.\n\
2489combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2490
2491static PyTypeObject cwr_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002492 PyVarObject_HEAD_INIT(NULL, 0)
2493 "itertools.combinations_with_replacement", /* tp_name */
2494 sizeof(cwrobject), /* tp_basicsize */
2495 0, /* tp_itemsize */
2496 /* methods */
2497 (destructor)cwr_dealloc, /* tp_dealloc */
2498 0, /* tp_print */
2499 0, /* tp_getattr */
2500 0, /* tp_setattr */
2501 0, /* tp_compare */
2502 0, /* tp_repr */
2503 0, /* tp_as_number */
2504 0, /* tp_as_sequence */
2505 0, /* tp_as_mapping */
2506 0, /* tp_hash */
2507 0, /* tp_call */
2508 0, /* tp_str */
2509 PyObject_GenericGetAttr, /* tp_getattro */
2510 0, /* tp_setattro */
2511 0, /* tp_as_buffer */
2512 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2513 Py_TPFLAGS_BASETYPE, /* tp_flags */
2514 cwr_doc, /* tp_doc */
2515 (traverseproc)cwr_traverse, /* tp_traverse */
2516 0, /* tp_clear */
2517 0, /* tp_richcompare */
2518 0, /* tp_weaklistoffset */
2519 PyObject_SelfIter, /* tp_iter */
2520 (iternextfunc)cwr_next, /* tp_iternext */
2521 0, /* tp_methods */
2522 0, /* tp_members */
2523 0, /* tp_getset */
2524 0, /* tp_base */
2525 0, /* tp_dict */
2526 0, /* tp_descr_get */
2527 0, /* tp_descr_set */
2528 0, /* tp_dictoffset */
2529 0, /* tp_init */
2530 0, /* tp_alloc */
2531 cwr_new, /* tp_new */
2532 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002533};
2534
2535
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002536/* permutations object ************************************************************
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002537
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002538def permutations(iterable, r=None):
2539 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2540 pool = tuple(iterable)
2541 n = len(pool)
2542 r = n if r is None else r
2543 indices = range(n)
2544 cycles = range(n-r+1, n+1)[::-1]
2545 yield tuple(pool[i] for i in indices[:r])
2546 while n:
Serhiy Storchakaf0aa88f2015-06-11 00:06:27 +03002547 for i in reversed(range(r)):
2548 cycles[i] -= 1
2549 if cycles[i] == 0:
2550 indices[i:] = indices[i+1:] + indices[i:i+1]
2551 cycles[i] = n - i
2552 else:
2553 j = cycles[i]
2554 indices[i], indices[-j] = indices[-j], indices[i]
2555 yield tuple(pool[i] for i in indices[:r])
2556 break
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002557 else:
Serhiy Storchakaf0aa88f2015-06-11 00:06:27 +03002558 return
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002559*/
2560
2561typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002562 PyObject_HEAD
2563 PyObject *pool; /* input converted to a tuple */
2564 Py_ssize_t *indices; /* one index per element in the pool */
2565 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2566 PyObject *result; /* most recently returned result tuple */
2567 Py_ssize_t r; /* size of result tuple */
2568 int stopped; /* set to 1 when the permutations iterator is exhausted */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002569} permutationsobject;
2570
2571static PyTypeObject permutations_type;
2572
2573static PyObject *
2574permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2575{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002576 permutationsobject *po;
2577 Py_ssize_t n;
2578 Py_ssize_t r;
2579 PyObject *robj = Py_None;
2580 PyObject *pool = NULL;
2581 PyObject *iterable = NULL;
2582 Py_ssize_t *indices = NULL;
2583 Py_ssize_t *cycles = NULL;
2584 Py_ssize_t i;
2585 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002586
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002587 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2588 &iterable, &robj))
2589 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002590
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002591 pool = PySequence_Tuple(iterable);
2592 if (pool == NULL)
2593 goto error;
2594 n = PyTuple_GET_SIZE(pool);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002595
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002596 r = n;
2597 if (robj != Py_None) {
2598 r = PyInt_AsSsize_t(robj);
2599 if (r == -1 && PyErr_Occurred())
2600 goto error;
2601 }
2602 if (r < 0) {
2603 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2604 goto error;
2605 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002606
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02002607 indices = PyMem_New(Py_ssize_t, n);
2608 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002609 if (indices == NULL || cycles == NULL) {
2610 PyErr_NoMemory();
2611 goto error;
2612 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002613
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002614 for (i=0 ; i<n ; i++)
2615 indices[i] = i;
2616 for (i=0 ; i<r ; i++)
2617 cycles[i] = n - i;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002618
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002619 /* create permutationsobject structure */
2620 po = (permutationsobject *)type->tp_alloc(type, 0);
2621 if (po == NULL)
2622 goto error;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002623
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002624 po->pool = pool;
2625 po->indices = indices;
2626 po->cycles = cycles;
2627 po->result = NULL;
2628 po->r = r;
2629 po->stopped = r > n ? 1 : 0;
2630
2631 return (PyObject *)po;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002632
2633error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002634 if (indices != NULL)
2635 PyMem_Free(indices);
2636 if (cycles != NULL)
2637 PyMem_Free(cycles);
2638 Py_XDECREF(pool);
2639 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002640}
2641
2642static void
2643permutations_dealloc(permutationsobject *po)
2644{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002645 PyObject_GC_UnTrack(po);
2646 Py_XDECREF(po->pool);
2647 Py_XDECREF(po->result);
2648 PyMem_Free(po->indices);
2649 PyMem_Free(po->cycles);
2650 Py_TYPE(po)->tp_free(po);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002651}
2652
2653static int
2654permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2655{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002656 Py_VISIT(po->pool);
2657 Py_VISIT(po->result);
2658 return 0;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002659}
2660
2661static PyObject *
2662permutations_next(permutationsobject *po)
2663{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002664 PyObject *elem;
2665 PyObject *oldelem;
2666 PyObject *pool = po->pool;
2667 Py_ssize_t *indices = po->indices;
2668 Py_ssize_t *cycles = po->cycles;
2669 PyObject *result = po->result;
2670 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2671 Py_ssize_t r = po->r;
2672 Py_ssize_t i, j, k, index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002673
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002674 if (po->stopped)
2675 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002676
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002677 if (result == NULL) {
2678 /* On the first pass, initialize result tuple using the indices */
2679 result = PyTuple_New(r);
2680 if (result == NULL)
2681 goto empty;
2682 po->result = result;
2683 for (i=0; i<r ; i++) {
2684 index = indices[i];
2685 elem = PyTuple_GET_ITEM(pool, index);
2686 Py_INCREF(elem);
2687 PyTuple_SET_ITEM(result, i, elem);
2688 }
2689 } else {
2690 if (n == 0)
2691 goto empty;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002692
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002693 /* Copy the previous result tuple or re-use it if available */
2694 if (Py_REFCNT(result) > 1) {
2695 PyObject *old_result = result;
2696 result = PyTuple_New(r);
2697 if (result == NULL)
2698 goto empty;
2699 po->result = result;
2700 for (i=0; i<r ; i++) {
2701 elem = PyTuple_GET_ITEM(old_result, i);
2702 Py_INCREF(elem);
2703 PyTuple_SET_ITEM(result, i, elem);
2704 }
2705 Py_DECREF(old_result);
2706 }
2707 /* Now, we've got the only copy so we can update it in-place */
2708 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002709
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002710 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2711 for (i=r-1 ; i>=0 ; i--) {
2712 cycles[i] -= 1;
2713 if (cycles[i] == 0) {
2714 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2715 index = indices[i];
2716 for (j=i ; j<n-1 ; j++)
2717 indices[j] = indices[j+1];
2718 indices[n-1] = index;
2719 cycles[i] = n - i;
2720 } else {
2721 j = cycles[i];
2722 index = indices[i];
2723 indices[i] = indices[n-j];
2724 indices[n-j] = index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002725
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002726 for (k=i; k<r ; k++) {
2727 /* start with i, the leftmost element that changed */
2728 /* yield tuple(pool[k] for k in indices[:r]) */
2729 index = indices[k];
2730 elem = PyTuple_GET_ITEM(pool, index);
2731 Py_INCREF(elem);
2732 oldelem = PyTuple_GET_ITEM(result, k);
2733 PyTuple_SET_ITEM(result, k, elem);
2734 Py_DECREF(oldelem);
2735 }
2736 break;
2737 }
2738 }
2739 /* If i is negative, then the cycles have all
2740 rolled-over and we're done. */
2741 if (i < 0)
2742 goto empty;
2743 }
2744 Py_INCREF(result);
2745 return result;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002746
2747empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002748 po->stopped = 1;
2749 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002750}
2751
2752PyDoc_STRVAR(permutations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +00002753"permutations(iterable[, r]) --> permutations object\n\
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002754\n\
2755Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl8550d692008-06-10 12:46:39 +00002756permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002757
2758static PyTypeObject permutations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002759 PyVarObject_HEAD_INIT(NULL, 0)
2760 "itertools.permutations", /* tp_name */
2761 sizeof(permutationsobject), /* tp_basicsize */
2762 0, /* tp_itemsize */
2763 /* methods */
2764 (destructor)permutations_dealloc, /* tp_dealloc */
2765 0, /* tp_print */
2766 0, /* tp_getattr */
2767 0, /* tp_setattr */
2768 0, /* tp_compare */
2769 0, /* tp_repr */
2770 0, /* tp_as_number */
2771 0, /* tp_as_sequence */
2772 0, /* tp_as_mapping */
2773 0, /* tp_hash */
2774 0, /* tp_call */
2775 0, /* tp_str */
2776 PyObject_GenericGetAttr, /* tp_getattro */
2777 0, /* tp_setattro */
2778 0, /* tp_as_buffer */
2779 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2780 Py_TPFLAGS_BASETYPE, /* tp_flags */
2781 permutations_doc, /* tp_doc */
2782 (traverseproc)permutations_traverse, /* tp_traverse */
2783 0, /* tp_clear */
2784 0, /* tp_richcompare */
2785 0, /* tp_weaklistoffset */
2786 PyObject_SelfIter, /* tp_iter */
2787 (iternextfunc)permutations_next, /* tp_iternext */
2788 0, /* tp_methods */
2789 0, /* tp_members */
2790 0, /* tp_getset */
2791 0, /* tp_base */
2792 0, /* tp_dict */
2793 0, /* tp_descr_get */
2794 0, /* tp_descr_set */
2795 0, /* tp_dictoffset */
2796 0, /* tp_init */
2797 0, /* tp_alloc */
2798 permutations_new, /* tp_new */
2799 PyObject_GC_Del, /* tp_free */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002800};
2801
2802
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002803/* compress object ************************************************************/
2804
2805/* Equivalent to:
2806
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002807 def compress(data, selectors):
2808 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2809 return (d for d, s in izip(data, selectors) if s)
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002810*/
2811
2812typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002813 PyObject_HEAD
2814 PyObject *data;
2815 PyObject *selectors;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002816} compressobject;
2817
2818static PyTypeObject compress_type;
2819
2820static PyObject *
2821compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2822{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002823 PyObject *seq1, *seq2;
2824 PyObject *data=NULL, *selectors=NULL;
2825 compressobject *lz;
2826 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002827
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002828 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2829 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002830
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002831 data = PyObject_GetIter(seq1);
2832 if (data == NULL)
2833 goto fail;
2834 selectors = PyObject_GetIter(seq2);
2835 if (selectors == NULL)
2836 goto fail;
2837
2838 /* create compressobject structure */
2839 lz = (compressobject *)type->tp_alloc(type, 0);
2840 if (lz == NULL)
2841 goto fail;
2842 lz->data = data;
2843 lz->selectors = selectors;
2844 return (PyObject *)lz;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002845
2846fail:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002847 Py_XDECREF(data);
2848 Py_XDECREF(selectors);
2849 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002850}
2851
2852static void
2853compress_dealloc(compressobject *lz)
2854{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002855 PyObject_GC_UnTrack(lz);
2856 Py_XDECREF(lz->data);
2857 Py_XDECREF(lz->selectors);
2858 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002859}
2860
2861static int
2862compress_traverse(compressobject *lz, visitproc visit, void *arg)
2863{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002864 Py_VISIT(lz->data);
2865 Py_VISIT(lz->selectors);
2866 return 0;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002867}
2868
2869static PyObject *
2870compress_next(compressobject *lz)
2871{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002872 PyObject *data = lz->data, *selectors = lz->selectors;
2873 PyObject *datum, *selector;
2874 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2875 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2876 int ok;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002877
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002878 while (1) {
2879 /* Steps: get datum, get selector, evaluate selector.
2880 Order is important (to match the pure python version
2881 in terms of which input gets a chance to raise an
2882 exception first).
2883 */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002884
Serhiy Storchakabb845652013-04-06 22:04:10 +03002885 datum = datanext(data);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03002886 if (datum == NULL)
Serhiy Storchakabb845652013-04-06 22:04:10 +03002887 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002888
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002889 selector = selectornext(selectors);
2890 if (selector == NULL) {
2891 Py_DECREF(datum);
2892 return NULL;
2893 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002894
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002895 ok = PyObject_IsTrue(selector);
2896 Py_DECREF(selector);
2897 if (ok == 1)
2898 return datum;
2899 Py_DECREF(datum);
2900 if (ok == -1)
2901 return NULL;
2902 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002903}
2904
2905PyDoc_STRVAR(compress_doc,
Raymond Hettinger2e2909f2009-02-19 02:15:14 +00002906"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002907\n\
2908Return data elements corresponding to true selector elements.\n\
2909Forms a shorter iterator from selected data elements using the\n\
2910selectors to choose the data elements.");
2911
2912static PyTypeObject compress_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002913 PyVarObject_HEAD_INIT(NULL, 0)
2914 "itertools.compress", /* tp_name */
2915 sizeof(compressobject), /* tp_basicsize */
2916 0, /* tp_itemsize */
2917 /* methods */
2918 (destructor)compress_dealloc, /* tp_dealloc */
2919 0, /* tp_print */
2920 0, /* tp_getattr */
2921 0, /* tp_setattr */
2922 0, /* tp_compare */
2923 0, /* tp_repr */
2924 0, /* tp_as_number */
2925 0, /* tp_as_sequence */
2926 0, /* tp_as_mapping */
2927 0, /* tp_hash */
2928 0, /* tp_call */
2929 0, /* tp_str */
2930 PyObject_GenericGetAttr, /* tp_getattro */
2931 0, /* tp_setattro */
2932 0, /* tp_as_buffer */
2933 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2934 Py_TPFLAGS_BASETYPE, /* tp_flags */
2935 compress_doc, /* tp_doc */
2936 (traverseproc)compress_traverse, /* tp_traverse */
2937 0, /* tp_clear */
2938 0, /* tp_richcompare */
2939 0, /* tp_weaklistoffset */
2940 PyObject_SelfIter, /* tp_iter */
2941 (iternextfunc)compress_next, /* tp_iternext */
2942 0, /* tp_methods */
2943 0, /* tp_members */
2944 0, /* tp_getset */
2945 0, /* tp_base */
2946 0, /* tp_dict */
2947 0, /* tp_descr_get */
2948 0, /* tp_descr_set */
2949 0, /* tp_dictoffset */
2950 0, /* tp_init */
2951 0, /* tp_alloc */
2952 compress_new, /* tp_new */
2953 PyObject_GC_Del, /* tp_free */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002954};
2955
2956
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002957/* ifilter object ************************************************************/
2958
2959typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002960 PyObject_HEAD
2961 PyObject *func;
2962 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002963} ifilterobject;
2964
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002965static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002966
2967static PyObject *
2968ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2969{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002970 PyObject *func, *seq;
2971 PyObject *it;
2972 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002973
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002974 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
2975 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002976
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002977 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
2978 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002979
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002980 /* Get iterator. */
2981 it = PyObject_GetIter(seq);
2982 if (it == NULL)
2983 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002984
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002985 /* create ifilterobject structure */
2986 lz = (ifilterobject *)type->tp_alloc(type, 0);
2987 if (lz == NULL) {
2988 Py_DECREF(it);
2989 return NULL;
2990 }
2991 Py_INCREF(func);
2992 lz->func = func;
2993 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002994
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002995 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002996}
2997
2998static void
2999ifilter_dealloc(ifilterobject *lz)
3000{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003001 PyObject_GC_UnTrack(lz);
3002 Py_XDECREF(lz->func);
3003 Py_XDECREF(lz->it);
3004 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003005}
3006
3007static int
3008ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
3009{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003010 Py_VISIT(lz->it);
3011 Py_VISIT(lz->func);
3012 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003013}
3014
3015static PyObject *
3016ifilter_next(ifilterobject *lz)
3017{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003018 PyObject *item;
3019 PyObject *it = lz->it;
3020 long ok;
3021 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003022
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003023 iternext = *Py_TYPE(it)->tp_iternext;
3024 for (;;) {
3025 item = iternext(it);
3026 if (item == NULL)
3027 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003028
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003029 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3030 ok = PyObject_IsTrue(item);
3031 } else {
3032 PyObject *good;
3033 good = PyObject_CallFunctionObjArgs(lz->func,
3034 item, NULL);
3035 if (good == NULL) {
3036 Py_DECREF(item);
3037 return NULL;
3038 }
3039 ok = PyObject_IsTrue(good);
3040 Py_DECREF(good);
3041 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003042 if (ok > 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003043 return item;
3044 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003045 if (ok < 0)
3046 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003047 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003048}
3049
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003050PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003051"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003052\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003053Return those items of sequence for which function(item) is true.\n\
3054If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003055
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003056static PyTypeObject ifilter_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003057 PyVarObject_HEAD_INIT(NULL, 0)
3058 "itertools.ifilter", /* tp_name */
3059 sizeof(ifilterobject), /* tp_basicsize */
3060 0, /* tp_itemsize */
3061 /* methods */
3062 (destructor)ifilter_dealloc, /* tp_dealloc */
3063 0, /* tp_print */
3064 0, /* tp_getattr */
3065 0, /* tp_setattr */
3066 0, /* tp_compare */
3067 0, /* tp_repr */
3068 0, /* tp_as_number */
3069 0, /* tp_as_sequence */
3070 0, /* tp_as_mapping */
3071 0, /* tp_hash */
3072 0, /* tp_call */
3073 0, /* tp_str */
3074 PyObject_GenericGetAttr, /* tp_getattro */
3075 0, /* tp_setattro */
3076 0, /* tp_as_buffer */
3077 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3078 Py_TPFLAGS_BASETYPE, /* tp_flags */
3079 ifilter_doc, /* tp_doc */
3080 (traverseproc)ifilter_traverse, /* tp_traverse */
3081 0, /* tp_clear */
3082 0, /* tp_richcompare */
3083 0, /* tp_weaklistoffset */
3084 PyObject_SelfIter, /* tp_iter */
3085 (iternextfunc)ifilter_next, /* tp_iternext */
3086 0, /* tp_methods */
3087 0, /* tp_members */
3088 0, /* tp_getset */
3089 0, /* tp_base */
3090 0, /* tp_dict */
3091 0, /* tp_descr_get */
3092 0, /* tp_descr_set */
3093 0, /* tp_dictoffset */
3094 0, /* tp_init */
3095 0, /* tp_alloc */
3096 ifilter_new, /* tp_new */
3097 PyObject_GC_Del, /* tp_free */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003098};
3099
3100
3101/* ifilterfalse object ************************************************************/
3102
3103typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003104 PyObject_HEAD
3105 PyObject *func;
3106 PyObject *it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003107} ifilterfalseobject;
3108
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003109static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003110
3111static PyObject *
3112ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3113{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003114 PyObject *func, *seq;
3115 PyObject *it;
3116 ifilterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003117
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003118 if (type == &ifilterfalse_type &&
3119 !_PyArg_NoKeywords("ifilterfalse()", kwds))
3120 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003121
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003122 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
3123 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003124
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003125 /* Get iterator. */
3126 it = PyObject_GetIter(seq);
3127 if (it == NULL)
3128 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003129
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003130 /* create ifilterfalseobject structure */
3131 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
3132 if (lz == NULL) {
3133 Py_DECREF(it);
3134 return NULL;
3135 }
3136 Py_INCREF(func);
3137 lz->func = func;
3138 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003139
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003140 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003141}
3142
3143static void
3144ifilterfalse_dealloc(ifilterfalseobject *lz)
3145{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003146 PyObject_GC_UnTrack(lz);
3147 Py_XDECREF(lz->func);
3148 Py_XDECREF(lz->it);
3149 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003150}
3151
3152static int
3153ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
3154{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003155 Py_VISIT(lz->it);
3156 Py_VISIT(lz->func);
3157 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003158}
3159
3160static PyObject *
3161ifilterfalse_next(ifilterfalseobject *lz)
3162{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003163 PyObject *item;
3164 PyObject *it = lz->it;
3165 long ok;
3166 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003167
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003168 iternext = *Py_TYPE(it)->tp_iternext;
3169 for (;;) {
3170 item = iternext(it);
3171 if (item == NULL)
3172 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003173
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003174 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3175 ok = PyObject_IsTrue(item);
3176 } else {
3177 PyObject *good;
3178 good = PyObject_CallFunctionObjArgs(lz->func,
3179 item, NULL);
3180 if (good == NULL) {
3181 Py_DECREF(item);
3182 return NULL;
3183 }
3184 ok = PyObject_IsTrue(good);
3185 Py_DECREF(good);
3186 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003187 if (ok == 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003188 return item;
3189 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003190 if (ok < 0)
3191 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003192 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003193}
3194
Raymond Hettinger60eca932003-02-09 06:40:58 +00003195PyDoc_STRVAR(ifilterfalse_doc,
3196"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
3197\n\
3198Return those items of sequence for which function(item) is false.\n\
3199If function is None, return the items that are false.");
3200
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003201static PyTypeObject ifilterfalse_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003202 PyVarObject_HEAD_INIT(NULL, 0)
3203 "itertools.ifilterfalse", /* tp_name */
3204 sizeof(ifilterfalseobject), /* tp_basicsize */
3205 0, /* tp_itemsize */
3206 /* methods */
3207 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
3208 0, /* tp_print */
3209 0, /* tp_getattr */
3210 0, /* tp_setattr */
3211 0, /* tp_compare */
3212 0, /* tp_repr */
3213 0, /* tp_as_number */
3214 0, /* tp_as_sequence */
3215 0, /* tp_as_mapping */
3216 0, /* tp_hash */
3217 0, /* tp_call */
3218 0, /* tp_str */
3219 PyObject_GenericGetAttr, /* tp_getattro */
3220 0, /* tp_setattro */
3221 0, /* tp_as_buffer */
3222 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3223 Py_TPFLAGS_BASETYPE, /* tp_flags */
3224 ifilterfalse_doc, /* tp_doc */
3225 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
3226 0, /* tp_clear */
3227 0, /* tp_richcompare */
3228 0, /* tp_weaklistoffset */
3229 PyObject_SelfIter, /* tp_iter */
3230 (iternextfunc)ifilterfalse_next, /* tp_iternext */
3231 0, /* tp_methods */
3232 0, /* tp_members */
3233 0, /* tp_getset */
3234 0, /* tp_base */
3235 0, /* tp_dict */
3236 0, /* tp_descr_get */
3237 0, /* tp_descr_set */
3238 0, /* tp_dictoffset */
3239 0, /* tp_init */
3240 0, /* tp_alloc */
3241 ifilterfalse_new, /* tp_new */
3242 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003243};
3244
3245
3246/* count object ************************************************************/
3247
3248typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003249 PyObject_HEAD
3250 Py_ssize_t cnt;
3251 PyObject *long_cnt;
3252 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003253} countobject;
3254
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003255/* Counting logic and invariants:
3256
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003257fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003258
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003259 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3260 Advances with: cnt += 1
3261 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003262
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003263slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003264
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003265 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3266 All counting is done with python objects (no overflows or underflows).
3267 Advances with: long_cnt += long_step
3268 Step may be zero -- effectively a slow version of repeat(cnt).
3269 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003270*/
3271
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003272static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003273
3274static PyObject *
3275count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3276{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003277 countobject *lz;
3278 int slow_mode = 0;
3279 Py_ssize_t cnt = 0;
3280 PyObject *long_cnt = NULL;
3281 PyObject *long_step = NULL;
3282 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003283
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003284 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3285 kwlist, &long_cnt, &long_step))
3286 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003287
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003288 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3289 (long_step != NULL && !PyNumber_Check(long_step))) {
3290 PyErr_SetString(PyExc_TypeError, "a number is required");
3291 return NULL;
3292 }
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003293
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003294 if (long_cnt != NULL) {
3295 cnt = PyInt_AsSsize_t(long_cnt);
3296 if ((cnt == -1 && PyErr_Occurred()) || !PyInt_Check(long_cnt)) {
3297 PyErr_Clear();
3298 slow_mode = 1;
3299 }
3300 Py_INCREF(long_cnt);
3301 } else {
3302 cnt = 0;
3303 long_cnt = PyInt_FromLong(0);
3304 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003305
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003306 /* If not specified, step defaults to 1 */
3307 if (long_step == NULL) {
3308 long_step = PyInt_FromLong(1);
3309 if (long_step == NULL) {
3310 Py_DECREF(long_cnt);
3311 return NULL;
3312 }
3313 } else
3314 Py_INCREF(long_step);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003315
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003316 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003317
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003318 /* Fast mode only works when the step is 1 */
3319 if (!PyInt_Check(long_step) ||
3320 PyInt_AS_LONG(long_step) != 1) {
3321 slow_mode = 1;
3322 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003323
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003324 if (slow_mode)
3325 cnt = PY_SSIZE_T_MAX;
3326 else
3327 Py_CLEAR(long_cnt);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003328
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003329 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3330 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3331 assert(slow_mode ||
3332 (PyInt_Check(long_step) && PyInt_AS_LONG(long_step) == 1));
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003333
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003334 /* create countobject structure */
3335 lz = (countobject *)type->tp_alloc(type, 0);
3336 if (lz == NULL) {
3337 Py_XDECREF(long_cnt);
3338 return NULL;
3339 }
3340 lz->cnt = cnt;
3341 lz->long_cnt = long_cnt;
3342 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003343
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003344 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003345}
3346
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003347static void
3348count_dealloc(countobject *lz)
3349{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003350 PyObject_GC_UnTrack(lz);
3351 Py_XDECREF(lz->long_cnt);
3352 Py_XDECREF(lz->long_step);
3353 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003354}
3355
Benjamin Peterson062a7c32009-02-16 21:07:52 +00003356static int
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003357count_traverse(countobject *lz, visitproc visit, void *arg)
3358{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003359 Py_VISIT(lz->long_cnt);
3360 Py_VISIT(lz->long_step);
3361 return 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003362}
3363
3364static PyObject *
3365count_nextlong(countobject *lz)
3366{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003367 PyObject *long_cnt;
3368 PyObject *stepped_up;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003369
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003370 long_cnt = lz->long_cnt;
3371 if (long_cnt == NULL) {
3372 /* Switch to slow_mode */
3373 long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
3374 if (long_cnt == NULL)
3375 return NULL;
3376 }
3377 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003378
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003379 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3380 if (stepped_up == NULL)
3381 return NULL;
3382 lz->long_cnt = stepped_up;
3383 return long_cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003384}
3385
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003386static PyObject *
3387count_next(countobject *lz)
3388{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003389 if (lz->cnt == PY_SSIZE_T_MAX)
3390 return count_nextlong(lz);
3391 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003392}
3393
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003394static PyObject *
3395count_repr(countobject *lz)
3396{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003397 PyObject *cnt_repr, *step_repr = NULL;
3398 PyObject *result = NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003399
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003400 if (lz->cnt != PY_SSIZE_T_MAX)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003401 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003402
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003403 cnt_repr = PyObject_Repr(lz->long_cnt);
3404 if (cnt_repr == NULL)
3405 return NULL;
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003406
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003407 if (PyInt_Check(lz->long_step) && PyInt_AS_LONG(lz->long_step) == 1) {
3408 /* Don't display step when it is an integer equal to 1 */
3409 result = PyString_FromFormat("count(%s)",
3410 PyString_AS_STRING(cnt_repr));
3411 } else {
3412 step_repr = PyObject_Repr(lz->long_step);
3413 if (step_repr != NULL)
3414 result = PyString_FromFormat("count(%s, %s)",
3415 PyString_AS_STRING(cnt_repr),
3416 PyString_AS_STRING(step_repr));
3417 }
3418 Py_DECREF(cnt_repr);
3419 Py_XDECREF(step_repr);
3420 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003421}
3422
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003423static PyObject *
3424count_reduce(countobject *lz)
3425{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003426 if (lz->cnt == PY_SSIZE_T_MAX)
3427 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3428 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003429}
3430
3431PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3432
3433static PyMethodDef count_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003434 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3435 count_reduce_doc},
3436 {NULL, NULL} /* sentinel */
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003437};
3438
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003439PyDoc_STRVAR(count_doc,
Ezio Melottifdc1e0d2010-06-11 02:21:25 +00003440 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003441\n\
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003442Return a count object whose .next() method returns consecutive values.\n\
3443Equivalent to:\n\n\
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003444 def count(firstval=0, step=1):\n\
Eli Benderskye6802db2013-09-03 06:41:58 -07003445 x = firstval\n\
3446 while 1:\n\
3447 yield x\n\
3448 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003449
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003450static PyTypeObject count_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003451 PyVarObject_HEAD_INIT(NULL, 0)
3452 "itertools.count", /* tp_name */
3453 sizeof(countobject), /* tp_basicsize */
3454 0, /* tp_itemsize */
3455 /* methods */
3456 (destructor)count_dealloc, /* tp_dealloc */
3457 0, /* tp_print */
3458 0, /* tp_getattr */
3459 0, /* tp_setattr */
3460 0, /* tp_compare */
3461 (reprfunc)count_repr, /* tp_repr */
3462 0, /* tp_as_number */
3463 0, /* tp_as_sequence */
3464 0, /* tp_as_mapping */
3465 0, /* tp_hash */
3466 0, /* tp_call */
3467 0, /* tp_str */
3468 PyObject_GenericGetAttr, /* tp_getattro */
3469 0, /* tp_setattro */
3470 0, /* tp_as_buffer */
3471 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3472 Py_TPFLAGS_BASETYPE, /* tp_flags */
3473 count_doc, /* tp_doc */
3474 (traverseproc)count_traverse, /* tp_traverse */
3475 0, /* tp_clear */
3476 0, /* tp_richcompare */
3477 0, /* tp_weaklistoffset */
3478 PyObject_SelfIter, /* tp_iter */
3479 (iternextfunc)count_next, /* tp_iternext */
3480 count_methods, /* tp_methods */
3481 0, /* tp_members */
3482 0, /* tp_getset */
3483 0, /* tp_base */
3484 0, /* tp_dict */
3485 0, /* tp_descr_get */
3486 0, /* tp_descr_set */
3487 0, /* tp_dictoffset */
3488 0, /* tp_init */
3489 0, /* tp_alloc */
3490 count_new, /* tp_new */
3491 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003492};
3493
3494
3495/* izip object ************************************************************/
3496
3497#include "Python.h"
3498
3499typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003500 PyObject_HEAD
3501 Py_ssize_t tuplesize;
3502 PyObject *ittuple; /* tuple of iterators */
3503 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003504} izipobject;
3505
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003506static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003507
3508static PyObject *
3509izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3510{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003511 izipobject *lz;
3512 Py_ssize_t i;
3513 PyObject *ittuple; /* tuple of iterators */
3514 PyObject *result;
3515 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003516
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003517 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
3518 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003519
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003520 /* args must be a tuple */
3521 assert(PyTuple_Check(args));
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003522
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003523 /* obtain iterators */
3524 ittuple = PyTuple_New(tuplesize);
3525 if (ittuple == NULL)
3526 return NULL;
3527 for (i=0; i < tuplesize; ++i) {
3528 PyObject *item = PyTuple_GET_ITEM(args, i);
3529 PyObject *it = PyObject_GetIter(item);
3530 if (it == NULL) {
3531 if (PyErr_ExceptionMatches(PyExc_TypeError))
3532 PyErr_Format(PyExc_TypeError,
3533 "izip argument #%zd must support iteration",
3534 i+1);
3535 Py_DECREF(ittuple);
3536 return NULL;
3537 }
3538 PyTuple_SET_ITEM(ittuple, i, it);
3539 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003540
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003541 /* create a result holder */
3542 result = PyTuple_New(tuplesize);
3543 if (result == NULL) {
3544 Py_DECREF(ittuple);
3545 return NULL;
3546 }
3547 for (i=0 ; i < tuplesize ; i++) {
3548 Py_INCREF(Py_None);
3549 PyTuple_SET_ITEM(result, i, Py_None);
3550 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003551
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003552 /* create izipobject structure */
3553 lz = (izipobject *)type->tp_alloc(type, 0);
3554 if (lz == NULL) {
3555 Py_DECREF(ittuple);
3556 Py_DECREF(result);
3557 return NULL;
3558 }
3559 lz->ittuple = ittuple;
3560 lz->tuplesize = tuplesize;
3561 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003562
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003563 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003564}
3565
3566static void
3567izip_dealloc(izipobject *lz)
3568{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003569 PyObject_GC_UnTrack(lz);
3570 Py_XDECREF(lz->ittuple);
3571 Py_XDECREF(lz->result);
3572 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003573}
3574
3575static int
3576izip_traverse(izipobject *lz, visitproc visit, void *arg)
3577{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003578 Py_VISIT(lz->ittuple);
3579 Py_VISIT(lz->result);
3580 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003581}
3582
3583static PyObject *
3584izip_next(izipobject *lz)
3585{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003586 Py_ssize_t i;
3587 Py_ssize_t tuplesize = lz->tuplesize;
3588 PyObject *result = lz->result;
3589 PyObject *it;
3590 PyObject *item;
3591 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003592
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003593 if (tuplesize == 0)
3594 return NULL;
3595 if (Py_REFCNT(result) == 1) {
3596 Py_INCREF(result);
3597 for (i=0 ; i < tuplesize ; i++) {
3598 it = PyTuple_GET_ITEM(lz->ittuple, i);
3599 item = (*Py_TYPE(it)->tp_iternext)(it);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003600 if (item == NULL) {
3601 Py_DECREF(result);
3602 return NULL;
3603 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003604 olditem = PyTuple_GET_ITEM(result, i);
3605 PyTuple_SET_ITEM(result, i, item);
3606 Py_DECREF(olditem);
3607 }
3608 } else {
3609 result = PyTuple_New(tuplesize);
3610 if (result == NULL)
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003611 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003612 for (i=0 ; i < tuplesize ; i++) {
3613 it = PyTuple_GET_ITEM(lz->ittuple, i);
3614 item = (*Py_TYPE(it)->tp_iternext)(it);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003615 if (item == NULL) {
3616 Py_DECREF(result);
3617 return NULL;
3618 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003619 PyTuple_SET_ITEM(result, i, item);
3620 }
3621 }
3622 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003623}
3624
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003625PyDoc_STRVAR(izip_doc,
3626"izip(iter1 [,iter2 [...]]) --> izip object\n\
3627\n\
3628Return a izip object whose .next() method returns a tuple where\n\
3629the i-th element comes from the i-th iterable argument. The .next()\n\
3630method continues until the shortest iterable in the argument sequence\n\
3631is exhausted and then it raises StopIteration. Works like the zip()\n\
3632function but consumes less memory by returning an iterator instead of\n\
3633a list.");
3634
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003635static PyTypeObject izip_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003636 PyVarObject_HEAD_INIT(NULL, 0)
3637 "itertools.izip", /* tp_name */
3638 sizeof(izipobject), /* tp_basicsize */
3639 0, /* tp_itemsize */
3640 /* methods */
3641 (destructor)izip_dealloc, /* tp_dealloc */
3642 0, /* tp_print */
3643 0, /* tp_getattr */
3644 0, /* tp_setattr */
3645 0, /* tp_compare */
3646 0, /* tp_repr */
3647 0, /* tp_as_number */
3648 0, /* tp_as_sequence */
3649 0, /* tp_as_mapping */
3650 0, /* tp_hash */
3651 0, /* tp_call */
3652 0, /* tp_str */
3653 PyObject_GenericGetAttr, /* tp_getattro */
3654 0, /* tp_setattro */
3655 0, /* tp_as_buffer */
3656 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3657 Py_TPFLAGS_BASETYPE, /* tp_flags */
3658 izip_doc, /* tp_doc */
3659 (traverseproc)izip_traverse, /* tp_traverse */
3660 0, /* tp_clear */
3661 0, /* tp_richcompare */
3662 0, /* tp_weaklistoffset */
3663 PyObject_SelfIter, /* tp_iter */
3664 (iternextfunc)izip_next, /* tp_iternext */
3665 0, /* tp_methods */
3666 0, /* tp_members */
3667 0, /* tp_getset */
3668 0, /* tp_base */
3669 0, /* tp_dict */
3670 0, /* tp_descr_get */
3671 0, /* tp_descr_set */
3672 0, /* tp_dictoffset */
3673 0, /* tp_init */
3674 0, /* tp_alloc */
3675 izip_new, /* tp_new */
3676 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003677};
3678
3679
3680/* repeat object ************************************************************/
3681
3682typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003683 PyObject_HEAD
3684 PyObject *element;
3685 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003686} repeatobject;
3687
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003688static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003689
3690static PyObject *
3691repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3692{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003693 repeatobject *ro;
3694 PyObject *element;
Raymond Hettinger58ad2452014-06-24 21:53:45 -07003695 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003696 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003697
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003698 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3699 &element, &cnt))
3700 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003701
Raymond Hettinger58ad2452014-06-24 21:53:45 -07003702 if (kwds != NULL)
3703 n_kwds = PyDict_Size(kwds);
3704 /* Does user supply times argument? */
3705 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003706 cnt = 0;
3707
3708 ro = (repeatobject *)type->tp_alloc(type, 0);
3709 if (ro == NULL)
3710 return NULL;
3711 Py_INCREF(element);
3712 ro->element = element;
3713 ro->cnt = cnt;
3714 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003715}
3716
3717static void
3718repeat_dealloc(repeatobject *ro)
3719{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003720 PyObject_GC_UnTrack(ro);
3721 Py_XDECREF(ro->element);
3722 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003723}
3724
3725static int
3726repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3727{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003728 Py_VISIT(ro->element);
3729 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003730}
3731
3732static PyObject *
3733repeat_next(repeatobject *ro)
3734{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003735 if (ro->cnt == 0)
3736 return NULL;
3737 if (ro->cnt > 0)
3738 ro->cnt--;
3739 Py_INCREF(ro->element);
3740 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003741}
3742
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003743static PyObject *
3744repeat_repr(repeatobject *ro)
3745{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003746 PyObject *result, *objrepr;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003747
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003748 objrepr = PyObject_Repr(ro->element);
3749 if (objrepr == NULL)
3750 return NULL;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003751
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003752 if (ro->cnt == -1)
3753 result = PyString_FromFormat("repeat(%s)",
3754 PyString_AS_STRING(objrepr));
3755 else
3756 result = PyString_FromFormat("repeat(%s, %zd)",
3757 PyString_AS_STRING(objrepr), ro->cnt);
3758 Py_DECREF(objrepr);
3759 return result;
3760}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003761
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003762static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003763repeat_len(repeatobject *ro)
3764{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003765 if (ro->cnt == -1) {
3766 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3767 return NULL;
3768 }
3769 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003770}
3771
Armin Rigof5b3e362006-02-11 21:32:43 +00003772PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003773
3774static PyMethodDef repeat_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003775 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3776 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003777};
3778
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003779PyDoc_STRVAR(repeat_doc,
Raymond Hettinger182edae2009-02-19 02:38:25 +00003780"repeat(object [,times]) -> create an iterator which returns the object\n\
3781for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003782endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003783
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003784static PyTypeObject repeat_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003785 PyVarObject_HEAD_INIT(NULL, 0)
3786 "itertools.repeat", /* tp_name */
3787 sizeof(repeatobject), /* tp_basicsize */
3788 0, /* tp_itemsize */
3789 /* methods */
3790 (destructor)repeat_dealloc, /* tp_dealloc */
3791 0, /* tp_print */
3792 0, /* tp_getattr */
3793 0, /* tp_setattr */
3794 0, /* tp_compare */
3795 (reprfunc)repeat_repr, /* tp_repr */
3796 0, /* tp_as_number */
3797 0, /* tp_as_sequence */
3798 0, /* tp_as_mapping */
3799 0, /* tp_hash */
3800 0, /* tp_call */
3801 0, /* tp_str */
3802 PyObject_GenericGetAttr, /* tp_getattro */
3803 0, /* tp_setattro */
3804 0, /* tp_as_buffer */
3805 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3806 Py_TPFLAGS_BASETYPE, /* tp_flags */
3807 repeat_doc, /* tp_doc */
3808 (traverseproc)repeat_traverse, /* tp_traverse */
3809 0, /* tp_clear */
3810 0, /* tp_richcompare */
3811 0, /* tp_weaklistoffset */
3812 PyObject_SelfIter, /* tp_iter */
3813 (iternextfunc)repeat_next, /* tp_iternext */
3814 repeat_methods, /* tp_methods */
3815 0, /* tp_members */
3816 0, /* tp_getset */
3817 0, /* tp_base */
3818 0, /* tp_dict */
3819 0, /* tp_descr_get */
3820 0, /* tp_descr_set */
3821 0, /* tp_dictoffset */
3822 0, /* tp_init */
3823 0, /* tp_alloc */
3824 repeat_new, /* tp_new */
3825 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003826};
3827
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003828/* iziplongest object ************************************************************/
3829
3830#include "Python.h"
3831
3832typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003833 PyObject_HEAD
3834 Py_ssize_t tuplesize;
3835 Py_ssize_t numactive;
3836 PyObject *ittuple; /* tuple of iterators */
3837 PyObject *result;
3838 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003839} iziplongestobject;
3840
3841static PyTypeObject iziplongest_type;
3842
3843static PyObject *
3844izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3845{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003846 iziplongestobject *lz;
3847 Py_ssize_t i;
3848 PyObject *ittuple; /* tuple of iterators */
3849 PyObject *result;
3850 PyObject *fillvalue = Py_None;
3851 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003852
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003853 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3854 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3855 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3856 PyErr_SetString(PyExc_TypeError,
3857 "izip_longest() got an unexpected keyword argument");
3858 return NULL;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003859 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003860 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003861
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003862 /* args must be a tuple */
3863 assert(PyTuple_Check(args));
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003864
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003865 /* obtain iterators */
3866 ittuple = PyTuple_New(tuplesize);
3867 if (ittuple == NULL)
3868 return NULL;
3869 for (i=0; i < tuplesize; ++i) {
3870 PyObject *item = PyTuple_GET_ITEM(args, i);
3871 PyObject *it = PyObject_GetIter(item);
3872 if (it == NULL) {
3873 if (PyErr_ExceptionMatches(PyExc_TypeError))
3874 PyErr_Format(PyExc_TypeError,
3875 "izip_longest argument #%zd must support iteration",
3876 i+1);
3877 Py_DECREF(ittuple);
3878 return NULL;
3879 }
3880 PyTuple_SET_ITEM(ittuple, i, it);
3881 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003882
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003883 /* create a result holder */
3884 result = PyTuple_New(tuplesize);
3885 if (result == NULL) {
3886 Py_DECREF(ittuple);
3887 return NULL;
3888 }
3889 for (i=0 ; i < tuplesize ; i++) {
3890 Py_INCREF(Py_None);
3891 PyTuple_SET_ITEM(result, i, Py_None);
3892 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003893
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003894 /* create iziplongestobject structure */
3895 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3896 if (lz == NULL) {
3897 Py_DECREF(ittuple);
3898 Py_DECREF(result);
3899 return NULL;
3900 }
3901 lz->ittuple = ittuple;
3902 lz->tuplesize = tuplesize;
3903 lz->numactive = tuplesize;
3904 lz->result = result;
3905 Py_INCREF(fillvalue);
3906 lz->fillvalue = fillvalue;
3907 return (PyObject *)lz;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003908}
3909
3910static void
3911izip_longest_dealloc(iziplongestobject *lz)
3912{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003913 PyObject_GC_UnTrack(lz);
3914 Py_XDECREF(lz->ittuple);
3915 Py_XDECREF(lz->result);
3916 Py_XDECREF(lz->fillvalue);
3917 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003918}
3919
3920static int
3921izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3922{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003923 Py_VISIT(lz->ittuple);
3924 Py_VISIT(lz->result);
3925 Py_VISIT(lz->fillvalue);
3926 return 0;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003927}
3928
3929static PyObject *
3930izip_longest_next(iziplongestobject *lz)
3931{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003932 Py_ssize_t i;
3933 Py_ssize_t tuplesize = lz->tuplesize;
3934 PyObject *result = lz->result;
3935 PyObject *it;
3936 PyObject *item;
3937 PyObject *olditem;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003938
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003939 if (tuplesize == 0)
3940 return NULL;
3941 if (lz->numactive == 0)
3942 return NULL;
3943 if (Py_REFCNT(result) == 1) {
3944 Py_INCREF(result);
3945 for (i=0 ; i < tuplesize ; i++) {
3946 it = PyTuple_GET_ITEM(lz->ittuple, i);
3947 if (it == NULL) {
3948 Py_INCREF(lz->fillvalue);
3949 item = lz->fillvalue;
3950 } else {
3951 item = PyIter_Next(it);
3952 if (item == NULL) {
3953 lz->numactive -= 1;
3954 if (lz->numactive == 0 || PyErr_Occurred()) {
3955 lz->numactive = 0;
3956 Py_DECREF(result);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003957 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003958 } else {
3959 Py_INCREF(lz->fillvalue);
3960 item = lz->fillvalue;
3961 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3962 Py_DECREF(it);
3963 }
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003964 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003965 }
3966 olditem = PyTuple_GET_ITEM(result, i);
3967 PyTuple_SET_ITEM(result, i, item);
3968 Py_DECREF(olditem);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003969 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003970 } else {
3971 result = PyTuple_New(tuplesize);
3972 if (result == NULL)
3973 return NULL;
3974 for (i=0 ; i < tuplesize ; i++) {
3975 it = PyTuple_GET_ITEM(lz->ittuple, i);
3976 if (it == NULL) {
3977 Py_INCREF(lz->fillvalue);
3978 item = lz->fillvalue;
3979 } else {
3980 item = PyIter_Next(it);
3981 if (item == NULL) {
3982 lz->numactive -= 1;
3983 if (lz->numactive == 0 || PyErr_Occurred()) {
3984 lz->numactive = 0;
3985 Py_DECREF(result);
3986 return NULL;
3987 } else {
3988 Py_INCREF(lz->fillvalue);
3989 item = lz->fillvalue;
3990 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3991 Py_DECREF(it);
3992 }
3993 }
3994 }
3995 PyTuple_SET_ITEM(result, i, item);
3996 }
3997 }
3998 return result;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003999}
4000
4001PyDoc_STRVAR(izip_longest_doc,
4002"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
4003\n\
4004Return an izip_longest object whose .next() method returns a tuple where\n\
4005the i-th element comes from the i-th iterable argument. The .next()\n\
4006method continues until the longest iterable in the argument sequence\n\
4007is exhausted and then it raises StopIteration. When the shorter iterables\n\
4008are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4009defaults to None or can be specified by a keyword argument.\n\
4010");
4011
4012static PyTypeObject iziplongest_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004013 PyVarObject_HEAD_INIT(NULL, 0)
4014 "itertools.izip_longest", /* tp_name */
4015 sizeof(iziplongestobject), /* tp_basicsize */
4016 0, /* tp_itemsize */
4017 /* methods */
4018 (destructor)izip_longest_dealloc, /* tp_dealloc */
4019 0, /* tp_print */
4020 0, /* tp_getattr */
4021 0, /* tp_setattr */
4022 0, /* tp_compare */
4023 0, /* tp_repr */
4024 0, /* tp_as_number */
4025 0, /* tp_as_sequence */
4026 0, /* tp_as_mapping */
4027 0, /* tp_hash */
4028 0, /* tp_call */
4029 0, /* tp_str */
4030 PyObject_GenericGetAttr, /* tp_getattro */
4031 0, /* tp_setattro */
4032 0, /* tp_as_buffer */
4033 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4034 Py_TPFLAGS_BASETYPE, /* tp_flags */
4035 izip_longest_doc, /* tp_doc */
4036 (traverseproc)izip_longest_traverse, /* tp_traverse */
4037 0, /* tp_clear */
4038 0, /* tp_richcompare */
4039 0, /* tp_weaklistoffset */
4040 PyObject_SelfIter, /* tp_iter */
4041 (iternextfunc)izip_longest_next, /* tp_iternext */
4042 0, /* tp_methods */
4043 0, /* tp_members */
4044 0, /* tp_getset */
4045 0, /* tp_base */
4046 0, /* tp_dict */
4047 0, /* tp_descr_get */
4048 0, /* tp_descr_set */
4049 0, /* tp_dictoffset */
4050 0, /* tp_init */
4051 0, /* tp_alloc */
4052 izip_longest_new, /* tp_new */
4053 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd36862c2007-02-21 05:20:38 +00004054};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004055
4056/* module level code ********************************************************/
4057
4058PyDoc_STRVAR(module_doc,
4059"Functional tools for creating and using iterators.\n\
4060\n\
4061Infinite iterators:\n\
4062count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004063cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004064repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004065\n\
4066Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004067chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettingere76816b2009-01-29 03:43:44 +00004068compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004069dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4070groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00004071ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
4072ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004073islice(seq, [start,] stop [, step]) --> elements from\n\
4074 seq[start:stop:step]\n\
4075imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
4076starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004077tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004078takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004079izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4080izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4081\n\
4082Combinatoric generators:\n\
4083product(p, q, ... [repeat=1]) --> cartesian product\n\
4084permutations(p[, r])\n\
Raymond Hettinger9eac1192009-11-19 01:22:04 +00004085combinations(p, r)\n\
4086combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004087");
4088
4089
Raymond Hettingerad983e72003-11-12 14:32:26 +00004090static PyMethodDef module_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004091 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4092 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004093};
4094
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004095PyMODINIT_FUNC
4096inititertools(void)
4097{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004098 int i;
4099 PyObject *m;
4100 char *name;
4101 PyTypeObject *typelist[] = {
4102 &combinations_type,
4103 &cwr_type,
4104 &cycle_type,
4105 &dropwhile_type,
4106 &takewhile_type,
4107 &islice_type,
4108 &starmap_type,
4109 &imap_type,
4110 &chain_type,
4111 &compress_type,
4112 &ifilter_type,
4113 &ifilterfalse_type,
4114 &count_type,
4115 &izip_type,
4116 &iziplongest_type,
4117 &permutations_type,
4118 &product_type,
4119 &repeat_type,
4120 &groupby_type,
4121 NULL
4122 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004123
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004124 Py_TYPE(&teedataobject_type) = &PyType_Type;
4125 m = Py_InitModule3("itertools", module_methods, module_doc);
4126 if (m == NULL)
4127 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004128
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004129 for (i=0 ; typelist[i] != NULL ; i++) {
4130 if (PyType_Ready(typelist[i]) < 0)
4131 return;
4132 name = strchr(typelist[i]->tp_name, '.');
4133 assert (name != NULL);
4134 Py_INCREF(typelist[i]);
4135 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4136 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004137
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004138 if (PyType_Ready(&teedataobject_type) < 0)
4139 return;
4140 if (PyType_Ready(&tee_type) < 0)
4141 return;
4142 if (PyType_Ready(&_grouper_type) < 0)
4143 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004144}