blob: 73c5e30db01d4711b7cf91c0057de37fed0271c1 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
2#include "Python.h"
Fred Drake08ebfec2004-10-17 19:36:57 +00003#include "structmember.h"
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004
Antoine Pitrouc83ea132010-05-09 14:46:46 +00005/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00006 by Raymond D. Hettinger <python@rcn.com>
7 Copyright (c) 2003 Python Software Foundation.
8 All rights reserved.
9*/
10
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000011
12/* groupby object ***********************************************************/
13
14typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +000015 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000021} groupbyobject;
22
23static PyTypeObject groupby_type;
24static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26static PyObject *
27groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28{
Antoine Pitrouc83ea132010-05-09 14:46:46 +000029 static char *kwargs[] = {"iterable", "key", NULL};
30 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000032
Antoine Pitrouc83ea132010-05-09 14:46:46 +000033 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
34 &it, &keyfunc))
35 return NULL;
36
37 gbo = (groupbyobject *)type->tp_alloc(type, 0);
38 if (gbo == NULL)
39 return NULL;
40 gbo->tgtkey = NULL;
41 gbo->currkey = NULL;
42 gbo->currvalue = NULL;
43 gbo->keyfunc = keyfunc;
44 Py_INCREF(keyfunc);
45 gbo->it = PyObject_GetIter(it);
46 if (gbo->it == NULL) {
47 Py_DECREF(gbo);
48 return NULL;
49 }
50 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000051}
52
53static void
54groupby_dealloc(groupbyobject *gbo)
55{
Antoine Pitrouc83ea132010-05-09 14:46:46 +000056 PyObject_GC_UnTrack(gbo);
Serhiy Storchakabb845652013-04-06 22:04:10 +030057 Py_TRASHCAN_SAFE_BEGIN(gbo)
Antoine Pitrouc83ea132010-05-09 14:46:46 +000058 Py_XDECREF(gbo->it);
59 Py_XDECREF(gbo->keyfunc);
60 Py_XDECREF(gbo->tgtkey);
61 Py_XDECREF(gbo->currkey);
62 Py_XDECREF(gbo->currvalue);
63 Py_TYPE(gbo)->tp_free(gbo);
Serhiy Storchakabb845652013-04-06 22:04:10 +030064 Py_TRASHCAN_SAFE_END(gbo)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000065}
66
67static int
68groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
69{
Antoine Pitrouc83ea132010-05-09 14:46:46 +000070 Py_VISIT(gbo->it);
71 Py_VISIT(gbo->keyfunc);
72 Py_VISIT(gbo->tgtkey);
73 Py_VISIT(gbo->currkey);
74 Py_VISIT(gbo->currvalue);
75 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000076}
77
78static PyObject *
79groupby_next(groupbyobject *gbo)
80{
Antoine Pitrouc83ea132010-05-09 14:46:46 +000081 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000082
Antoine Pitrouc83ea132010-05-09 14:46:46 +000083 /* skip to next iteration group */
84 for (;;) {
85 if (gbo->currkey == NULL)
86 /* pass */;
87 else if (gbo->tgtkey == NULL)
88 break;
89 else {
90 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000091
Antoine Pitrouc83ea132010-05-09 14:46:46 +000092 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
93 gbo->currkey, Py_EQ);
94 if (rcmp == -1)
95 return NULL;
96 else if (rcmp == 0)
97 break;
98 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000099
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000100 newvalue = PyIter_Next(gbo->it);
101 if (newvalue == NULL)
102 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000103
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000104 if (gbo->keyfunc == Py_None) {
105 newkey = newvalue;
106 Py_INCREF(newvalue);
107 } else {
108 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
109 newvalue, NULL);
110 if (newkey == NULL) {
111 Py_DECREF(newvalue);
112 return NULL;
113 }
114 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000115
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000116 tmp = gbo->currkey;
117 gbo->currkey = newkey;
118 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000119
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000120 tmp = gbo->currvalue;
121 gbo->currvalue = newvalue;
122 Py_XDECREF(tmp);
123 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000124
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000125 Py_INCREF(gbo->currkey);
126 tmp = gbo->tgtkey;
127 gbo->tgtkey = gbo->currkey;
128 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000129
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000130 grouper = _grouper_create(gbo, gbo->tgtkey);
131 if (grouper == NULL)
132 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000133
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000134 r = PyTuple_Pack(2, gbo->currkey, grouper);
135 Py_DECREF(grouper);
136 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000137}
138
139PyDoc_STRVAR(groupby_doc,
140"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
141(key, sub-iterator) grouped by each value of key(value).\n");
142
143static PyTypeObject groupby_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000144 PyVarObject_HEAD_INIT(NULL, 0)
145 "itertools.groupby", /* tp_name */
146 sizeof(groupbyobject), /* tp_basicsize */
147 0, /* tp_itemsize */
148 /* methods */
149 (destructor)groupby_dealloc, /* tp_dealloc */
150 0, /* tp_print */
151 0, /* tp_getattr */
152 0, /* tp_setattr */
153 0, /* tp_compare */
154 0, /* tp_repr */
155 0, /* tp_as_number */
156 0, /* tp_as_sequence */
157 0, /* tp_as_mapping */
158 0, /* tp_hash */
159 0, /* tp_call */
160 0, /* tp_str */
161 PyObject_GenericGetAttr, /* tp_getattro */
162 0, /* tp_setattro */
163 0, /* tp_as_buffer */
164 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
165 Py_TPFLAGS_BASETYPE, /* tp_flags */
166 groupby_doc, /* tp_doc */
167 (traverseproc)groupby_traverse, /* tp_traverse */
168 0, /* tp_clear */
169 0, /* tp_richcompare */
170 0, /* tp_weaklistoffset */
171 PyObject_SelfIter, /* tp_iter */
172 (iternextfunc)groupby_next, /* tp_iternext */
173 0, /* tp_methods */
174 0, /* tp_members */
175 0, /* tp_getset */
176 0, /* tp_base */
177 0, /* tp_dict */
178 0, /* tp_descr_get */
179 0, /* tp_descr_set */
180 0, /* tp_dictoffset */
181 0, /* tp_init */
182 0, /* tp_alloc */
183 groupby_new, /* tp_new */
184 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000185};
186
187
188/* _grouper object (internal) ************************************************/
189
190typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000191 PyObject_HEAD
192 PyObject *parent;
193 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000194} _grouperobject;
195
196static PyTypeObject _grouper_type;
197
198static PyObject *
199_grouper_create(groupbyobject *parent, PyObject *tgtkey)
200{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000201 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000202
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000203 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
204 if (igo == NULL)
205 return NULL;
206 igo->parent = (PyObject *)parent;
207 Py_INCREF(parent);
208 igo->tgtkey = tgtkey;
209 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000210
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000211 PyObject_GC_Track(igo);
212 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000213}
214
215static void
216_grouper_dealloc(_grouperobject *igo)
217{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000218 PyObject_GC_UnTrack(igo);
219 Py_DECREF(igo->parent);
220 Py_DECREF(igo->tgtkey);
221 PyObject_GC_Del(igo);
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000222}
223
224static int
225_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
226{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000227 Py_VISIT(igo->parent);
228 Py_VISIT(igo->tgtkey);
229 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000230}
231
232static PyObject *
233_grouper_next(_grouperobject *igo)
234{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000235 groupbyobject *gbo = (groupbyobject *)igo->parent;
236 PyObject *newvalue, *newkey, *r;
237 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000238
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000239 if (gbo->currvalue == NULL) {
240 newvalue = PyIter_Next(gbo->it);
241 if (newvalue == NULL)
242 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000243
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000244 if (gbo->keyfunc == Py_None) {
245 newkey = newvalue;
246 Py_INCREF(newvalue);
247 } else {
248 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
249 newvalue, NULL);
250 if (newkey == NULL) {
251 Py_DECREF(newvalue);
252 return NULL;
253 }
254 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000255
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000256 assert(gbo->currkey == NULL);
257 gbo->currkey = newkey;
258 gbo->currvalue = newvalue;
259 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000260
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000261 assert(gbo->currkey != NULL);
262 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
263 if (rcmp <= 0)
264 /* got any error or current group is end */
265 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000266
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000267 r = gbo->currvalue;
268 gbo->currvalue = NULL;
269 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000270
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000271 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000272}
273
274static PyTypeObject _grouper_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000275 PyVarObject_HEAD_INIT(NULL, 0)
276 "itertools._grouper", /* tp_name */
277 sizeof(_grouperobject), /* tp_basicsize */
278 0, /* tp_itemsize */
279 /* methods */
280 (destructor)_grouper_dealloc, /* tp_dealloc */
281 0, /* tp_print */
282 0, /* tp_getattr */
283 0, /* tp_setattr */
284 0, /* tp_compare */
285 0, /* tp_repr */
286 0, /* tp_as_number */
287 0, /* tp_as_sequence */
288 0, /* tp_as_mapping */
289 0, /* tp_hash */
290 0, /* tp_call */
291 0, /* tp_str */
292 PyObject_GenericGetAttr, /* tp_getattro */
293 0, /* tp_setattro */
294 0, /* tp_as_buffer */
295 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
296 0, /* tp_doc */
297 (traverseproc)_grouper_traverse,/* tp_traverse */
298 0, /* tp_clear */
299 0, /* tp_richcompare */
300 0, /* tp_weaklistoffset */
301 PyObject_SelfIter, /* tp_iter */
302 (iternextfunc)_grouper_next, /* tp_iternext */
303 0, /* tp_methods */
304 0, /* tp_members */
305 0, /* tp_getset */
306 0, /* tp_base */
307 0, /* tp_dict */
308 0, /* tp_descr_get */
309 0, /* tp_descr_set */
310 0, /* tp_dictoffset */
311 0, /* tp_init */
312 0, /* tp_alloc */
313 0, /* tp_new */
314 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000315};
316
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000317
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000318
Raymond Hettingerad983e72003-11-12 14:32:26 +0000319/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000320
Raymond Hettingerad983e72003-11-12 14:32:26 +0000321/* The teedataobject pre-allocates space for LINKCELLS number of objects.
322 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000323 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000324 the other structure members including PyHEAD overhead. The larger the
325 value, the less memory overhead per object and the less time spent
326 allocating/deallocating new links. The smaller the number, the less
327 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000328*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000329#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000330
331typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000332 PyObject_HEAD
333 PyObject *it;
334 int numread;
335 PyObject *nextlink;
336 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000337} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000338
339typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000340 PyObject_HEAD
341 teedataobject *dataobj;
342 int index;
343 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000344} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000345
Raymond Hettingerad983e72003-11-12 14:32:26 +0000346static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000347
348static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000349teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000350{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000351 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000352
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000353 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
354 if (tdo == NULL)
355 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000356
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000357 tdo->numread = 0;
358 tdo->nextlink = NULL;
359 Py_INCREF(it);
360 tdo->it = it;
361 PyObject_GC_Track(tdo);
362 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000363}
364
365static PyObject *
366teedataobject_jumplink(teedataobject *tdo)
367{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000368 if (tdo->nextlink == NULL)
369 tdo->nextlink = teedataobject_new(tdo->it);
370 Py_XINCREF(tdo->nextlink);
371 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000372}
373
374static PyObject *
375teedataobject_getitem(teedataobject *tdo, int i)
376{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000377 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000378
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000379 assert(i < LINKCELLS);
380 if (i < tdo->numread)
381 value = tdo->values[i];
382 else {
383 /* this is the lead iterator, so fetch more data */
384 assert(i == tdo->numread);
385 value = PyIter_Next(tdo->it);
386 if (value == NULL)
387 return NULL;
388 tdo->numread++;
389 tdo->values[i] = value;
390 }
391 Py_INCREF(value);
392 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000393}
394
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000395static int
396teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
397{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000398 int i;
399 Py_VISIT(tdo->it);
400 for (i = 0; i < tdo->numread; i++)
401 Py_VISIT(tdo->values[i]);
402 Py_VISIT(tdo->nextlink);
403 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000404}
405
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200406static void
407teedataobject_safe_decref(PyObject *obj)
408{
409 while (obj && Py_TYPE(obj) == &teedataobject_type &&
410 Py_REFCNT(obj) == 1) {
411 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
412 ((teedataobject *)obj)->nextlink = NULL;
413 Py_DECREF(obj);
414 obj = nextlink;
415 }
416 Py_XDECREF(obj);
417}
418
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000419static int
420teedataobject_clear(teedataobject *tdo)
421{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000422 int i;
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200423 PyObject *tmp;
424
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000425 Py_CLEAR(tdo->it);
426 for (i=0 ; i<tdo->numread ; i++)
427 Py_CLEAR(tdo->values[i]);
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200428 tmp = tdo->nextlink;
429 tdo->nextlink = NULL;
430 teedataobject_safe_decref(tmp);
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000431 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000432}
433
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000434static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000435teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000436{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000437 PyObject_GC_UnTrack(tdo);
438 teedataobject_clear(tdo);
439 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000440}
441
Raymond Hettingerad983e72003-11-12 14:32:26 +0000442PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000443
Raymond Hettingerad983e72003-11-12 14:32:26 +0000444static PyTypeObject teedataobject_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000445 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
446 "itertools.tee_dataobject", /* tp_name */
447 sizeof(teedataobject), /* tp_basicsize */
448 0, /* tp_itemsize */
449 /* methods */
450 (destructor)teedataobject_dealloc, /* tp_dealloc */
451 0, /* tp_print */
452 0, /* tp_getattr */
453 0, /* tp_setattr */
454 0, /* tp_compare */
455 0, /* tp_repr */
456 0, /* tp_as_number */
457 0, /* tp_as_sequence */
458 0, /* tp_as_mapping */
459 0, /* tp_hash */
460 0, /* tp_call */
461 0, /* tp_str */
462 PyObject_GenericGetAttr, /* tp_getattro */
463 0, /* tp_setattro */
464 0, /* tp_as_buffer */
465 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
466 teedataobject_doc, /* tp_doc */
467 (traverseproc)teedataobject_traverse, /* tp_traverse */
468 (inquiry)teedataobject_clear, /* tp_clear */
469 0, /* tp_richcompare */
470 0, /* tp_weaklistoffset */
471 0, /* tp_iter */
472 0, /* tp_iternext */
473 0, /* tp_methods */
474 0, /* tp_members */
475 0, /* tp_getset */
476 0, /* tp_base */
477 0, /* tp_dict */
478 0, /* tp_descr_get */
479 0, /* tp_descr_set */
480 0, /* tp_dictoffset */
481 0, /* tp_init */
482 0, /* tp_alloc */
483 0, /* tp_new */
484 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000485};
486
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000487
488static PyTypeObject tee_type;
489
490static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000491tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000492{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000493 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000494
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000495 if (to->index >= LINKCELLS) {
496 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200497 if (link == NULL)
498 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000499 Py_DECREF(to->dataobj);
500 to->dataobj = (teedataobject *)link;
501 to->index = 0;
502 }
503 value = teedataobject_getitem(to->dataobj, to->index);
504 if (value == NULL)
505 return NULL;
506 to->index++;
507 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000508}
509
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000510static int
511tee_traverse(teeobject *to, visitproc visit, void *arg)
512{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000513 Py_VISIT((PyObject *)to->dataobj);
514 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000515}
516
Raymond Hettingerad983e72003-11-12 14:32:26 +0000517static PyObject *
518tee_copy(teeobject *to)
519{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000520 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000521
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000522 newto = PyObject_GC_New(teeobject, &tee_type);
523 if (newto == NULL)
524 return NULL;
525 Py_INCREF(to->dataobj);
526 newto->dataobj = to->dataobj;
527 newto->index = to->index;
528 newto->weakreflist = NULL;
529 PyObject_GC_Track(newto);
530 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000531}
532
533PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
534
535static PyObject *
536tee_fromiterable(PyObject *iterable)
537{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000538 teeobject *to;
539 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000540
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000541 it = PyObject_GetIter(iterable);
542 if (it == NULL)
543 return NULL;
544 if (PyObject_TypeCheck(it, &tee_type)) {
545 to = (teeobject *)tee_copy((teeobject *)it);
546 goto done;
547 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000548
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000549 to = PyObject_GC_New(teeobject, &tee_type);
550 if (to == NULL)
551 goto done;
552 to->dataobj = (teedataobject *)teedataobject_new(it);
553 if (!to->dataobj) {
554 PyObject_GC_Del(to);
555 to = NULL;
556 goto done;
557 }
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000558
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000559 to->index = 0;
560 to->weakreflist = NULL;
561 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000562done:
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000563 Py_XDECREF(it);
564 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000565}
566
567static PyObject *
568tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
569{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000570 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000571
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000572 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
573 return NULL;
574 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000575}
576
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000577static int
578tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000579{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000580 if (to->weakreflist != NULL)
581 PyObject_ClearWeakRefs((PyObject *) to);
582 Py_CLEAR(to->dataobj);
583 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000584}
585
586static void
587tee_dealloc(teeobject *to)
588{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000589 PyObject_GC_UnTrack(to);
590 tee_clear(to);
591 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000592}
593
Raymond Hettingerad983e72003-11-12 14:32:26 +0000594PyDoc_STRVAR(teeobject_doc,
595"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000596
Raymond Hettingerad983e72003-11-12 14:32:26 +0000597static PyMethodDef tee_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000598 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
599 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000600};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000601
602static PyTypeObject tee_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000603 PyVarObject_HEAD_INIT(NULL, 0)
604 "itertools.tee", /* tp_name */
605 sizeof(teeobject), /* tp_basicsize */
606 0, /* tp_itemsize */
607 /* methods */
608 (destructor)tee_dealloc, /* tp_dealloc */
609 0, /* tp_print */
610 0, /* tp_getattr */
611 0, /* tp_setattr */
612 0, /* tp_compare */
613 0, /* tp_repr */
614 0, /* tp_as_number */
615 0, /* tp_as_sequence */
616 0, /* tp_as_mapping */
617 0, /* tp_hash */
618 0, /* tp_call */
619 0, /* tp_str */
620 0, /* tp_getattro */
621 0, /* tp_setattro */
622 0, /* tp_as_buffer */
623 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
624 teeobject_doc, /* tp_doc */
625 (traverseproc)tee_traverse, /* tp_traverse */
626 (inquiry)tee_clear, /* tp_clear */
627 0, /* tp_richcompare */
628 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
629 PyObject_SelfIter, /* tp_iter */
630 (iternextfunc)tee_next, /* tp_iternext */
631 tee_methods, /* tp_methods */
632 0, /* tp_members */
633 0, /* tp_getset */
634 0, /* tp_base */
635 0, /* tp_dict */
636 0, /* tp_descr_get */
637 0, /* tp_descr_set */
638 0, /* tp_dictoffset */
639 0, /* tp_init */
640 0, /* tp_alloc */
641 tee_new, /* tp_new */
642 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000643};
644
Raymond Hettingerad983e72003-11-12 14:32:26 +0000645static PyObject *
646tee(PyObject *self, PyObject *args)
647{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000648 Py_ssize_t i, n=2;
649 PyObject *it, *iterable, *copyable, *result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000650
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000651 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
652 return NULL;
653 if (n < 0) {
654 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
655 return NULL;
656 }
657 result = PyTuple_New(n);
658 if (result == NULL)
659 return NULL;
660 if (n == 0)
661 return result;
662 it = PyObject_GetIter(iterable);
663 if (it == NULL) {
664 Py_DECREF(result);
665 return NULL;
666 }
667 if (!PyObject_HasAttrString(it, "__copy__")) {
668 copyable = tee_fromiterable(it);
669 Py_DECREF(it);
670 if (copyable == NULL) {
671 Py_DECREF(result);
672 return NULL;
673 }
674 } else
675 copyable = it;
676 PyTuple_SET_ITEM(result, 0, copyable);
677 for (i=1 ; i<n ; i++) {
678 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
679 if (copyable == NULL) {
680 Py_DECREF(result);
681 return NULL;
682 }
683 PyTuple_SET_ITEM(result, i, copyable);
684 }
685 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000686}
687
688PyDoc_STRVAR(tee_doc,
689"tee(iterable, n=2) --> tuple of n independent iterators.");
690
691
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000692/* cycle object **********************************************************/
693
694typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000695 PyObject_HEAD
696 PyObject *it;
697 PyObject *saved;
698 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000699} cycleobject;
700
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000701static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000702
703static PyObject *
704cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
705{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000706 PyObject *it;
707 PyObject *iterable;
708 PyObject *saved;
709 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000710
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000711 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
712 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000713
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000714 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
715 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000716
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000717 /* Get iterator. */
718 it = PyObject_GetIter(iterable);
719 if (it == NULL)
720 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000721
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000722 saved = PyList_New(0);
723 if (saved == NULL) {
724 Py_DECREF(it);
725 return NULL;
726 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000727
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000728 /* create cycleobject structure */
729 lz = (cycleobject *)type->tp_alloc(type, 0);
730 if (lz == NULL) {
731 Py_DECREF(it);
732 Py_DECREF(saved);
733 return NULL;
734 }
735 lz->it = it;
736 lz->saved = saved;
737 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000738
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000739 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000740}
741
742static void
743cycle_dealloc(cycleobject *lz)
744{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000745 PyObject_GC_UnTrack(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +0300746 Py_TRASHCAN_SAFE_BEGIN(lz)
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000747 Py_XDECREF(lz->saved);
748 Py_XDECREF(lz->it);
749 Py_TYPE(lz)->tp_free(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +0300750 Py_TRASHCAN_SAFE_END(lz)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000751}
752
753static int
754cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
755{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000756 Py_VISIT(lz->it);
757 Py_VISIT(lz->saved);
758 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000759}
760
761static PyObject *
762cycle_next(cycleobject *lz)
763{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000764 PyObject *item;
765 PyObject *it;
766 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000767
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000768 while (1) {
769 item = PyIter_Next(lz->it);
770 if (item != NULL) {
771 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
772 Py_DECREF(item);
773 return NULL;
774 }
775 return item;
776 }
777 if (PyErr_Occurred()) {
778 if (PyErr_ExceptionMatches(PyExc_StopIteration))
779 PyErr_Clear();
780 else
781 return NULL;
782 }
783 if (PyList_Size(lz->saved) == 0)
784 return NULL;
785 it = PyObject_GetIter(lz->saved);
786 if (it == NULL)
787 return NULL;
788 tmp = lz->it;
789 lz->it = it;
790 lz->firstpass = 1;
791 Py_DECREF(tmp);
792 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000793}
794
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000795PyDoc_STRVAR(cycle_doc,
796"cycle(iterable) --> cycle object\n\
797\n\
798Return elements from the iterable until it is exhausted.\n\
799Then repeat the sequence indefinitely.");
800
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000801static PyTypeObject cycle_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000802 PyVarObject_HEAD_INIT(NULL, 0)
803 "itertools.cycle", /* tp_name */
804 sizeof(cycleobject), /* tp_basicsize */
805 0, /* tp_itemsize */
806 /* methods */
807 (destructor)cycle_dealloc, /* tp_dealloc */
808 0, /* tp_print */
809 0, /* tp_getattr */
810 0, /* tp_setattr */
811 0, /* tp_compare */
812 0, /* tp_repr */
813 0, /* tp_as_number */
814 0, /* tp_as_sequence */
815 0, /* tp_as_mapping */
816 0, /* tp_hash */
817 0, /* tp_call */
818 0, /* tp_str */
819 PyObject_GenericGetAttr, /* tp_getattro */
820 0, /* tp_setattro */
821 0, /* tp_as_buffer */
822 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
823 Py_TPFLAGS_BASETYPE, /* tp_flags */
824 cycle_doc, /* tp_doc */
825 (traverseproc)cycle_traverse, /* tp_traverse */
826 0, /* tp_clear */
827 0, /* tp_richcompare */
828 0, /* tp_weaklistoffset */
829 PyObject_SelfIter, /* tp_iter */
830 (iternextfunc)cycle_next, /* tp_iternext */
831 0, /* tp_methods */
832 0, /* tp_members */
833 0, /* tp_getset */
834 0, /* tp_base */
835 0, /* tp_dict */
836 0, /* tp_descr_get */
837 0, /* tp_descr_set */
838 0, /* tp_dictoffset */
839 0, /* tp_init */
840 0, /* tp_alloc */
841 cycle_new, /* tp_new */
842 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000843};
844
845
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000846/* dropwhile object **********************************************************/
847
848typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000849 PyObject_HEAD
850 PyObject *func;
851 PyObject *it;
852 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000853} dropwhileobject;
854
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000855static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000856
857static PyObject *
858dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
859{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000860 PyObject *func, *seq;
861 PyObject *it;
862 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000863
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000864 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
865 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000866
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000867 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
868 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000869
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000870 /* Get iterator. */
871 it = PyObject_GetIter(seq);
872 if (it == NULL)
873 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000874
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000875 /* create dropwhileobject structure */
876 lz = (dropwhileobject *)type->tp_alloc(type, 0);
877 if (lz == NULL) {
878 Py_DECREF(it);
879 return NULL;
880 }
881 Py_INCREF(func);
882 lz->func = func;
883 lz->it = it;
884 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000885
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000886 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000887}
888
889static void
890dropwhile_dealloc(dropwhileobject *lz)
891{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000892 PyObject_GC_UnTrack(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +0300893 Py_TRASHCAN_SAFE_BEGIN(lz)
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000894 Py_XDECREF(lz->func);
895 Py_XDECREF(lz->it);
896 Py_TYPE(lz)->tp_free(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +0300897 Py_TRASHCAN_SAFE_END(lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000898}
899
900static int
901dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
902{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000903 Py_VISIT(lz->it);
904 Py_VISIT(lz->func);
905 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000906}
907
908static PyObject *
909dropwhile_next(dropwhileobject *lz)
910{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000911 PyObject *item, *good;
912 PyObject *it = lz->it;
913 long ok;
914 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000915
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000916 iternext = *Py_TYPE(it)->tp_iternext;
917 for (;;) {
Serhiy Storchakabb845652013-04-06 22:04:10 +0300918 if (Py_EnterRecursiveCall(" while iterating"))
919 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000920 item = iternext(it);
Serhiy Storchakabb845652013-04-06 22:04:10 +0300921 Py_LeaveRecursiveCall();
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000922 if (item == NULL)
923 return NULL;
924 if (lz->start == 1)
925 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000926
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000927 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
928 if (good == NULL) {
929 Py_DECREF(item);
930 return NULL;
931 }
932 ok = PyObject_IsTrue(good);
933 Py_DECREF(good);
Antoine Pitrouc5bef752012-08-15 23:16:51 +0200934 if (ok == 0) {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000935 lz->start = 1;
936 return item;
937 }
938 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +0200939 if (ok < 0)
940 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000941 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000942}
943
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000944PyDoc_STRVAR(dropwhile_doc,
945"dropwhile(predicate, iterable) --> dropwhile object\n\
946\n\
947Drop items from the iterable while predicate(item) is true.\n\
948Afterwards, return every element until the iterable is exhausted.");
949
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000950static PyTypeObject dropwhile_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000951 PyVarObject_HEAD_INIT(NULL, 0)
952 "itertools.dropwhile", /* tp_name */
953 sizeof(dropwhileobject), /* tp_basicsize */
954 0, /* tp_itemsize */
955 /* methods */
956 (destructor)dropwhile_dealloc, /* tp_dealloc */
957 0, /* tp_print */
958 0, /* tp_getattr */
959 0, /* tp_setattr */
960 0, /* tp_compare */
961 0, /* tp_repr */
962 0, /* tp_as_number */
963 0, /* tp_as_sequence */
964 0, /* tp_as_mapping */
965 0, /* tp_hash */
966 0, /* tp_call */
967 0, /* tp_str */
968 PyObject_GenericGetAttr, /* tp_getattro */
969 0, /* tp_setattro */
970 0, /* tp_as_buffer */
971 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
972 Py_TPFLAGS_BASETYPE, /* tp_flags */
973 dropwhile_doc, /* tp_doc */
974 (traverseproc)dropwhile_traverse, /* tp_traverse */
975 0, /* tp_clear */
976 0, /* tp_richcompare */
977 0, /* tp_weaklistoffset */
978 PyObject_SelfIter, /* tp_iter */
979 (iternextfunc)dropwhile_next, /* tp_iternext */
980 0, /* tp_methods */
981 0, /* tp_members */
982 0, /* tp_getset */
983 0, /* tp_base */
984 0, /* tp_dict */
985 0, /* tp_descr_get */
986 0, /* tp_descr_set */
987 0, /* tp_dictoffset */
988 0, /* tp_init */
989 0, /* tp_alloc */
990 dropwhile_new, /* tp_new */
991 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000992};
993
994
995/* takewhile object **********************************************************/
996
997typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000998 PyObject_HEAD
999 PyObject *func;
1000 PyObject *it;
1001 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001002} takewhileobject;
1003
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001004static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001005
1006static PyObject *
1007takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1008{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001009 PyObject *func, *seq;
1010 PyObject *it;
1011 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001012
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001013 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1014 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001015
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001016 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1017 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001018
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001019 /* Get iterator. */
1020 it = PyObject_GetIter(seq);
1021 if (it == NULL)
1022 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001023
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001024 /* create takewhileobject structure */
1025 lz = (takewhileobject *)type->tp_alloc(type, 0);
1026 if (lz == NULL) {
1027 Py_DECREF(it);
1028 return NULL;
1029 }
1030 Py_INCREF(func);
1031 lz->func = func;
1032 lz->it = it;
1033 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001034
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001035 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001036}
1037
1038static void
1039takewhile_dealloc(takewhileobject *lz)
1040{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001041 PyObject_GC_UnTrack(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03001042 Py_TRASHCAN_SAFE_BEGIN(lz)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001043 Py_XDECREF(lz->func);
1044 Py_XDECREF(lz->it);
1045 Py_TYPE(lz)->tp_free(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03001046 Py_TRASHCAN_SAFE_END(lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001047}
1048
1049static int
1050takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1051{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001052 Py_VISIT(lz->it);
1053 Py_VISIT(lz->func);
1054 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001055}
1056
1057static PyObject *
1058takewhile_next(takewhileobject *lz)
1059{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001060 PyObject *item, *good;
1061 PyObject *it = lz->it;
1062 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001063
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001064 if (lz->stop == 1)
1065 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001066
Serhiy Storchakabb845652013-04-06 22:04:10 +03001067 if (Py_EnterRecursiveCall(" while iterating"))
1068 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001069 item = (*Py_TYPE(it)->tp_iternext)(it);
Serhiy Storchakabb845652013-04-06 22:04:10 +03001070 Py_LeaveRecursiveCall();
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001071 if (item == NULL)
1072 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001073
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001074 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1075 if (good == NULL) {
1076 Py_DECREF(item);
1077 return NULL;
1078 }
1079 ok = PyObject_IsTrue(good);
1080 Py_DECREF(good);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02001081 if (ok > 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001082 return item;
1083 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02001084 if (ok == 0)
1085 lz->stop = 1;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001086 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001087}
1088
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001089PyDoc_STRVAR(takewhile_doc,
1090"takewhile(predicate, iterable) --> takewhile object\n\
1091\n\
1092Return successive entries from an iterable as long as the \n\
1093predicate evaluates to true for each entry.");
1094
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001095static PyTypeObject takewhile_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001096 PyVarObject_HEAD_INIT(NULL, 0)
1097 "itertools.takewhile", /* tp_name */
1098 sizeof(takewhileobject), /* tp_basicsize */
1099 0, /* tp_itemsize */
1100 /* methods */
1101 (destructor)takewhile_dealloc, /* tp_dealloc */
1102 0, /* tp_print */
1103 0, /* tp_getattr */
1104 0, /* tp_setattr */
1105 0, /* tp_compare */
1106 0, /* tp_repr */
1107 0, /* tp_as_number */
1108 0, /* tp_as_sequence */
1109 0, /* tp_as_mapping */
1110 0, /* tp_hash */
1111 0, /* tp_call */
1112 0, /* tp_str */
1113 PyObject_GenericGetAttr, /* tp_getattro */
1114 0, /* tp_setattro */
1115 0, /* tp_as_buffer */
1116 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1117 Py_TPFLAGS_BASETYPE, /* tp_flags */
1118 takewhile_doc, /* tp_doc */
1119 (traverseproc)takewhile_traverse, /* tp_traverse */
1120 0, /* tp_clear */
1121 0, /* tp_richcompare */
1122 0, /* tp_weaklistoffset */
1123 PyObject_SelfIter, /* tp_iter */
1124 (iternextfunc)takewhile_next, /* tp_iternext */
1125 0, /* tp_methods */
1126 0, /* tp_members */
1127 0, /* tp_getset */
1128 0, /* tp_base */
1129 0, /* tp_dict */
1130 0, /* tp_descr_get */
1131 0, /* tp_descr_set */
1132 0, /* tp_dictoffset */
1133 0, /* tp_init */
1134 0, /* tp_alloc */
1135 takewhile_new, /* tp_new */
1136 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001137};
1138
1139
1140/* islice object ************************************************************/
1141
1142typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001143 PyObject_HEAD
1144 PyObject *it;
1145 Py_ssize_t next;
1146 Py_ssize_t stop;
1147 Py_ssize_t step;
1148 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001149} isliceobject;
1150
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001151static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001152
1153static PyObject *
1154islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1155{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001156 PyObject *seq;
1157 Py_ssize_t start=0, stop=-1, step=1;
1158 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1159 Py_ssize_t numargs;
1160 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001161
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001162 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1163 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001164
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001165 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1166 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001167
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001168 numargs = PyTuple_Size(args);
1169 if (numargs == 2) {
1170 if (a1 != Py_None) {
1171 stop = PyInt_AsSsize_t(a1);
1172 if (stop == -1) {
1173 if (PyErr_Occurred())
1174 PyErr_Clear();
1175 PyErr_SetString(PyExc_ValueError,
1176 "Stop argument for islice() must be None or an integer: 0 <= x <= maxint.");
1177 return NULL;
1178 }
1179 }
1180 } else {
1181 if (a1 != Py_None)
1182 start = PyInt_AsSsize_t(a1);
1183 if (start == -1 && PyErr_Occurred())
1184 PyErr_Clear();
1185 if (a2 != Py_None) {
1186 stop = PyInt_AsSsize_t(a2);
1187 if (stop == -1) {
1188 if (PyErr_Occurred())
1189 PyErr_Clear();
1190 PyErr_SetString(PyExc_ValueError,
1191 "Stop argument for islice() must be None or an integer: 0 <= x <= maxint.");
1192 return NULL;
1193 }
1194 }
1195 }
1196 if (start<0 || stop<-1) {
1197 PyErr_SetString(PyExc_ValueError,
1198 "Indices for islice() must be None or an integer: 0 <= x <= maxint.");
1199 return NULL;
1200 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001201
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001202 if (a3 != NULL) {
1203 if (a3 != Py_None)
1204 step = PyInt_AsSsize_t(a3);
1205 if (step == -1 && PyErr_Occurred())
1206 PyErr_Clear();
1207 }
1208 if (step<1) {
1209 PyErr_SetString(PyExc_ValueError,
1210 "Step for islice() must be a positive integer or None.");
1211 return NULL;
1212 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001213
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001214 /* Get iterator. */
1215 it = PyObject_GetIter(seq);
1216 if (it == NULL)
1217 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001218
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001219 /* create isliceobject structure */
1220 lz = (isliceobject *)type->tp_alloc(type, 0);
1221 if (lz == NULL) {
1222 Py_DECREF(it);
1223 return NULL;
1224 }
1225 lz->it = it;
1226 lz->next = start;
1227 lz->stop = stop;
1228 lz->step = step;
1229 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001230
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001231 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001232}
1233
1234static void
1235islice_dealloc(isliceobject *lz)
1236{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001237 PyObject_GC_UnTrack(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03001238 Py_TRASHCAN_SAFE_BEGIN(lz)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001239 Py_XDECREF(lz->it);
1240 Py_TYPE(lz)->tp_free(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03001241 Py_TRASHCAN_SAFE_END(lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001242}
1243
1244static int
1245islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1246{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001247 Py_VISIT(lz->it);
1248 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001249}
1250
1251static PyObject *
1252islice_next(isliceobject *lz)
1253{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001254 PyObject *item;
1255 PyObject *it = lz->it;
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001256 Py_ssize_t stop = lz->stop;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001257 Py_ssize_t oldnext;
1258 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001259
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001260 iternext = *Py_TYPE(it)->tp_iternext;
1261 while (lz->cnt < lz->next) {
Serhiy Storchakabb845652013-04-06 22:04:10 +03001262 if (Py_EnterRecursiveCall(" while iterating"))
1263 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001264 item = iternext(it);
Serhiy Storchakabb845652013-04-06 22:04:10 +03001265 Py_LeaveRecursiveCall();
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001266 if (item == NULL)
1267 return NULL;
1268 Py_DECREF(item);
1269 lz->cnt++;
1270 }
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001271 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001272 return NULL;
Serhiy Storchakabb845652013-04-06 22:04:10 +03001273 if (Py_EnterRecursiveCall(" while iterating"))
1274 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001275 item = iternext(it);
Serhiy Storchakabb845652013-04-06 22:04:10 +03001276 Py_LeaveRecursiveCall();
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001277 if (item == NULL)
1278 return NULL;
1279 lz->cnt++;
1280 oldnext = lz->next;
Mark Dickinsona96b0d12011-09-24 09:01:16 +01001281 /* The (size_t) cast below avoids the danger of undefined
1282 behaviour from signed integer overflow. */
1283 lz->next += (size_t)lz->step;
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001284 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1285 lz->next = stop;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001286 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001287}
1288
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001289PyDoc_STRVAR(islice_doc,
1290"islice(iterable, [start,] stop [, step]) --> islice object\n\
1291\n\
1292Return an iterator whose next() method returns selected values from an\n\
1293iterable. If start is specified, will skip all preceding elements;\n\
1294otherwise, start defaults to zero. Step defaults to one. If\n\
1295specified as another value, step determines how many values are \n\
1296skipped between successive calls. Works like a slice() on a list\n\
1297but returns an iterator.");
1298
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001299static PyTypeObject islice_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001300 PyVarObject_HEAD_INIT(NULL, 0)
1301 "itertools.islice", /* tp_name */
1302 sizeof(isliceobject), /* tp_basicsize */
1303 0, /* tp_itemsize */
1304 /* methods */
1305 (destructor)islice_dealloc, /* tp_dealloc */
1306 0, /* tp_print */
1307 0, /* tp_getattr */
1308 0, /* tp_setattr */
1309 0, /* tp_compare */
1310 0, /* tp_repr */
1311 0, /* tp_as_number */
1312 0, /* tp_as_sequence */
1313 0, /* tp_as_mapping */
1314 0, /* tp_hash */
1315 0, /* tp_call */
1316 0, /* tp_str */
1317 PyObject_GenericGetAttr, /* tp_getattro */
1318 0, /* tp_setattro */
1319 0, /* tp_as_buffer */
1320 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1321 Py_TPFLAGS_BASETYPE, /* tp_flags */
1322 islice_doc, /* tp_doc */
1323 (traverseproc)islice_traverse, /* tp_traverse */
1324 0, /* tp_clear */
1325 0, /* tp_richcompare */
1326 0, /* tp_weaklistoffset */
1327 PyObject_SelfIter, /* tp_iter */
1328 (iternextfunc)islice_next, /* tp_iternext */
1329 0, /* tp_methods */
1330 0, /* tp_members */
1331 0, /* tp_getset */
1332 0, /* tp_base */
1333 0, /* tp_dict */
1334 0, /* tp_descr_get */
1335 0, /* tp_descr_set */
1336 0, /* tp_dictoffset */
1337 0, /* tp_init */
1338 0, /* tp_alloc */
1339 islice_new, /* tp_new */
1340 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001341};
1342
1343
1344/* starmap object ************************************************************/
1345
1346typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001347 PyObject_HEAD
1348 PyObject *func;
1349 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001350} starmapobject;
1351
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001352static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001353
1354static PyObject *
1355starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1356{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001357 PyObject *func, *seq;
1358 PyObject *it;
1359 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001360
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001361 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1362 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001363
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001364 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1365 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001366
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001367 /* Get iterator. */
1368 it = PyObject_GetIter(seq);
1369 if (it == NULL)
1370 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001371
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001372 /* create starmapobject structure */
1373 lz = (starmapobject *)type->tp_alloc(type, 0);
1374 if (lz == NULL) {
1375 Py_DECREF(it);
1376 return NULL;
1377 }
1378 Py_INCREF(func);
1379 lz->func = func;
1380 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001381
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001382 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001383}
1384
1385static void
1386starmap_dealloc(starmapobject *lz)
1387{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001388 PyObject_GC_UnTrack(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03001389 Py_TRASHCAN_SAFE_BEGIN(lz)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001390 Py_XDECREF(lz->func);
1391 Py_XDECREF(lz->it);
1392 Py_TYPE(lz)->tp_free(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03001393 Py_TRASHCAN_SAFE_END(lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001394}
1395
1396static int
1397starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1398{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001399 Py_VISIT(lz->it);
1400 Py_VISIT(lz->func);
1401 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001402}
1403
1404static PyObject *
1405starmap_next(starmapobject *lz)
1406{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001407 PyObject *args;
1408 PyObject *result;
1409 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001410
Serhiy Storchakabb845652013-04-06 22:04:10 +03001411 if (Py_EnterRecursiveCall(" while iterating"))
1412 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001413 args = (*Py_TYPE(it)->tp_iternext)(it);
Serhiy Storchakabb845652013-04-06 22:04:10 +03001414 Py_LeaveRecursiveCall();
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001415 if (args == NULL)
1416 return NULL;
1417 if (!PyTuple_CheckExact(args)) {
1418 PyObject *newargs = PySequence_Tuple(args);
1419 Py_DECREF(args);
1420 if (newargs == NULL)
1421 return NULL;
1422 args = newargs;
1423 }
1424 result = PyObject_Call(lz->func, args, NULL);
1425 Py_DECREF(args);
1426 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001427}
1428
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001429PyDoc_STRVAR(starmap_doc,
1430"starmap(function, sequence) --> starmap object\n\
1431\n\
1432Return an iterator whose values are returned from the function evaluated\n\
1433with a argument tuple taken from the given sequence.");
1434
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001435static PyTypeObject starmap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001436 PyVarObject_HEAD_INIT(NULL, 0)
1437 "itertools.starmap", /* tp_name */
1438 sizeof(starmapobject), /* tp_basicsize */
1439 0, /* tp_itemsize */
1440 /* methods */
1441 (destructor)starmap_dealloc, /* tp_dealloc */
1442 0, /* tp_print */
1443 0, /* tp_getattr */
1444 0, /* tp_setattr */
1445 0, /* tp_compare */
1446 0, /* tp_repr */
1447 0, /* tp_as_number */
1448 0, /* tp_as_sequence */
1449 0, /* tp_as_mapping */
1450 0, /* tp_hash */
1451 0, /* tp_call */
1452 0, /* tp_str */
1453 PyObject_GenericGetAttr, /* tp_getattro */
1454 0, /* tp_setattro */
1455 0, /* tp_as_buffer */
1456 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1457 Py_TPFLAGS_BASETYPE, /* tp_flags */
1458 starmap_doc, /* tp_doc */
1459 (traverseproc)starmap_traverse, /* tp_traverse */
1460 0, /* tp_clear */
1461 0, /* tp_richcompare */
1462 0, /* tp_weaklistoffset */
1463 PyObject_SelfIter, /* tp_iter */
1464 (iternextfunc)starmap_next, /* tp_iternext */
1465 0, /* tp_methods */
1466 0, /* tp_members */
1467 0, /* tp_getset */
1468 0, /* tp_base */
1469 0, /* tp_dict */
1470 0, /* tp_descr_get */
1471 0, /* tp_descr_set */
1472 0, /* tp_dictoffset */
1473 0, /* tp_init */
1474 0, /* tp_alloc */
1475 starmap_new, /* tp_new */
1476 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001477};
1478
1479
1480/* imap object ************************************************************/
1481
1482typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001483 PyObject_HEAD
1484 PyObject *iters;
1485 PyObject *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001486} imapobject;
1487
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001488static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001489
1490static PyObject *
1491imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1492{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001493 PyObject *it, *iters, *func;
1494 imapobject *lz;
1495 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001496
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001497 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
1498 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001499
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001500 numargs = PyTuple_Size(args);
1501 if (numargs < 2) {
1502 PyErr_SetString(PyExc_TypeError,
1503 "imap() must have at least two arguments.");
1504 return NULL;
1505 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001506
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001507 iters = PyTuple_New(numargs-1);
1508 if (iters == NULL)
1509 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001510
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001511 for (i=1 ; i<numargs ; i++) {
1512 /* Get iterator. */
1513 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1514 if (it == NULL) {
1515 Py_DECREF(iters);
1516 return NULL;
1517 }
1518 PyTuple_SET_ITEM(iters, i-1, it);
1519 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001520
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001521 /* create imapobject structure */
1522 lz = (imapobject *)type->tp_alloc(type, 0);
1523 if (lz == NULL) {
1524 Py_DECREF(iters);
1525 return NULL;
1526 }
1527 lz->iters = iters;
1528 func = PyTuple_GET_ITEM(args, 0);
1529 Py_INCREF(func);
1530 lz->func = func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001531
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001532 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001533}
1534
1535static void
1536imap_dealloc(imapobject *lz)
1537{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001538 PyObject_GC_UnTrack(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03001539 Py_TRASHCAN_SAFE_BEGIN(lz)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001540 Py_XDECREF(lz->iters);
1541 Py_XDECREF(lz->func);
1542 Py_TYPE(lz)->tp_free(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03001543 Py_TRASHCAN_SAFE_END(lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001544}
1545
1546static int
1547imap_traverse(imapobject *lz, visitproc visit, void *arg)
1548{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001549 Py_VISIT(lz->iters);
1550 Py_VISIT(lz->func);
1551 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001552}
1553
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001554/*
Raymond Hettinger2012f172003-02-07 05:32:58 +00001555imap() is an iterator version of __builtins__.map() except that it does
1556not have the None fill-in feature. That was intentionally left out for
1557the following reasons:
1558
1559 1) Itertools are designed to be easily combined and chained together.
1560 Having all tools stop with the shortest input is a unifying principle
1561 that makes it easier to combine finite iterators (supplying data) with
1562 infinite iterators like count() and repeat() (for supplying sequential
1563 or constant arguments to a function).
1564
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001565 2) In typical use cases for combining itertools, having one finite data
Raymond Hettinger2012f172003-02-07 05:32:58 +00001566 supplier run out before another is likely to be an error condition which
1567 should not pass silently by automatically supplying None.
1568
1569 3) The use cases for automatic None fill-in are rare -- not many functions
1570 do something useful when a parameter suddenly switches type and becomes
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001571 None.
Raymond Hettinger2012f172003-02-07 05:32:58 +00001572
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001573 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001574 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001575
1576 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1577*/
1578
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001579static PyObject *
1580imap_next(imapobject *lz)
1581{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001582 PyObject *val;
1583 PyObject *argtuple;
1584 PyObject *result;
1585 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001586
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001587 numargs = PyTuple_Size(lz->iters);
1588 argtuple = PyTuple_New(numargs);
1589 if (argtuple == NULL)
1590 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001591
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001592 for (i=0 ; i<numargs ; i++) {
1593 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1594 if (val == NULL) {
1595 Py_DECREF(argtuple);
1596 return NULL;
1597 }
1598 PyTuple_SET_ITEM(argtuple, i, val);
1599 }
1600 if (lz->func == Py_None)
1601 return argtuple;
1602 result = PyObject_Call(lz->func, argtuple, NULL);
1603 Py_DECREF(argtuple);
1604 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001605}
1606
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001607PyDoc_STRVAR(imap_doc,
1608"imap(func, *iterables) --> imap object\n\
1609\n\
1610Make an iterator that computes the function using arguments from\n\
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001611each of the iterables. Like map() except that it returns\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001612an iterator instead of a list and that it stops when the shortest\n\
1613iterable is exhausted instead of filling in None for shorter\n\
1614iterables.");
1615
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001616static PyTypeObject imap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001617 PyVarObject_HEAD_INIT(NULL, 0)
1618 "itertools.imap", /* tp_name */
1619 sizeof(imapobject), /* tp_basicsize */
1620 0, /* tp_itemsize */
1621 /* methods */
1622 (destructor)imap_dealloc, /* tp_dealloc */
1623 0, /* tp_print */
1624 0, /* tp_getattr */
1625 0, /* tp_setattr */
1626 0, /* tp_compare */
1627 0, /* tp_repr */
1628 0, /* tp_as_number */
1629 0, /* tp_as_sequence */
1630 0, /* tp_as_mapping */
1631 0, /* tp_hash */
1632 0, /* tp_call */
1633 0, /* tp_str */
1634 PyObject_GenericGetAttr, /* tp_getattro */
1635 0, /* tp_setattro */
1636 0, /* tp_as_buffer */
1637 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1638 Py_TPFLAGS_BASETYPE, /* tp_flags */
1639 imap_doc, /* tp_doc */
1640 (traverseproc)imap_traverse, /* tp_traverse */
1641 0, /* tp_clear */
1642 0, /* tp_richcompare */
1643 0, /* tp_weaklistoffset */
1644 PyObject_SelfIter, /* tp_iter */
1645 (iternextfunc)imap_next, /* tp_iternext */
1646 0, /* tp_methods */
1647 0, /* tp_members */
1648 0, /* tp_getset */
1649 0, /* tp_base */
1650 0, /* tp_dict */
1651 0, /* tp_descr_get */
1652 0, /* tp_descr_set */
1653 0, /* tp_dictoffset */
1654 0, /* tp_init */
1655 0, /* tp_alloc */
1656 imap_new, /* tp_new */
1657 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001658};
1659
1660
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001661/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001662
1663typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001664 PyObject_HEAD
1665 PyObject *source; /* Iterator over input iterables */
1666 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001667} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001668
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001669static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001670
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001671static PyObject *
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001672chain_new_internal(PyTypeObject *type, PyObject *source)
1673{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001674 chainobject *lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001675
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001676 lz = (chainobject *)type->tp_alloc(type, 0);
1677 if (lz == NULL) {
1678 Py_DECREF(source);
1679 return NULL;
1680 }
1681
1682 lz->source = source;
1683 lz->active = NULL;
1684 return (PyObject *)lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001685}
1686
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001687static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001688chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001689{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001690 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001691
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001692 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1693 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001694
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001695 source = PyObject_GetIter(args);
1696 if (source == NULL)
1697 return NULL;
1698
1699 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001700}
1701
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001702static PyObject *
1703chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1704{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001705 PyObject *source;
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001706
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001707 source = PyObject_GetIter(arg);
1708 if (source == NULL)
1709 return NULL;
1710
1711 return chain_new_internal(type, source);
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001712}
1713
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001714static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001715chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001716{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001717 PyObject_GC_UnTrack(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03001718 Py_TRASHCAN_SAFE_BEGIN(lz)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001719 Py_XDECREF(lz->active);
1720 Py_XDECREF(lz->source);
1721 Py_TYPE(lz)->tp_free(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03001722 Py_TRASHCAN_SAFE_END(lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001723}
1724
Raymond Hettinger2012f172003-02-07 05:32:58 +00001725static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001726chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001727{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001728 Py_VISIT(lz->source);
1729 Py_VISIT(lz->active);
1730 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001731}
1732
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001733static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001734chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001735{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001736 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001737
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001738 if (lz->source == NULL)
1739 return NULL; /* already stopped */
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001740
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001741 if (lz->active == NULL) {
1742 PyObject *iterable = PyIter_Next(lz->source);
1743 if (iterable == NULL) {
1744 Py_CLEAR(lz->source);
1745 return NULL; /* no more input sources */
1746 }
1747 lz->active = PyObject_GetIter(iterable);
1748 Py_DECREF(iterable);
1749 if (lz->active == NULL) {
1750 Py_CLEAR(lz->source);
1751 return NULL; /* input not iterable */
1752 }
1753 }
1754 item = PyIter_Next(lz->active);
1755 if (item != NULL)
1756 return item;
1757 if (PyErr_Occurred()) {
1758 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1759 PyErr_Clear();
1760 else
1761 return NULL; /* input raised an exception */
1762 }
1763 Py_CLEAR(lz->active);
1764 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001765}
1766
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001767PyDoc_STRVAR(chain_doc,
1768"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001769\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001770Return a chain object whose .next() method returns elements from the\n\
1771first iterable until it is exhausted, then elements from the next\n\
1772iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001773
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001774PyDoc_STRVAR(chain_from_iterable_doc,
1775"chain.from_iterable(iterable) --> chain object\n\
1776\n\
1777Alternate chain() contructor taking a single iterable argument\n\
1778that evaluates lazily.");
1779
1780static PyMethodDef chain_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001781 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1782 chain_from_iterable_doc},
1783 {NULL, NULL} /* sentinel */
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001784};
1785
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001786static PyTypeObject chain_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001787 PyVarObject_HEAD_INIT(NULL, 0)
1788 "itertools.chain", /* tp_name */
1789 sizeof(chainobject), /* tp_basicsize */
1790 0, /* tp_itemsize */
1791 /* methods */
1792 (destructor)chain_dealloc, /* tp_dealloc */
1793 0, /* tp_print */
1794 0, /* tp_getattr */
1795 0, /* tp_setattr */
1796 0, /* tp_compare */
1797 0, /* tp_repr */
1798 0, /* tp_as_number */
1799 0, /* tp_as_sequence */
1800 0, /* tp_as_mapping */
1801 0, /* tp_hash */
1802 0, /* tp_call */
1803 0, /* tp_str */
1804 PyObject_GenericGetAttr, /* tp_getattro */
1805 0, /* tp_setattro */
1806 0, /* tp_as_buffer */
1807 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1808 Py_TPFLAGS_BASETYPE, /* tp_flags */
1809 chain_doc, /* tp_doc */
1810 (traverseproc)chain_traverse, /* tp_traverse */
1811 0, /* tp_clear */
1812 0, /* tp_richcompare */
1813 0, /* tp_weaklistoffset */
1814 PyObject_SelfIter, /* tp_iter */
1815 (iternextfunc)chain_next, /* tp_iternext */
1816 chain_methods, /* tp_methods */
1817 0, /* tp_members */
1818 0, /* tp_getset */
1819 0, /* tp_base */
1820 0, /* tp_dict */
1821 0, /* tp_descr_get */
1822 0, /* tp_descr_set */
1823 0, /* tp_dictoffset */
1824 0, /* tp_init */
1825 0, /* tp_alloc */
1826 chain_new, /* tp_new */
1827 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001828};
1829
1830
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001831/* product object ************************************************************/
1832
1833typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001834 PyObject_HEAD
1835 PyObject *pools; /* tuple of pool tuples */
1836 Py_ssize_t *indices; /* one index per pool */
1837 PyObject *result; /* most recently returned result tuple */
1838 int stopped; /* set to 1 when the product iterator is exhausted */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001839} productobject;
1840
1841static PyTypeObject product_type;
1842
1843static PyObject *
1844product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1845{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001846 productobject *lz;
1847 Py_ssize_t nargs, npools, repeat=1;
1848 PyObject *pools = NULL;
1849 Py_ssize_t *indices = NULL;
1850 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001851
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001852 if (kwds != NULL) {
1853 char *kwlist[] = {"repeat", 0};
1854 PyObject *tmpargs = PyTuple_New(0);
1855 if (tmpargs == NULL)
1856 return NULL;
1857 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1858 Py_DECREF(tmpargs);
1859 return NULL;
1860 }
1861 Py_DECREF(tmpargs);
1862 if (repeat < 0) {
1863 PyErr_SetString(PyExc_ValueError,
1864 "repeat argument cannot be negative");
1865 return NULL;
1866 }
1867 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001868
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001869 assert(PyTuple_Check(args));
1870 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1871 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001872
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001873 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1874 if (indices == NULL) {
1875 PyErr_NoMemory();
1876 goto error;
1877 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001878
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001879 pools = PyTuple_New(npools);
1880 if (pools == NULL)
1881 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001882
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001883 for (i=0; i < nargs ; ++i) {
1884 PyObject *item = PyTuple_GET_ITEM(args, i);
1885 PyObject *pool = PySequence_Tuple(item);
1886 if (pool == NULL)
1887 goto error;
1888 PyTuple_SET_ITEM(pools, i, pool);
1889 indices[i] = 0;
1890 }
1891 for ( ; i < npools; ++i) {
1892 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1893 Py_INCREF(pool);
1894 PyTuple_SET_ITEM(pools, i, pool);
1895 indices[i] = 0;
1896 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001897
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001898 /* create productobject structure */
1899 lz = (productobject *)type->tp_alloc(type, 0);
1900 if (lz == NULL)
1901 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001902
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001903 lz->pools = pools;
1904 lz->indices = indices;
1905 lz->result = NULL;
1906 lz->stopped = 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001907
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001908 return (PyObject *)lz;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001909
1910error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001911 if (indices != NULL)
1912 PyMem_Free(indices);
1913 Py_XDECREF(pools);
1914 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001915}
1916
1917static void
1918product_dealloc(productobject *lz)
1919{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001920 PyObject_GC_UnTrack(lz);
1921 Py_XDECREF(lz->pools);
1922 Py_XDECREF(lz->result);
1923 if (lz->indices != NULL)
1924 PyMem_Free(lz->indices);
1925 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001926}
1927
1928static int
1929product_traverse(productobject *lz, visitproc visit, void *arg)
1930{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001931 Py_VISIT(lz->pools);
1932 Py_VISIT(lz->result);
1933 return 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001934}
1935
1936static PyObject *
1937product_next(productobject *lz)
1938{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001939 PyObject *pool;
1940 PyObject *elem;
1941 PyObject *oldelem;
1942 PyObject *pools = lz->pools;
1943 PyObject *result = lz->result;
1944 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1945 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001946
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001947 if (lz->stopped)
1948 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001949
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001950 if (result == NULL) {
1951 /* On the first pass, return an initial tuple filled with the
1952 first element from each pool. */
1953 result = PyTuple_New(npools);
1954 if (result == NULL)
1955 goto empty;
1956 lz->result = result;
1957 for (i=0; i < npools; i++) {
1958 pool = PyTuple_GET_ITEM(pools, i);
1959 if (PyTuple_GET_SIZE(pool) == 0)
1960 goto empty;
1961 elem = PyTuple_GET_ITEM(pool, 0);
1962 Py_INCREF(elem);
1963 PyTuple_SET_ITEM(result, i, elem);
1964 }
1965 } else {
1966 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001967
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001968 /* Copy the previous result tuple or re-use it if available */
1969 if (Py_REFCNT(result) > 1) {
1970 PyObject *old_result = result;
1971 result = PyTuple_New(npools);
1972 if (result == NULL)
1973 goto empty;
1974 lz->result = result;
1975 for (i=0; i < npools; i++) {
1976 elem = PyTuple_GET_ITEM(old_result, i);
1977 Py_INCREF(elem);
1978 PyTuple_SET_ITEM(result, i, elem);
1979 }
1980 Py_DECREF(old_result);
1981 }
1982 /* Now, we've got the only copy so we can update it in-place */
1983 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001984
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001985 /* Update the pool indices right-to-left. Only advance to the
1986 next pool when the previous one rolls-over */
1987 for (i=npools-1 ; i >= 0 ; i--) {
1988 pool = PyTuple_GET_ITEM(pools, i);
1989 indices[i]++;
1990 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1991 /* Roll-over and advance to next pool */
1992 indices[i] = 0;
1993 elem = PyTuple_GET_ITEM(pool, 0);
1994 Py_INCREF(elem);
1995 oldelem = PyTuple_GET_ITEM(result, i);
1996 PyTuple_SET_ITEM(result, i, elem);
1997 Py_DECREF(oldelem);
1998 } else {
1999 /* No rollover. Just increment and stop here. */
2000 elem = PyTuple_GET_ITEM(pool, indices[i]);
2001 Py_INCREF(elem);
2002 oldelem = PyTuple_GET_ITEM(result, i);
2003 PyTuple_SET_ITEM(result, i, elem);
2004 Py_DECREF(oldelem);
2005 break;
2006 }
2007 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00002008
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002009 /* If i is negative, then the indices have all rolled-over
2010 and we're done. */
2011 if (i < 0)
2012 goto empty;
2013 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002014
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002015 Py_INCREF(result);
2016 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002017
2018empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002019 lz->stopped = 1;
2020 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002021}
2022
2023PyDoc_STRVAR(product_doc,
2024"product(*iterables) --> product object\n\
2025\n\
Raymond Hettinger73d79632008-02-23 02:20:41 +00002026Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002027For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2028The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2029cycle in a manner similar to an odometer (with the rightmost element changing\n\
2030on every iteration).\n\n\
Raymond Hettinger94034ea2009-02-09 18:39:41 +00002031To compute the product of an iterable with itself, specify the number\n\
2032of repetitions with the optional repeat keyword argument. For example,\n\
2033product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002034product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2035product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2036
2037static PyTypeObject product_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002038 PyVarObject_HEAD_INIT(NULL, 0)
2039 "itertools.product", /* tp_name */
2040 sizeof(productobject), /* tp_basicsize */
2041 0, /* tp_itemsize */
2042 /* methods */
2043 (destructor)product_dealloc, /* tp_dealloc */
2044 0, /* tp_print */
2045 0, /* tp_getattr */
2046 0, /* tp_setattr */
2047 0, /* tp_compare */
2048 0, /* tp_repr */
2049 0, /* tp_as_number */
2050 0, /* tp_as_sequence */
2051 0, /* tp_as_mapping */
2052 0, /* tp_hash */
2053 0, /* tp_call */
2054 0, /* tp_str */
2055 PyObject_GenericGetAttr, /* tp_getattro */
2056 0, /* tp_setattro */
2057 0, /* tp_as_buffer */
2058 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2059 Py_TPFLAGS_BASETYPE, /* tp_flags */
2060 product_doc, /* tp_doc */
2061 (traverseproc)product_traverse, /* tp_traverse */
2062 0, /* tp_clear */
2063 0, /* tp_richcompare */
2064 0, /* tp_weaklistoffset */
2065 PyObject_SelfIter, /* tp_iter */
2066 (iternextfunc)product_next, /* tp_iternext */
2067 0, /* tp_methods */
2068 0, /* tp_members */
2069 0, /* tp_getset */
2070 0, /* tp_base */
2071 0, /* tp_dict */
2072 0, /* tp_descr_get */
2073 0, /* tp_descr_set */
2074 0, /* tp_dictoffset */
2075 0, /* tp_init */
2076 0, /* tp_alloc */
2077 product_new, /* tp_new */
2078 PyObject_GC_Del, /* tp_free */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002079};
2080
2081
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002082/* combinations object ************************************************************/
2083
2084typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002085 PyObject_HEAD
2086 PyObject *pool; /* input converted to a tuple */
2087 Py_ssize_t *indices; /* one index per result element */
2088 PyObject *result; /* most recently returned result tuple */
2089 Py_ssize_t r; /* size of result tuple */
2090 int stopped; /* set to 1 when the combinations iterator is exhausted */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002091} combinationsobject;
2092
2093static PyTypeObject combinations_type;
2094
2095static PyObject *
2096combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2097{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002098 combinationsobject *co;
2099 Py_ssize_t n;
2100 Py_ssize_t r;
2101 PyObject *pool = NULL;
2102 PyObject *iterable = NULL;
2103 Py_ssize_t *indices = NULL;
2104 Py_ssize_t i;
2105 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002106
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002107 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2108 &iterable, &r))
2109 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002110
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002111 pool = PySequence_Tuple(iterable);
2112 if (pool == NULL)
2113 goto error;
2114 n = PyTuple_GET_SIZE(pool);
2115 if (r < 0) {
2116 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2117 goto error;
2118 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002119
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002120 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2121 if (indices == NULL) {
2122 PyErr_NoMemory();
2123 goto error;
2124 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002125
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002126 for (i=0 ; i<r ; i++)
2127 indices[i] = i;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002128
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002129 /* create combinationsobject structure */
2130 co = (combinationsobject *)type->tp_alloc(type, 0);
2131 if (co == NULL)
2132 goto error;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002133
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002134 co->pool = pool;
2135 co->indices = indices;
2136 co->result = NULL;
2137 co->r = r;
2138 co->stopped = r > n ? 1 : 0;
2139
2140 return (PyObject *)co;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002141
2142error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002143 if (indices != NULL)
2144 PyMem_Free(indices);
2145 Py_XDECREF(pool);
2146 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002147}
2148
2149static void
2150combinations_dealloc(combinationsobject *co)
2151{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002152 PyObject_GC_UnTrack(co);
2153 Py_XDECREF(co->pool);
2154 Py_XDECREF(co->result);
2155 if (co->indices != NULL)
2156 PyMem_Free(co->indices);
2157 Py_TYPE(co)->tp_free(co);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002158}
2159
2160static int
2161combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2162{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002163 Py_VISIT(co->pool);
2164 Py_VISIT(co->result);
2165 return 0;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002166}
2167
2168static PyObject *
2169combinations_next(combinationsobject *co)
2170{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002171 PyObject *elem;
2172 PyObject *oldelem;
2173 PyObject *pool = co->pool;
2174 Py_ssize_t *indices = co->indices;
2175 PyObject *result = co->result;
2176 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2177 Py_ssize_t r = co->r;
2178 Py_ssize_t i, j, index;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002179
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002180 if (co->stopped)
2181 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002182
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002183 if (result == NULL) {
2184 /* On the first pass, initialize result tuple using the indices */
2185 result = PyTuple_New(r);
2186 if (result == NULL)
2187 goto empty;
2188 co->result = result;
2189 for (i=0; i<r ; i++) {
2190 index = indices[i];
2191 elem = PyTuple_GET_ITEM(pool, index);
2192 Py_INCREF(elem);
2193 PyTuple_SET_ITEM(result, i, elem);
2194 }
2195 } else {
2196 /* Copy the previous result tuple or re-use it if available */
2197 if (Py_REFCNT(result) > 1) {
2198 PyObject *old_result = result;
2199 result = PyTuple_New(r);
2200 if (result == NULL)
2201 goto empty;
2202 co->result = result;
2203 for (i=0; i<r ; i++) {
2204 elem = PyTuple_GET_ITEM(old_result, i);
2205 Py_INCREF(elem);
2206 PyTuple_SET_ITEM(result, i, elem);
2207 }
2208 Py_DECREF(old_result);
2209 }
2210 /* Now, we've got the only copy so we can update it in-place
2211 * CPython's empty tuple is a singleton and cached in
2212 * PyTuple's freelist.
2213 */
2214 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002215
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002216 /* Scan indices right-to-left until finding one that is not
2217 at its maximum (i + n - r). */
2218 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2219 ;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002220
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002221 /* If i is negative, then the indices are all at
2222 their maximum value and we're done. */
2223 if (i < 0)
2224 goto empty;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002225
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002226 /* Increment the current index which we know is not at its
2227 maximum. Then move back to the right setting each index
2228 to its lowest possible value (one higher than the index
2229 to its left -- this maintains the sort order invariant). */
2230 indices[i]++;
2231 for (j=i+1 ; j<r ; j++)
2232 indices[j] = indices[j-1] + 1;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002233
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002234 /* Update the result tuple for the new indices
2235 starting with i, the leftmost index that changed */
2236 for ( ; i<r ; i++) {
2237 index = indices[i];
2238 elem = PyTuple_GET_ITEM(pool, index);
2239 Py_INCREF(elem);
2240 oldelem = PyTuple_GET_ITEM(result, i);
2241 PyTuple_SET_ITEM(result, i, elem);
2242 Py_DECREF(oldelem);
2243 }
2244 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002245
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002246 Py_INCREF(result);
2247 return result;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002248
2249empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002250 co->stopped = 1;
2251 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002252}
2253
2254PyDoc_STRVAR(combinations_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002255"combinations(iterable, r) --> combinations object\n\
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002256\n\
2257Return successive r-length combinations of elements in the iterable.\n\n\
2258combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2259
2260static PyTypeObject combinations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002261 PyVarObject_HEAD_INIT(NULL, 0)
2262 "itertools.combinations", /* tp_name */
2263 sizeof(combinationsobject), /* tp_basicsize */
2264 0, /* tp_itemsize */
2265 /* methods */
2266 (destructor)combinations_dealloc, /* tp_dealloc */
2267 0, /* tp_print */
2268 0, /* tp_getattr */
2269 0, /* tp_setattr */
2270 0, /* tp_compare */
2271 0, /* tp_repr */
2272 0, /* tp_as_number */
2273 0, /* tp_as_sequence */
2274 0, /* tp_as_mapping */
2275 0, /* tp_hash */
2276 0, /* tp_call */
2277 0, /* tp_str */
2278 PyObject_GenericGetAttr, /* tp_getattro */
2279 0, /* tp_setattro */
2280 0, /* tp_as_buffer */
2281 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2282 Py_TPFLAGS_BASETYPE, /* tp_flags */
2283 combinations_doc, /* tp_doc */
2284 (traverseproc)combinations_traverse, /* tp_traverse */
2285 0, /* tp_clear */
2286 0, /* tp_richcompare */
2287 0, /* tp_weaklistoffset */
2288 PyObject_SelfIter, /* tp_iter */
2289 (iternextfunc)combinations_next, /* tp_iternext */
2290 0, /* tp_methods */
2291 0, /* tp_members */
2292 0, /* tp_getset */
2293 0, /* tp_base */
2294 0, /* tp_dict */
2295 0, /* tp_descr_get */
2296 0, /* tp_descr_set */
2297 0, /* tp_dictoffset */
2298 0, /* tp_init */
2299 0, /* tp_alloc */
2300 combinations_new, /* tp_new */
2301 PyObject_GC_Del, /* tp_free */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002302};
2303
2304
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002305/* combinations with replacement object *******************************************/
2306
2307/* Equivalent to:
2308
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002309 def combinations_with_replacement(iterable, r):
2310 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2311 # number items returned: (n+r-1)! / r! / (n-1)!
2312 pool = tuple(iterable)
2313 n = len(pool)
2314 indices = [0] * r
2315 yield tuple(pool[i] for i in indices)
2316 while 1:
2317 for i in reversed(range(r)):
2318 if indices[i] != n - 1:
2319 break
2320 else:
2321 return
2322 indices[i:] = [indices[i] + 1] * (r - i)
2323 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002324
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002325 def combinations_with_replacement2(iterable, r):
2326 'Alternate version that filters from product()'
2327 pool = tuple(iterable)
2328 n = len(pool)
2329 for indices in product(range(n), repeat=r):
2330 if sorted(indices) == list(indices):
2331 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002332*/
2333typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002334 PyObject_HEAD
2335 PyObject *pool; /* input converted to a tuple */
2336 Py_ssize_t *indices; /* one index per result element */
2337 PyObject *result; /* most recently returned result tuple */
2338 Py_ssize_t r; /* size of result tuple */
2339 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002340} cwrobject;
2341
2342static PyTypeObject cwr_type;
2343
2344static PyObject *
2345cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2346{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002347 cwrobject *co;
2348 Py_ssize_t n;
2349 Py_ssize_t r;
2350 PyObject *pool = NULL;
2351 PyObject *iterable = NULL;
2352 Py_ssize_t *indices = NULL;
2353 Py_ssize_t i;
2354 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002355
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002356 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2357 &iterable, &r))
2358 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002359
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002360 pool = PySequence_Tuple(iterable);
2361 if (pool == NULL)
2362 goto error;
2363 n = PyTuple_GET_SIZE(pool);
2364 if (r < 0) {
2365 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2366 goto error;
2367 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002368
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002369 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2370 if (indices == NULL) {
2371 PyErr_NoMemory();
2372 goto error;
2373 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002374
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002375 for (i=0 ; i<r ; i++)
2376 indices[i] = 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002377
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002378 /* create cwrobject structure */
2379 co = (cwrobject *)type->tp_alloc(type, 0);
2380 if (co == NULL)
2381 goto error;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002382
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002383 co->pool = pool;
2384 co->indices = indices;
2385 co->result = NULL;
2386 co->r = r;
2387 co->stopped = !n && r;
2388
2389 return (PyObject *)co;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002390
2391error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002392 if (indices != NULL)
2393 PyMem_Free(indices);
2394 Py_XDECREF(pool);
2395 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002396}
2397
2398static void
2399cwr_dealloc(cwrobject *co)
2400{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002401 PyObject_GC_UnTrack(co);
2402 Py_XDECREF(co->pool);
2403 Py_XDECREF(co->result);
2404 if (co->indices != NULL)
2405 PyMem_Free(co->indices);
2406 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002407}
2408
2409static int
2410cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2411{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002412 Py_VISIT(co->pool);
2413 Py_VISIT(co->result);
2414 return 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002415}
2416
2417static PyObject *
2418cwr_next(cwrobject *co)
2419{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002420 PyObject *elem;
2421 PyObject *oldelem;
2422 PyObject *pool = co->pool;
2423 Py_ssize_t *indices = co->indices;
2424 PyObject *result = co->result;
2425 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2426 Py_ssize_t r = co->r;
2427 Py_ssize_t i, j, index;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002428
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002429 if (co->stopped)
2430 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002431
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002432 if (result == NULL) {
2433 /* On the first pass, initialize result tuple using the indices */
2434 result = PyTuple_New(r);
2435 if (result == NULL)
2436 goto empty;
2437 co->result = result;
2438 for (i=0; i<r ; i++) {
2439 index = indices[i];
2440 elem = PyTuple_GET_ITEM(pool, index);
2441 Py_INCREF(elem);
2442 PyTuple_SET_ITEM(result, i, elem);
2443 }
2444 } else {
2445 /* Copy the previous result tuple or re-use it if available */
2446 if (Py_REFCNT(result) > 1) {
2447 PyObject *old_result = result;
2448 result = PyTuple_New(r);
2449 if (result == NULL)
2450 goto empty;
2451 co->result = result;
2452 for (i=0; i<r ; i++) {
2453 elem = PyTuple_GET_ITEM(old_result, i);
2454 Py_INCREF(elem);
2455 PyTuple_SET_ITEM(result, i, elem);
2456 }
2457 Py_DECREF(old_result);
2458 }
2459 /* Now, we've got the only copy so we can update it in-place CPython's
2460 empty tuple is a singleton and cached in PyTuple's freelist. */
2461 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002462
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002463 /* Scan indices right-to-left until finding one that is not
2464 * at its maximum (n-1). */
2465 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2466 ;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002467
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002468 /* If i is negative, then the indices are all at
2469 their maximum value and we're done. */
2470 if (i < 0)
2471 goto empty;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002472
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002473 /* Increment the current index which we know is not at its
2474 maximum. Then set all to the right to the same value. */
2475 indices[i]++;
2476 for (j=i+1 ; j<r ; j++)
2477 indices[j] = indices[j-1];
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002478
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002479 /* Update the result tuple for the new indices
2480 starting with i, the leftmost index that changed */
2481 for ( ; i<r ; i++) {
2482 index = indices[i];
2483 elem = PyTuple_GET_ITEM(pool, index);
2484 Py_INCREF(elem);
2485 oldelem = PyTuple_GET_ITEM(result, i);
2486 PyTuple_SET_ITEM(result, i, elem);
2487 Py_DECREF(oldelem);
2488 }
2489 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002490
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002491 Py_INCREF(result);
2492 return result;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002493
2494empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002495 co->stopped = 1;
2496 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002497}
2498
2499PyDoc_STRVAR(cwr_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002500"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002501\n\
2502Return successive r-length combinations of elements in the iterable\n\
2503allowing individual elements to have successive repeats.\n\
2504combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2505
2506static PyTypeObject cwr_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002507 PyVarObject_HEAD_INIT(NULL, 0)
2508 "itertools.combinations_with_replacement", /* tp_name */
2509 sizeof(cwrobject), /* tp_basicsize */
2510 0, /* tp_itemsize */
2511 /* methods */
2512 (destructor)cwr_dealloc, /* tp_dealloc */
2513 0, /* tp_print */
2514 0, /* tp_getattr */
2515 0, /* tp_setattr */
2516 0, /* tp_compare */
2517 0, /* tp_repr */
2518 0, /* tp_as_number */
2519 0, /* tp_as_sequence */
2520 0, /* tp_as_mapping */
2521 0, /* tp_hash */
2522 0, /* tp_call */
2523 0, /* tp_str */
2524 PyObject_GenericGetAttr, /* tp_getattro */
2525 0, /* tp_setattro */
2526 0, /* tp_as_buffer */
2527 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2528 Py_TPFLAGS_BASETYPE, /* tp_flags */
2529 cwr_doc, /* tp_doc */
2530 (traverseproc)cwr_traverse, /* tp_traverse */
2531 0, /* tp_clear */
2532 0, /* tp_richcompare */
2533 0, /* tp_weaklistoffset */
2534 PyObject_SelfIter, /* tp_iter */
2535 (iternextfunc)cwr_next, /* tp_iternext */
2536 0, /* tp_methods */
2537 0, /* tp_members */
2538 0, /* tp_getset */
2539 0, /* tp_base */
2540 0, /* tp_dict */
2541 0, /* tp_descr_get */
2542 0, /* tp_descr_set */
2543 0, /* tp_dictoffset */
2544 0, /* tp_init */
2545 0, /* tp_alloc */
2546 cwr_new, /* tp_new */
2547 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002548};
2549
2550
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002551/* permutations object ************************************************************
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002552
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002553def permutations(iterable, r=None):
2554 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2555 pool = tuple(iterable)
2556 n = len(pool)
2557 r = n if r is None else r
2558 indices = range(n)
2559 cycles = range(n-r+1, n+1)[::-1]
2560 yield tuple(pool[i] for i in indices[:r])
2561 while n:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002562 for i in reversed(range(r)):
2563 cycles[i] -= 1
2564 if cycles[i] == 0:
2565 indices[i:] = indices[i+1:] + indices[i:i+1]
2566 cycles[i] = n - i
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002567 else:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002568 j = cycles[i]
2569 indices[i], indices[-j] = indices[-j], indices[i]
2570 yield tuple(pool[i] for i in indices[:r])
2571 break
2572 else:
2573 return
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002574*/
2575
2576typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002577 PyObject_HEAD
2578 PyObject *pool; /* input converted to a tuple */
2579 Py_ssize_t *indices; /* one index per element in the pool */
2580 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2581 PyObject *result; /* most recently returned result tuple */
2582 Py_ssize_t r; /* size of result tuple */
2583 int stopped; /* set to 1 when the permutations iterator is exhausted */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002584} permutationsobject;
2585
2586static PyTypeObject permutations_type;
2587
2588static PyObject *
2589permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2590{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002591 permutationsobject *po;
2592 Py_ssize_t n;
2593 Py_ssize_t r;
2594 PyObject *robj = Py_None;
2595 PyObject *pool = NULL;
2596 PyObject *iterable = NULL;
2597 Py_ssize_t *indices = NULL;
2598 Py_ssize_t *cycles = NULL;
2599 Py_ssize_t i;
2600 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002601
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002602 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2603 &iterable, &robj))
2604 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002605
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002606 pool = PySequence_Tuple(iterable);
2607 if (pool == NULL)
2608 goto error;
2609 n = PyTuple_GET_SIZE(pool);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002610
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002611 r = n;
2612 if (robj != Py_None) {
2613 r = PyInt_AsSsize_t(robj);
2614 if (r == -1 && PyErr_Occurred())
2615 goto error;
2616 }
2617 if (r < 0) {
2618 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2619 goto error;
2620 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002621
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002622 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2623 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2624 if (indices == NULL || cycles == NULL) {
2625 PyErr_NoMemory();
2626 goto error;
2627 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002628
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002629 for (i=0 ; i<n ; i++)
2630 indices[i] = i;
2631 for (i=0 ; i<r ; i++)
2632 cycles[i] = n - i;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002633
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002634 /* create permutationsobject structure */
2635 po = (permutationsobject *)type->tp_alloc(type, 0);
2636 if (po == NULL)
2637 goto error;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002638
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002639 po->pool = pool;
2640 po->indices = indices;
2641 po->cycles = cycles;
2642 po->result = NULL;
2643 po->r = r;
2644 po->stopped = r > n ? 1 : 0;
2645
2646 return (PyObject *)po;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002647
2648error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002649 if (indices != NULL)
2650 PyMem_Free(indices);
2651 if (cycles != NULL)
2652 PyMem_Free(cycles);
2653 Py_XDECREF(pool);
2654 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002655}
2656
2657static void
2658permutations_dealloc(permutationsobject *po)
2659{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002660 PyObject_GC_UnTrack(po);
2661 Py_XDECREF(po->pool);
2662 Py_XDECREF(po->result);
2663 PyMem_Free(po->indices);
2664 PyMem_Free(po->cycles);
2665 Py_TYPE(po)->tp_free(po);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002666}
2667
2668static int
2669permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2670{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002671 Py_VISIT(po->pool);
2672 Py_VISIT(po->result);
2673 return 0;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002674}
2675
2676static PyObject *
2677permutations_next(permutationsobject *po)
2678{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002679 PyObject *elem;
2680 PyObject *oldelem;
2681 PyObject *pool = po->pool;
2682 Py_ssize_t *indices = po->indices;
2683 Py_ssize_t *cycles = po->cycles;
2684 PyObject *result = po->result;
2685 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2686 Py_ssize_t r = po->r;
2687 Py_ssize_t i, j, k, index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002688
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002689 if (po->stopped)
2690 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002691
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002692 if (result == NULL) {
2693 /* On the first pass, initialize result tuple using the indices */
2694 result = PyTuple_New(r);
2695 if (result == NULL)
2696 goto empty;
2697 po->result = result;
2698 for (i=0; i<r ; i++) {
2699 index = indices[i];
2700 elem = PyTuple_GET_ITEM(pool, index);
2701 Py_INCREF(elem);
2702 PyTuple_SET_ITEM(result, i, elem);
2703 }
2704 } else {
2705 if (n == 0)
2706 goto empty;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002707
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002708 /* Copy the previous result tuple or re-use it if available */
2709 if (Py_REFCNT(result) > 1) {
2710 PyObject *old_result = result;
2711 result = PyTuple_New(r);
2712 if (result == NULL)
2713 goto empty;
2714 po->result = result;
2715 for (i=0; i<r ; i++) {
2716 elem = PyTuple_GET_ITEM(old_result, i);
2717 Py_INCREF(elem);
2718 PyTuple_SET_ITEM(result, i, elem);
2719 }
2720 Py_DECREF(old_result);
2721 }
2722 /* Now, we've got the only copy so we can update it in-place */
2723 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002724
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002725 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2726 for (i=r-1 ; i>=0 ; i--) {
2727 cycles[i] -= 1;
2728 if (cycles[i] == 0) {
2729 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2730 index = indices[i];
2731 for (j=i ; j<n-1 ; j++)
2732 indices[j] = indices[j+1];
2733 indices[n-1] = index;
2734 cycles[i] = n - i;
2735 } else {
2736 j = cycles[i];
2737 index = indices[i];
2738 indices[i] = indices[n-j];
2739 indices[n-j] = index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002740
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002741 for (k=i; k<r ; k++) {
2742 /* start with i, the leftmost element that changed */
2743 /* yield tuple(pool[k] for k in indices[:r]) */
2744 index = indices[k];
2745 elem = PyTuple_GET_ITEM(pool, index);
2746 Py_INCREF(elem);
2747 oldelem = PyTuple_GET_ITEM(result, k);
2748 PyTuple_SET_ITEM(result, k, elem);
2749 Py_DECREF(oldelem);
2750 }
2751 break;
2752 }
2753 }
2754 /* If i is negative, then the cycles have all
2755 rolled-over and we're done. */
2756 if (i < 0)
2757 goto empty;
2758 }
2759 Py_INCREF(result);
2760 return result;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002761
2762empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002763 po->stopped = 1;
2764 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002765}
2766
2767PyDoc_STRVAR(permutations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +00002768"permutations(iterable[, r]) --> permutations object\n\
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002769\n\
2770Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl8550d692008-06-10 12:46:39 +00002771permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002772
2773static PyTypeObject permutations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002774 PyVarObject_HEAD_INIT(NULL, 0)
2775 "itertools.permutations", /* tp_name */
2776 sizeof(permutationsobject), /* tp_basicsize */
2777 0, /* tp_itemsize */
2778 /* methods */
2779 (destructor)permutations_dealloc, /* tp_dealloc */
2780 0, /* tp_print */
2781 0, /* tp_getattr */
2782 0, /* tp_setattr */
2783 0, /* tp_compare */
2784 0, /* tp_repr */
2785 0, /* tp_as_number */
2786 0, /* tp_as_sequence */
2787 0, /* tp_as_mapping */
2788 0, /* tp_hash */
2789 0, /* tp_call */
2790 0, /* tp_str */
2791 PyObject_GenericGetAttr, /* tp_getattro */
2792 0, /* tp_setattro */
2793 0, /* tp_as_buffer */
2794 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2795 Py_TPFLAGS_BASETYPE, /* tp_flags */
2796 permutations_doc, /* tp_doc */
2797 (traverseproc)permutations_traverse, /* tp_traverse */
2798 0, /* tp_clear */
2799 0, /* tp_richcompare */
2800 0, /* tp_weaklistoffset */
2801 PyObject_SelfIter, /* tp_iter */
2802 (iternextfunc)permutations_next, /* tp_iternext */
2803 0, /* tp_methods */
2804 0, /* tp_members */
2805 0, /* tp_getset */
2806 0, /* tp_base */
2807 0, /* tp_dict */
2808 0, /* tp_descr_get */
2809 0, /* tp_descr_set */
2810 0, /* tp_dictoffset */
2811 0, /* tp_init */
2812 0, /* tp_alloc */
2813 permutations_new, /* tp_new */
2814 PyObject_GC_Del, /* tp_free */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002815};
2816
2817
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002818/* compress object ************************************************************/
2819
2820/* Equivalent to:
2821
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002822 def compress(data, selectors):
2823 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2824 return (d for d, s in izip(data, selectors) if s)
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002825*/
2826
2827typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002828 PyObject_HEAD
2829 PyObject *data;
2830 PyObject *selectors;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002831} compressobject;
2832
2833static PyTypeObject compress_type;
2834
2835static PyObject *
2836compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2837{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002838 PyObject *seq1, *seq2;
2839 PyObject *data=NULL, *selectors=NULL;
2840 compressobject *lz;
2841 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002842
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002843 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2844 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002845
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002846 data = PyObject_GetIter(seq1);
2847 if (data == NULL)
2848 goto fail;
2849 selectors = PyObject_GetIter(seq2);
2850 if (selectors == NULL)
2851 goto fail;
2852
2853 /* create compressobject structure */
2854 lz = (compressobject *)type->tp_alloc(type, 0);
2855 if (lz == NULL)
2856 goto fail;
2857 lz->data = data;
2858 lz->selectors = selectors;
2859 return (PyObject *)lz;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002860
2861fail:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002862 Py_XDECREF(data);
2863 Py_XDECREF(selectors);
2864 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002865}
2866
2867static void
2868compress_dealloc(compressobject *lz)
2869{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002870 PyObject_GC_UnTrack(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03002871 Py_TRASHCAN_SAFE_BEGIN(lz)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002872 Py_XDECREF(lz->data);
2873 Py_XDECREF(lz->selectors);
2874 Py_TYPE(lz)->tp_free(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03002875 Py_TRASHCAN_SAFE_END(lz)
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002876}
2877
2878static int
2879compress_traverse(compressobject *lz, visitproc visit, void *arg)
2880{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002881 Py_VISIT(lz->data);
2882 Py_VISIT(lz->selectors);
2883 return 0;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002884}
2885
2886static PyObject *
2887compress_next(compressobject *lz)
2888{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002889 PyObject *data = lz->data, *selectors = lz->selectors;
2890 PyObject *datum, *selector;
2891 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2892 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2893 int ok;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002894
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002895 while (1) {
2896 /* Steps: get datum, get selector, evaluate selector.
2897 Order is important (to match the pure python version
2898 in terms of which input gets a chance to raise an
2899 exception first).
2900 */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002901
Serhiy Storchakabb845652013-04-06 22:04:10 +03002902 if (Py_EnterRecursiveCall(" while iterating"))
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002903 return NULL;
Serhiy Storchakabb845652013-04-06 22:04:10 +03002904 datum = datanext(data);
2905 if (datum == NULL) {
2906 Py_LeaveRecursiveCall();
2907 return NULL;
2908 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002909
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002910 selector = selectornext(selectors);
Serhiy Storchakabb845652013-04-06 22:04:10 +03002911 Py_LeaveRecursiveCall();
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002912 if (selector == NULL) {
2913 Py_DECREF(datum);
2914 return NULL;
2915 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002916
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002917 ok = PyObject_IsTrue(selector);
2918 Py_DECREF(selector);
2919 if (ok == 1)
2920 return datum;
2921 Py_DECREF(datum);
2922 if (ok == -1)
2923 return NULL;
2924 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002925}
2926
2927PyDoc_STRVAR(compress_doc,
Raymond Hettinger2e2909f2009-02-19 02:15:14 +00002928"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002929\n\
2930Return data elements corresponding to true selector elements.\n\
2931Forms a shorter iterator from selected data elements using the\n\
2932selectors to choose the data elements.");
2933
2934static PyTypeObject compress_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002935 PyVarObject_HEAD_INIT(NULL, 0)
2936 "itertools.compress", /* tp_name */
2937 sizeof(compressobject), /* tp_basicsize */
2938 0, /* tp_itemsize */
2939 /* methods */
2940 (destructor)compress_dealloc, /* tp_dealloc */
2941 0, /* tp_print */
2942 0, /* tp_getattr */
2943 0, /* tp_setattr */
2944 0, /* tp_compare */
2945 0, /* tp_repr */
2946 0, /* tp_as_number */
2947 0, /* tp_as_sequence */
2948 0, /* tp_as_mapping */
2949 0, /* tp_hash */
2950 0, /* tp_call */
2951 0, /* tp_str */
2952 PyObject_GenericGetAttr, /* tp_getattro */
2953 0, /* tp_setattro */
2954 0, /* tp_as_buffer */
2955 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2956 Py_TPFLAGS_BASETYPE, /* tp_flags */
2957 compress_doc, /* tp_doc */
2958 (traverseproc)compress_traverse, /* tp_traverse */
2959 0, /* tp_clear */
2960 0, /* tp_richcompare */
2961 0, /* tp_weaklistoffset */
2962 PyObject_SelfIter, /* tp_iter */
2963 (iternextfunc)compress_next, /* tp_iternext */
2964 0, /* tp_methods */
2965 0, /* tp_members */
2966 0, /* tp_getset */
2967 0, /* tp_base */
2968 0, /* tp_dict */
2969 0, /* tp_descr_get */
2970 0, /* tp_descr_set */
2971 0, /* tp_dictoffset */
2972 0, /* tp_init */
2973 0, /* tp_alloc */
2974 compress_new, /* tp_new */
2975 PyObject_GC_Del, /* tp_free */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002976};
2977
2978
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002979/* ifilter object ************************************************************/
2980
2981typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002982 PyObject_HEAD
2983 PyObject *func;
2984 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002985} ifilterobject;
2986
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002987static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002988
2989static PyObject *
2990ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2991{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002992 PyObject *func, *seq;
2993 PyObject *it;
2994 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002995
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002996 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
2997 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002998
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002999 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
3000 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003001
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003002 /* Get iterator. */
3003 it = PyObject_GetIter(seq);
3004 if (it == NULL)
3005 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003006
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003007 /* create ifilterobject structure */
3008 lz = (ifilterobject *)type->tp_alloc(type, 0);
3009 if (lz == NULL) {
3010 Py_DECREF(it);
3011 return NULL;
3012 }
3013 Py_INCREF(func);
3014 lz->func = func;
3015 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003016
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003017 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003018}
3019
3020static void
3021ifilter_dealloc(ifilterobject *lz)
3022{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003023 PyObject_GC_UnTrack(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03003024 Py_TRASHCAN_SAFE_BEGIN(lz)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003025 Py_XDECREF(lz->func);
3026 Py_XDECREF(lz->it);
3027 Py_TYPE(lz)->tp_free(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03003028 Py_TRASHCAN_SAFE_END(lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003029}
3030
3031static int
3032ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
3033{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003034 Py_VISIT(lz->it);
3035 Py_VISIT(lz->func);
3036 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003037}
3038
3039static PyObject *
3040ifilter_next(ifilterobject *lz)
3041{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003042 PyObject *item;
3043 PyObject *it = lz->it;
3044 long ok;
3045 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003046
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003047 iternext = *Py_TYPE(it)->tp_iternext;
3048 for (;;) {
Serhiy Storchakabb845652013-04-06 22:04:10 +03003049 if (Py_EnterRecursiveCall(" while iterating"))
3050 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003051 item = iternext(it);
Serhiy Storchakabb845652013-04-06 22:04:10 +03003052 Py_LeaveRecursiveCall();
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003053 if (item == NULL)
3054 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003055
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003056 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3057 ok = PyObject_IsTrue(item);
3058 } else {
3059 PyObject *good;
3060 good = PyObject_CallFunctionObjArgs(lz->func,
3061 item, NULL);
3062 if (good == NULL) {
3063 Py_DECREF(item);
3064 return NULL;
3065 }
3066 ok = PyObject_IsTrue(good);
3067 Py_DECREF(good);
3068 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003069 if (ok > 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003070 return item;
3071 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003072 if (ok < 0)
3073 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003074 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003075}
3076
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003077PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003078"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003079\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003080Return those items of sequence for which function(item) is true.\n\
3081If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003082
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003083static PyTypeObject ifilter_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003084 PyVarObject_HEAD_INIT(NULL, 0)
3085 "itertools.ifilter", /* tp_name */
3086 sizeof(ifilterobject), /* tp_basicsize */
3087 0, /* tp_itemsize */
3088 /* methods */
3089 (destructor)ifilter_dealloc, /* tp_dealloc */
3090 0, /* tp_print */
3091 0, /* tp_getattr */
3092 0, /* tp_setattr */
3093 0, /* tp_compare */
3094 0, /* tp_repr */
3095 0, /* tp_as_number */
3096 0, /* tp_as_sequence */
3097 0, /* tp_as_mapping */
3098 0, /* tp_hash */
3099 0, /* tp_call */
3100 0, /* tp_str */
3101 PyObject_GenericGetAttr, /* tp_getattro */
3102 0, /* tp_setattro */
3103 0, /* tp_as_buffer */
3104 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3105 Py_TPFLAGS_BASETYPE, /* tp_flags */
3106 ifilter_doc, /* tp_doc */
3107 (traverseproc)ifilter_traverse, /* tp_traverse */
3108 0, /* tp_clear */
3109 0, /* tp_richcompare */
3110 0, /* tp_weaklistoffset */
3111 PyObject_SelfIter, /* tp_iter */
3112 (iternextfunc)ifilter_next, /* tp_iternext */
3113 0, /* tp_methods */
3114 0, /* tp_members */
3115 0, /* tp_getset */
3116 0, /* tp_base */
3117 0, /* tp_dict */
3118 0, /* tp_descr_get */
3119 0, /* tp_descr_set */
3120 0, /* tp_dictoffset */
3121 0, /* tp_init */
3122 0, /* tp_alloc */
3123 ifilter_new, /* tp_new */
3124 PyObject_GC_Del, /* tp_free */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003125};
3126
3127
3128/* ifilterfalse object ************************************************************/
3129
3130typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003131 PyObject_HEAD
3132 PyObject *func;
3133 PyObject *it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003134} ifilterfalseobject;
3135
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003136static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003137
3138static PyObject *
3139ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3140{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003141 PyObject *func, *seq;
3142 PyObject *it;
3143 ifilterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003144
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003145 if (type == &ifilterfalse_type &&
3146 !_PyArg_NoKeywords("ifilterfalse()", kwds))
3147 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003148
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003149 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
3150 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003151
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003152 /* Get iterator. */
3153 it = PyObject_GetIter(seq);
3154 if (it == NULL)
3155 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003156
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003157 /* create ifilterfalseobject structure */
3158 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
3159 if (lz == NULL) {
3160 Py_DECREF(it);
3161 return NULL;
3162 }
3163 Py_INCREF(func);
3164 lz->func = func;
3165 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003166
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003167 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003168}
3169
3170static void
3171ifilterfalse_dealloc(ifilterfalseobject *lz)
3172{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003173 PyObject_GC_UnTrack(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03003174 Py_TRASHCAN_SAFE_BEGIN(lz)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003175 Py_XDECREF(lz->func);
3176 Py_XDECREF(lz->it);
3177 Py_TYPE(lz)->tp_free(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03003178 Py_TRASHCAN_SAFE_END(lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003179}
3180
3181static int
3182ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
3183{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003184 Py_VISIT(lz->it);
3185 Py_VISIT(lz->func);
3186 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003187}
3188
3189static PyObject *
3190ifilterfalse_next(ifilterfalseobject *lz)
3191{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003192 PyObject *item;
3193 PyObject *it = lz->it;
3194 long ok;
3195 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003196
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003197 iternext = *Py_TYPE(it)->tp_iternext;
3198 for (;;) {
Serhiy Storchakabb845652013-04-06 22:04:10 +03003199 if (Py_EnterRecursiveCall(" while iterating"))
3200 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003201 item = iternext(it);
Serhiy Storchakabb845652013-04-06 22:04:10 +03003202 Py_LeaveRecursiveCall();
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003203 if (item == NULL)
3204 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003205
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003206 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3207 ok = PyObject_IsTrue(item);
3208 } else {
3209 PyObject *good;
3210 good = PyObject_CallFunctionObjArgs(lz->func,
3211 item, NULL);
3212 if (good == NULL) {
3213 Py_DECREF(item);
3214 return NULL;
3215 }
3216 ok = PyObject_IsTrue(good);
3217 Py_DECREF(good);
3218 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003219 if (ok == 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003220 return item;
3221 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003222 if (ok < 0)
3223 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003224 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003225}
3226
Raymond Hettinger60eca932003-02-09 06:40:58 +00003227PyDoc_STRVAR(ifilterfalse_doc,
3228"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
3229\n\
3230Return those items of sequence for which function(item) is false.\n\
3231If function is None, return the items that are false.");
3232
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003233static PyTypeObject ifilterfalse_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003234 PyVarObject_HEAD_INIT(NULL, 0)
3235 "itertools.ifilterfalse", /* tp_name */
3236 sizeof(ifilterfalseobject), /* tp_basicsize */
3237 0, /* tp_itemsize */
3238 /* methods */
3239 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
3240 0, /* tp_print */
3241 0, /* tp_getattr */
3242 0, /* tp_setattr */
3243 0, /* tp_compare */
3244 0, /* tp_repr */
3245 0, /* tp_as_number */
3246 0, /* tp_as_sequence */
3247 0, /* tp_as_mapping */
3248 0, /* tp_hash */
3249 0, /* tp_call */
3250 0, /* tp_str */
3251 PyObject_GenericGetAttr, /* tp_getattro */
3252 0, /* tp_setattro */
3253 0, /* tp_as_buffer */
3254 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3255 Py_TPFLAGS_BASETYPE, /* tp_flags */
3256 ifilterfalse_doc, /* tp_doc */
3257 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
3258 0, /* tp_clear */
3259 0, /* tp_richcompare */
3260 0, /* tp_weaklistoffset */
3261 PyObject_SelfIter, /* tp_iter */
3262 (iternextfunc)ifilterfalse_next, /* tp_iternext */
3263 0, /* tp_methods */
3264 0, /* tp_members */
3265 0, /* tp_getset */
3266 0, /* tp_base */
3267 0, /* tp_dict */
3268 0, /* tp_descr_get */
3269 0, /* tp_descr_set */
3270 0, /* tp_dictoffset */
3271 0, /* tp_init */
3272 0, /* tp_alloc */
3273 ifilterfalse_new, /* tp_new */
3274 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003275};
3276
3277
3278/* count object ************************************************************/
3279
3280typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003281 PyObject_HEAD
3282 Py_ssize_t cnt;
3283 PyObject *long_cnt;
3284 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003285} countobject;
3286
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003287/* Counting logic and invariants:
3288
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003289fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003290
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003291 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3292 Advances with: cnt += 1
3293 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003294
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003295slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003296
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003297 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3298 All counting is done with python objects (no overflows or underflows).
3299 Advances with: long_cnt += long_step
3300 Step may be zero -- effectively a slow version of repeat(cnt).
3301 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003302*/
3303
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003304static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003305
3306static PyObject *
3307count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3308{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003309 countobject *lz;
3310 int slow_mode = 0;
3311 Py_ssize_t cnt = 0;
3312 PyObject *long_cnt = NULL;
3313 PyObject *long_step = NULL;
3314 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003315
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003316 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3317 kwlist, &long_cnt, &long_step))
3318 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003319
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003320 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3321 (long_step != NULL && !PyNumber_Check(long_step))) {
3322 PyErr_SetString(PyExc_TypeError, "a number is required");
3323 return NULL;
3324 }
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003325
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003326 if (long_cnt != NULL) {
3327 cnt = PyInt_AsSsize_t(long_cnt);
3328 if ((cnt == -1 && PyErr_Occurred()) || !PyInt_Check(long_cnt)) {
3329 PyErr_Clear();
3330 slow_mode = 1;
3331 }
3332 Py_INCREF(long_cnt);
3333 } else {
3334 cnt = 0;
3335 long_cnt = PyInt_FromLong(0);
3336 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003337
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003338 /* If not specified, step defaults to 1 */
3339 if (long_step == NULL) {
3340 long_step = PyInt_FromLong(1);
3341 if (long_step == NULL) {
3342 Py_DECREF(long_cnt);
3343 return NULL;
3344 }
3345 } else
3346 Py_INCREF(long_step);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003347
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003348 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003349
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003350 /* Fast mode only works when the step is 1 */
3351 if (!PyInt_Check(long_step) ||
3352 PyInt_AS_LONG(long_step) != 1) {
3353 slow_mode = 1;
3354 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003355
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003356 if (slow_mode)
3357 cnt = PY_SSIZE_T_MAX;
3358 else
3359 Py_CLEAR(long_cnt);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003360
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003361 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3362 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3363 assert(slow_mode ||
3364 (PyInt_Check(long_step) && PyInt_AS_LONG(long_step) == 1));
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003365
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003366 /* create countobject structure */
3367 lz = (countobject *)type->tp_alloc(type, 0);
3368 if (lz == NULL) {
3369 Py_XDECREF(long_cnt);
3370 return NULL;
3371 }
3372 lz->cnt = cnt;
3373 lz->long_cnt = long_cnt;
3374 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003375
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003376 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003377}
3378
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003379static void
3380count_dealloc(countobject *lz)
3381{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003382 PyObject_GC_UnTrack(lz);
3383 Py_XDECREF(lz->long_cnt);
3384 Py_XDECREF(lz->long_step);
3385 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003386}
3387
Benjamin Peterson062a7c32009-02-16 21:07:52 +00003388static int
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003389count_traverse(countobject *lz, visitproc visit, void *arg)
3390{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003391 Py_VISIT(lz->long_cnt);
3392 Py_VISIT(lz->long_step);
3393 return 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003394}
3395
3396static PyObject *
3397count_nextlong(countobject *lz)
3398{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003399 PyObject *long_cnt;
3400 PyObject *stepped_up;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003401
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003402 long_cnt = lz->long_cnt;
3403 if (long_cnt == NULL) {
3404 /* Switch to slow_mode */
3405 long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
3406 if (long_cnt == NULL)
3407 return NULL;
3408 }
3409 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003410
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003411 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3412 if (stepped_up == NULL)
3413 return NULL;
3414 lz->long_cnt = stepped_up;
3415 return long_cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003416}
3417
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003418static PyObject *
3419count_next(countobject *lz)
3420{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003421 if (lz->cnt == PY_SSIZE_T_MAX)
3422 return count_nextlong(lz);
3423 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003424}
3425
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003426static PyObject *
3427count_repr(countobject *lz)
3428{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003429 PyObject *cnt_repr, *step_repr = NULL;
3430 PyObject *result = NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003431
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003432 if (lz->cnt != PY_SSIZE_T_MAX)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003433 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003434
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003435 cnt_repr = PyObject_Repr(lz->long_cnt);
3436 if (cnt_repr == NULL)
3437 return NULL;
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003438
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003439 if (PyInt_Check(lz->long_step) && PyInt_AS_LONG(lz->long_step) == 1) {
3440 /* Don't display step when it is an integer equal to 1 */
3441 result = PyString_FromFormat("count(%s)",
3442 PyString_AS_STRING(cnt_repr));
3443 } else {
3444 step_repr = PyObject_Repr(lz->long_step);
3445 if (step_repr != NULL)
3446 result = PyString_FromFormat("count(%s, %s)",
3447 PyString_AS_STRING(cnt_repr),
3448 PyString_AS_STRING(step_repr));
3449 }
3450 Py_DECREF(cnt_repr);
3451 Py_XDECREF(step_repr);
3452 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003453}
3454
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003455static PyObject *
3456count_reduce(countobject *lz)
3457{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003458 if (lz->cnt == PY_SSIZE_T_MAX)
3459 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3460 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003461}
3462
3463PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3464
3465static PyMethodDef count_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003466 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3467 count_reduce_doc},
3468 {NULL, NULL} /* sentinel */
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003469};
3470
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003471PyDoc_STRVAR(count_doc,
Ezio Melottifdc1e0d2010-06-11 02:21:25 +00003472 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003473\n\
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003474Return a count object whose .next() method returns consecutive values.\n\
3475Equivalent to:\n\n\
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003476 def count(firstval=0, step=1):\n\
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003477 x = firstval\n\
3478 while 1:\n\
3479 yield x\n\
3480 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003481
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003482static PyTypeObject count_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003483 PyVarObject_HEAD_INIT(NULL, 0)
3484 "itertools.count", /* tp_name */
3485 sizeof(countobject), /* tp_basicsize */
3486 0, /* tp_itemsize */
3487 /* methods */
3488 (destructor)count_dealloc, /* tp_dealloc */
3489 0, /* tp_print */
3490 0, /* tp_getattr */
3491 0, /* tp_setattr */
3492 0, /* tp_compare */
3493 (reprfunc)count_repr, /* tp_repr */
3494 0, /* tp_as_number */
3495 0, /* tp_as_sequence */
3496 0, /* tp_as_mapping */
3497 0, /* tp_hash */
3498 0, /* tp_call */
3499 0, /* tp_str */
3500 PyObject_GenericGetAttr, /* tp_getattro */
3501 0, /* tp_setattro */
3502 0, /* tp_as_buffer */
3503 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3504 Py_TPFLAGS_BASETYPE, /* tp_flags */
3505 count_doc, /* tp_doc */
3506 (traverseproc)count_traverse, /* tp_traverse */
3507 0, /* tp_clear */
3508 0, /* tp_richcompare */
3509 0, /* tp_weaklistoffset */
3510 PyObject_SelfIter, /* tp_iter */
3511 (iternextfunc)count_next, /* tp_iternext */
3512 count_methods, /* tp_methods */
3513 0, /* tp_members */
3514 0, /* tp_getset */
3515 0, /* tp_base */
3516 0, /* tp_dict */
3517 0, /* tp_descr_get */
3518 0, /* tp_descr_set */
3519 0, /* tp_dictoffset */
3520 0, /* tp_init */
3521 0, /* tp_alloc */
3522 count_new, /* tp_new */
3523 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003524};
3525
3526
3527/* izip object ************************************************************/
3528
3529#include "Python.h"
3530
3531typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003532 PyObject_HEAD
3533 Py_ssize_t tuplesize;
3534 PyObject *ittuple; /* tuple of iterators */
3535 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003536} izipobject;
3537
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003538static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003539
3540static PyObject *
3541izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3542{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003543 izipobject *lz;
3544 Py_ssize_t i;
3545 PyObject *ittuple; /* tuple of iterators */
3546 PyObject *result;
3547 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003548
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003549 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
3550 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003551
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003552 /* args must be a tuple */
3553 assert(PyTuple_Check(args));
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003554
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003555 /* obtain iterators */
3556 ittuple = PyTuple_New(tuplesize);
3557 if (ittuple == NULL)
3558 return NULL;
3559 for (i=0; i < tuplesize; ++i) {
3560 PyObject *item = PyTuple_GET_ITEM(args, i);
3561 PyObject *it = PyObject_GetIter(item);
3562 if (it == NULL) {
3563 if (PyErr_ExceptionMatches(PyExc_TypeError))
3564 PyErr_Format(PyExc_TypeError,
3565 "izip argument #%zd must support iteration",
3566 i+1);
3567 Py_DECREF(ittuple);
3568 return NULL;
3569 }
3570 PyTuple_SET_ITEM(ittuple, i, it);
3571 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003572
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003573 /* create a result holder */
3574 result = PyTuple_New(tuplesize);
3575 if (result == NULL) {
3576 Py_DECREF(ittuple);
3577 return NULL;
3578 }
3579 for (i=0 ; i < tuplesize ; i++) {
3580 Py_INCREF(Py_None);
3581 PyTuple_SET_ITEM(result, i, Py_None);
3582 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003583
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003584 /* create izipobject structure */
3585 lz = (izipobject *)type->tp_alloc(type, 0);
3586 if (lz == NULL) {
3587 Py_DECREF(ittuple);
3588 Py_DECREF(result);
3589 return NULL;
3590 }
3591 lz->ittuple = ittuple;
3592 lz->tuplesize = tuplesize;
3593 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003594
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003595 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003596}
3597
3598static void
3599izip_dealloc(izipobject *lz)
3600{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003601 PyObject_GC_UnTrack(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03003602 Py_TRASHCAN_SAFE_BEGIN(lz)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003603 Py_XDECREF(lz->ittuple);
3604 Py_XDECREF(lz->result);
3605 Py_TYPE(lz)->tp_free(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03003606 Py_TRASHCAN_SAFE_END(lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003607}
3608
3609static int
3610izip_traverse(izipobject *lz, visitproc visit, void *arg)
3611{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003612 Py_VISIT(lz->ittuple);
3613 Py_VISIT(lz->result);
3614 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003615}
3616
3617static PyObject *
3618izip_next(izipobject *lz)
3619{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003620 Py_ssize_t i;
3621 Py_ssize_t tuplesize = lz->tuplesize;
3622 PyObject *result = lz->result;
3623 PyObject *it;
3624 PyObject *item;
3625 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003626
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003627 if (tuplesize == 0)
3628 return NULL;
Serhiy Storchakabb845652013-04-06 22:04:10 +03003629 if (Py_EnterRecursiveCall(" while iterating"))
3630 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003631 if (Py_REFCNT(result) == 1) {
3632 Py_INCREF(result);
3633 for (i=0 ; i < tuplesize ; i++) {
3634 it = PyTuple_GET_ITEM(lz->ittuple, i);
3635 item = (*Py_TYPE(it)->tp_iternext)(it);
Serhiy Storchakabb845652013-04-06 22:04:10 +03003636 if (item == NULL)
3637 goto error;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003638 olditem = PyTuple_GET_ITEM(result, i);
3639 PyTuple_SET_ITEM(result, i, item);
3640 Py_DECREF(olditem);
3641 }
3642 } else {
3643 result = PyTuple_New(tuplesize);
3644 if (result == NULL)
Serhiy Storchakabb845652013-04-06 22:04:10 +03003645 goto error;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003646 for (i=0 ; i < tuplesize ; i++) {
3647 it = PyTuple_GET_ITEM(lz->ittuple, i);
3648 item = (*Py_TYPE(it)->tp_iternext)(it);
Serhiy Storchakabb845652013-04-06 22:04:10 +03003649 if (item == NULL)
3650 goto error;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003651 PyTuple_SET_ITEM(result, i, item);
3652 }
3653 }
Serhiy Storchakabb845652013-04-06 22:04:10 +03003654 Py_LeaveRecursiveCall();
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003655 return result;
Serhiy Storchakabb845652013-04-06 22:04:10 +03003656error:
3657 Py_XDECREF(result);
3658 Py_LeaveRecursiveCall();
3659 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003660}
3661
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003662PyDoc_STRVAR(izip_doc,
3663"izip(iter1 [,iter2 [...]]) --> izip object\n\
3664\n\
3665Return a izip object whose .next() method returns a tuple where\n\
3666the i-th element comes from the i-th iterable argument. The .next()\n\
3667method continues until the shortest iterable in the argument sequence\n\
3668is exhausted and then it raises StopIteration. Works like the zip()\n\
3669function but consumes less memory by returning an iterator instead of\n\
3670a list.");
3671
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003672static PyTypeObject izip_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003673 PyVarObject_HEAD_INIT(NULL, 0)
3674 "itertools.izip", /* tp_name */
3675 sizeof(izipobject), /* tp_basicsize */
3676 0, /* tp_itemsize */
3677 /* methods */
3678 (destructor)izip_dealloc, /* tp_dealloc */
3679 0, /* tp_print */
3680 0, /* tp_getattr */
3681 0, /* tp_setattr */
3682 0, /* tp_compare */
3683 0, /* tp_repr */
3684 0, /* tp_as_number */
3685 0, /* tp_as_sequence */
3686 0, /* tp_as_mapping */
3687 0, /* tp_hash */
3688 0, /* tp_call */
3689 0, /* tp_str */
3690 PyObject_GenericGetAttr, /* tp_getattro */
3691 0, /* tp_setattro */
3692 0, /* tp_as_buffer */
3693 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3694 Py_TPFLAGS_BASETYPE, /* tp_flags */
3695 izip_doc, /* tp_doc */
3696 (traverseproc)izip_traverse, /* tp_traverse */
3697 0, /* tp_clear */
3698 0, /* tp_richcompare */
3699 0, /* tp_weaklistoffset */
3700 PyObject_SelfIter, /* tp_iter */
3701 (iternextfunc)izip_next, /* tp_iternext */
3702 0, /* tp_methods */
3703 0, /* tp_members */
3704 0, /* tp_getset */
3705 0, /* tp_base */
3706 0, /* tp_dict */
3707 0, /* tp_descr_get */
3708 0, /* tp_descr_set */
3709 0, /* tp_dictoffset */
3710 0, /* tp_init */
3711 0, /* tp_alloc */
3712 izip_new, /* tp_new */
3713 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003714};
3715
3716
3717/* repeat object ************************************************************/
3718
3719typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003720 PyObject_HEAD
3721 PyObject *element;
3722 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003723} repeatobject;
3724
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003725static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003726
3727static PyObject *
3728repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3729{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003730 repeatobject *ro;
3731 PyObject *element;
3732 Py_ssize_t cnt = -1;
3733 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003734
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003735 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3736 &element, &cnt))
3737 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003738
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003739 if (PyTuple_Size(args) == 2 && cnt < 0)
3740 cnt = 0;
3741
3742 ro = (repeatobject *)type->tp_alloc(type, 0);
3743 if (ro == NULL)
3744 return NULL;
3745 Py_INCREF(element);
3746 ro->element = element;
3747 ro->cnt = cnt;
3748 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003749}
3750
3751static void
3752repeat_dealloc(repeatobject *ro)
3753{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003754 PyObject_GC_UnTrack(ro);
3755 Py_XDECREF(ro->element);
3756 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003757}
3758
3759static int
3760repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3761{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003762 Py_VISIT(ro->element);
3763 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003764}
3765
3766static PyObject *
3767repeat_next(repeatobject *ro)
3768{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003769 if (ro->cnt == 0)
3770 return NULL;
3771 if (ro->cnt > 0)
3772 ro->cnt--;
3773 Py_INCREF(ro->element);
3774 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003775}
3776
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003777static PyObject *
3778repeat_repr(repeatobject *ro)
3779{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003780 PyObject *result, *objrepr;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003781
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003782 objrepr = PyObject_Repr(ro->element);
3783 if (objrepr == NULL)
3784 return NULL;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003785
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003786 if (ro->cnt == -1)
3787 result = PyString_FromFormat("repeat(%s)",
3788 PyString_AS_STRING(objrepr));
3789 else
3790 result = PyString_FromFormat("repeat(%s, %zd)",
3791 PyString_AS_STRING(objrepr), ro->cnt);
3792 Py_DECREF(objrepr);
3793 return result;
3794}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003795
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003796static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003797repeat_len(repeatobject *ro)
3798{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003799 if (ro->cnt == -1) {
3800 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3801 return NULL;
3802 }
3803 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003804}
3805
Armin Rigof5b3e362006-02-11 21:32:43 +00003806PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003807
3808static PyMethodDef repeat_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003809 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3810 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003811};
3812
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003813PyDoc_STRVAR(repeat_doc,
Raymond Hettinger182edae2009-02-19 02:38:25 +00003814"repeat(object [,times]) -> create an iterator which returns the object\n\
3815for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003816endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003817
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003818static PyTypeObject repeat_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003819 PyVarObject_HEAD_INIT(NULL, 0)
3820 "itertools.repeat", /* tp_name */
3821 sizeof(repeatobject), /* tp_basicsize */
3822 0, /* tp_itemsize */
3823 /* methods */
3824 (destructor)repeat_dealloc, /* tp_dealloc */
3825 0, /* tp_print */
3826 0, /* tp_getattr */
3827 0, /* tp_setattr */
3828 0, /* tp_compare */
3829 (reprfunc)repeat_repr, /* tp_repr */
3830 0, /* tp_as_number */
3831 0, /* tp_as_sequence */
3832 0, /* tp_as_mapping */
3833 0, /* tp_hash */
3834 0, /* tp_call */
3835 0, /* tp_str */
3836 PyObject_GenericGetAttr, /* tp_getattro */
3837 0, /* tp_setattro */
3838 0, /* tp_as_buffer */
3839 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3840 Py_TPFLAGS_BASETYPE, /* tp_flags */
3841 repeat_doc, /* tp_doc */
3842 (traverseproc)repeat_traverse, /* tp_traverse */
3843 0, /* tp_clear */
3844 0, /* tp_richcompare */
3845 0, /* tp_weaklistoffset */
3846 PyObject_SelfIter, /* tp_iter */
3847 (iternextfunc)repeat_next, /* tp_iternext */
3848 repeat_methods, /* tp_methods */
3849 0, /* tp_members */
3850 0, /* tp_getset */
3851 0, /* tp_base */
3852 0, /* tp_dict */
3853 0, /* tp_descr_get */
3854 0, /* tp_descr_set */
3855 0, /* tp_dictoffset */
3856 0, /* tp_init */
3857 0, /* tp_alloc */
3858 repeat_new, /* tp_new */
3859 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003860};
3861
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003862/* iziplongest object ************************************************************/
3863
3864#include "Python.h"
3865
3866typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003867 PyObject_HEAD
3868 Py_ssize_t tuplesize;
3869 Py_ssize_t numactive;
3870 PyObject *ittuple; /* tuple of iterators */
3871 PyObject *result;
3872 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003873} iziplongestobject;
3874
3875static PyTypeObject iziplongest_type;
3876
3877static PyObject *
3878izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3879{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003880 iziplongestobject *lz;
3881 Py_ssize_t i;
3882 PyObject *ittuple; /* tuple of iterators */
3883 PyObject *result;
3884 PyObject *fillvalue = Py_None;
3885 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003886
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003887 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3888 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3889 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3890 PyErr_SetString(PyExc_TypeError,
3891 "izip_longest() got an unexpected keyword argument");
3892 return NULL;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003893 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003894 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003895
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003896 /* args must be a tuple */
3897 assert(PyTuple_Check(args));
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003898
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003899 /* obtain iterators */
3900 ittuple = PyTuple_New(tuplesize);
3901 if (ittuple == NULL)
3902 return NULL;
3903 for (i=0; i < tuplesize; ++i) {
3904 PyObject *item = PyTuple_GET_ITEM(args, i);
3905 PyObject *it = PyObject_GetIter(item);
3906 if (it == NULL) {
3907 if (PyErr_ExceptionMatches(PyExc_TypeError))
3908 PyErr_Format(PyExc_TypeError,
3909 "izip_longest argument #%zd must support iteration",
3910 i+1);
3911 Py_DECREF(ittuple);
3912 return NULL;
3913 }
3914 PyTuple_SET_ITEM(ittuple, i, it);
3915 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003916
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003917 /* create a result holder */
3918 result = PyTuple_New(tuplesize);
3919 if (result == NULL) {
3920 Py_DECREF(ittuple);
3921 return NULL;
3922 }
3923 for (i=0 ; i < tuplesize ; i++) {
3924 Py_INCREF(Py_None);
3925 PyTuple_SET_ITEM(result, i, Py_None);
3926 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003927
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003928 /* create iziplongestobject structure */
3929 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3930 if (lz == NULL) {
3931 Py_DECREF(ittuple);
3932 Py_DECREF(result);
3933 return NULL;
3934 }
3935 lz->ittuple = ittuple;
3936 lz->tuplesize = tuplesize;
3937 lz->numactive = tuplesize;
3938 lz->result = result;
3939 Py_INCREF(fillvalue);
3940 lz->fillvalue = fillvalue;
3941 return (PyObject *)lz;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003942}
3943
3944static void
3945izip_longest_dealloc(iziplongestobject *lz)
3946{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003947 PyObject_GC_UnTrack(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03003948 Py_TRASHCAN_SAFE_BEGIN(lz)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003949 Py_XDECREF(lz->ittuple);
3950 Py_XDECREF(lz->result);
3951 Py_XDECREF(lz->fillvalue);
3952 Py_TYPE(lz)->tp_free(lz);
Serhiy Storchakabb845652013-04-06 22:04:10 +03003953 Py_TRASHCAN_SAFE_END(lz)
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003954}
3955
3956static int
3957izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3958{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003959 Py_VISIT(lz->ittuple);
3960 Py_VISIT(lz->result);
3961 Py_VISIT(lz->fillvalue);
3962 return 0;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003963}
3964
3965static PyObject *
3966izip_longest_next(iziplongestobject *lz)
3967{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003968 Py_ssize_t i;
3969 Py_ssize_t tuplesize = lz->tuplesize;
3970 PyObject *result = lz->result;
3971 PyObject *it;
3972 PyObject *item;
3973 PyObject *olditem;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003974
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003975 if (tuplesize == 0)
3976 return NULL;
3977 if (lz->numactive == 0)
3978 return NULL;
3979 if (Py_REFCNT(result) == 1) {
3980 Py_INCREF(result);
3981 for (i=0 ; i < tuplesize ; i++) {
3982 it = PyTuple_GET_ITEM(lz->ittuple, i);
3983 if (it == NULL) {
3984 Py_INCREF(lz->fillvalue);
3985 item = lz->fillvalue;
3986 } else {
3987 item = PyIter_Next(it);
3988 if (item == NULL) {
3989 lz->numactive -= 1;
3990 if (lz->numactive == 0 || PyErr_Occurred()) {
3991 lz->numactive = 0;
3992 Py_DECREF(result);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003993 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003994 } else {
3995 Py_INCREF(lz->fillvalue);
3996 item = lz->fillvalue;
3997 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3998 Py_DECREF(it);
3999 }
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00004000 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004001 }
4002 olditem = PyTuple_GET_ITEM(result, i);
4003 PyTuple_SET_ITEM(result, i, item);
4004 Py_DECREF(olditem);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00004005 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004006 } else {
4007 result = PyTuple_New(tuplesize);
4008 if (result == NULL)
4009 return NULL;
4010 for (i=0 ; i < tuplesize ; i++) {
4011 it = PyTuple_GET_ITEM(lz->ittuple, i);
4012 if (it == NULL) {
4013 Py_INCREF(lz->fillvalue);
4014 item = lz->fillvalue;
4015 } else {
4016 item = PyIter_Next(it);
4017 if (item == NULL) {
4018 lz->numactive -= 1;
4019 if (lz->numactive == 0 || PyErr_Occurred()) {
4020 lz->numactive = 0;
4021 Py_DECREF(result);
4022 return NULL;
4023 } else {
4024 Py_INCREF(lz->fillvalue);
4025 item = lz->fillvalue;
4026 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4027 Py_DECREF(it);
4028 }
4029 }
4030 }
4031 PyTuple_SET_ITEM(result, i, item);
4032 }
4033 }
4034 return result;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00004035}
4036
4037PyDoc_STRVAR(izip_longest_doc,
4038"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
4039\n\
4040Return an izip_longest object whose .next() method returns a tuple where\n\
4041the i-th element comes from the i-th iterable argument. The .next()\n\
4042method continues until the longest iterable in the argument sequence\n\
4043is exhausted and then it raises StopIteration. When the shorter iterables\n\
4044are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4045defaults to None or can be specified by a keyword argument.\n\
4046");
4047
4048static PyTypeObject iziplongest_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004049 PyVarObject_HEAD_INIT(NULL, 0)
4050 "itertools.izip_longest", /* tp_name */
4051 sizeof(iziplongestobject), /* tp_basicsize */
4052 0, /* tp_itemsize */
4053 /* methods */
4054 (destructor)izip_longest_dealloc, /* tp_dealloc */
4055 0, /* tp_print */
4056 0, /* tp_getattr */
4057 0, /* tp_setattr */
4058 0, /* tp_compare */
4059 0, /* tp_repr */
4060 0, /* tp_as_number */
4061 0, /* tp_as_sequence */
4062 0, /* tp_as_mapping */
4063 0, /* tp_hash */
4064 0, /* tp_call */
4065 0, /* tp_str */
4066 PyObject_GenericGetAttr, /* tp_getattro */
4067 0, /* tp_setattro */
4068 0, /* tp_as_buffer */
4069 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4070 Py_TPFLAGS_BASETYPE, /* tp_flags */
4071 izip_longest_doc, /* tp_doc */
4072 (traverseproc)izip_longest_traverse, /* tp_traverse */
4073 0, /* tp_clear */
4074 0, /* tp_richcompare */
4075 0, /* tp_weaklistoffset */
4076 PyObject_SelfIter, /* tp_iter */
4077 (iternextfunc)izip_longest_next, /* tp_iternext */
4078 0, /* tp_methods */
4079 0, /* tp_members */
4080 0, /* tp_getset */
4081 0, /* tp_base */
4082 0, /* tp_dict */
4083 0, /* tp_descr_get */
4084 0, /* tp_descr_set */
4085 0, /* tp_dictoffset */
4086 0, /* tp_init */
4087 0, /* tp_alloc */
4088 izip_longest_new, /* tp_new */
4089 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd36862c2007-02-21 05:20:38 +00004090};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004091
4092/* module level code ********************************************************/
4093
4094PyDoc_STRVAR(module_doc,
4095"Functional tools for creating and using iterators.\n\
4096\n\
4097Infinite iterators:\n\
4098count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004099cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004100repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004101\n\
4102Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004103chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettingere76816b2009-01-29 03:43:44 +00004104compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004105dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4106groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00004107ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
4108ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004109islice(seq, [start,] stop [, step]) --> elements from\n\
4110 seq[start:stop:step]\n\
4111imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
4112starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004113tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004114takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004115izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4116izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4117\n\
4118Combinatoric generators:\n\
4119product(p, q, ... [repeat=1]) --> cartesian product\n\
4120permutations(p[, r])\n\
Raymond Hettinger9eac1192009-11-19 01:22:04 +00004121combinations(p, r)\n\
4122combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004123");
4124
4125
Raymond Hettingerad983e72003-11-12 14:32:26 +00004126static PyMethodDef module_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004127 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4128 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004129};
4130
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004131PyMODINIT_FUNC
4132inititertools(void)
4133{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004134 int i;
4135 PyObject *m;
4136 char *name;
4137 PyTypeObject *typelist[] = {
4138 &combinations_type,
4139 &cwr_type,
4140 &cycle_type,
4141 &dropwhile_type,
4142 &takewhile_type,
4143 &islice_type,
4144 &starmap_type,
4145 &imap_type,
4146 &chain_type,
4147 &compress_type,
4148 &ifilter_type,
4149 &ifilterfalse_type,
4150 &count_type,
4151 &izip_type,
4152 &iziplongest_type,
4153 &permutations_type,
4154 &product_type,
4155 &repeat_type,
4156 &groupby_type,
4157 NULL
4158 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004159
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004160 Py_TYPE(&teedataobject_type) = &PyType_Type;
4161 m = Py_InitModule3("itertools", module_methods, module_doc);
4162 if (m == NULL)
4163 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004164
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004165 for (i=0 ; typelist[i] != NULL ; i++) {
4166 if (PyType_Ready(typelist[i]) < 0)
4167 return;
4168 name = strchr(typelist[i]->tp_name, '.');
4169 assert (name != NULL);
4170 Py_INCREF(typelist[i]);
4171 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4172 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004173
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004174 if (PyType_Ready(&teedataobject_type) < 0)
4175 return;
4176 if (PyType_Ready(&tee_type) < 0)
4177 return;
4178 if (PyType_Ready(&_grouper_type) < 0)
4179 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004180}