blob: 486c2e0a591608f44e5ce3e5a6ed8c3d83d15cb5 [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
5/* Itertools module written and maintained
6 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 {
15 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
21} groupbyobject;
22
23static PyTypeObject groupby_type;
24static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26static PyObject *
27groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28{
Martin v. Löwis02cbf4a2006-02-27 17:20:04 +000029 static char *kwargs[] = {"iterable", "key", NULL};
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000030 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
32
33 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;
51}
52
53static void
54groupby_dealloc(groupbyobject *gbo)
55{
56 PyObject_GC_UnTrack(gbo);
57 Py_XDECREF(gbo->it);
58 Py_XDECREF(gbo->keyfunc);
59 Py_XDECREF(gbo->tgtkey);
60 Py_XDECREF(gbo->currkey);
61 Py_XDECREF(gbo->currvalue);
Christian Heimese93237d2007-12-19 02:37:44 +000062 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000063}
64
65static int
66groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
67{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +000068 Py_VISIT(gbo->it);
69 Py_VISIT(gbo->keyfunc);
70 Py_VISIT(gbo->tgtkey);
71 Py_VISIT(gbo->currkey);
72 Py_VISIT(gbo->currvalue);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000073 return 0;
74}
75
76static PyObject *
77groupby_next(groupbyobject *gbo)
78{
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
81 /* skip to next iteration group */
82 for (;;) {
83 if (gbo->currkey == NULL)
84 /* pass */;
85 else if (gbo->tgtkey == NULL)
86 break;
87 else {
88 int rcmp;
89
90 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
91 gbo->currkey, Py_EQ);
92 if (rcmp == -1)
93 return NULL;
94 else if (rcmp == 0)
95 break;
96 }
97
98 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
101
102 if (gbo->keyfunc == Py_None) {
103 newkey = newvalue;
104 Py_INCREF(newvalue);
105 } else {
106 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
107 newvalue, NULL);
108 if (newkey == NULL) {
109 Py_DECREF(newvalue);
110 return NULL;
111 }
112 }
113
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000114 tmp = gbo->currkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000115 gbo->currkey = newkey;
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000116 Py_XDECREF(tmp);
117
118 tmp = gbo->currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000119 gbo->currvalue = newvalue;
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000120 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000121 }
122
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000123 Py_INCREF(gbo->currkey);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000124 tmp = gbo->tgtkey;
125 gbo->tgtkey = gbo->currkey;
126 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000127
128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
131
132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
135}
136
137PyDoc_STRVAR(groupby_doc,
138"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
139(key, sub-iterator) grouped by each value of key(value).\n");
140
141static PyTypeObject groupby_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000142 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000143 "itertools.groupby", /* tp_name */
144 sizeof(groupbyobject), /* tp_basicsize */
145 0, /* tp_itemsize */
146 /* methods */
147 (destructor)groupby_dealloc, /* tp_dealloc */
148 0, /* tp_print */
149 0, /* tp_getattr */
150 0, /* tp_setattr */
151 0, /* tp_compare */
152 0, /* tp_repr */
153 0, /* tp_as_number */
154 0, /* tp_as_sequence */
155 0, /* tp_as_mapping */
156 0, /* tp_hash */
157 0, /* tp_call */
158 0, /* tp_str */
159 PyObject_GenericGetAttr, /* tp_getattro */
160 0, /* tp_setattro */
161 0, /* tp_as_buffer */
162 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
163 Py_TPFLAGS_BASETYPE, /* tp_flags */
164 groupby_doc, /* tp_doc */
165 (traverseproc)groupby_traverse, /* tp_traverse */
166 0, /* tp_clear */
167 0, /* tp_richcompare */
168 0, /* tp_weaklistoffset */
169 PyObject_SelfIter, /* tp_iter */
170 (iternextfunc)groupby_next, /* tp_iternext */
171 0, /* tp_methods */
172 0, /* tp_members */
173 0, /* tp_getset */
174 0, /* tp_base */
175 0, /* tp_dict */
176 0, /* tp_descr_get */
177 0, /* tp_descr_set */
178 0, /* tp_dictoffset */
179 0, /* tp_init */
180 0, /* tp_alloc */
181 groupby_new, /* tp_new */
182 PyObject_GC_Del, /* tp_free */
183};
184
185
186/* _grouper object (internal) ************************************************/
187
188typedef struct {
189 PyObject_HEAD
190 PyObject *parent;
191 PyObject *tgtkey;
192} _grouperobject;
193
194static PyTypeObject _grouper_type;
195
196static PyObject *
197_grouper_create(groupbyobject *parent, PyObject *tgtkey)
198{
199 _grouperobject *igo;
200
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000201 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000202 if (igo == NULL)
203 return NULL;
204 igo->parent = (PyObject *)parent;
205 Py_INCREF(parent);
206 igo->tgtkey = tgtkey;
207 Py_INCREF(tgtkey);
208
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000209 PyObject_GC_Track(igo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000210 return (PyObject *)igo;
211}
212
213static void
214_grouper_dealloc(_grouperobject *igo)
215{
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000216 PyObject_GC_UnTrack(igo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000217 Py_DECREF(igo->parent);
218 Py_DECREF(igo->tgtkey);
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000219 PyObject_GC_Del(igo);
220}
221
222static int
223_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
224{
225 Py_VISIT(igo->parent);
226 Py_VISIT(igo->tgtkey);
227 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000228}
229
230static PyObject *
231_grouper_next(_grouperobject *igo)
232{
233 groupbyobject *gbo = (groupbyobject *)igo->parent;
234 PyObject *newvalue, *newkey, *r;
235 int rcmp;
236
237 if (gbo->currvalue == NULL) {
238 newvalue = PyIter_Next(gbo->it);
239 if (newvalue == NULL)
240 return NULL;
241
242 if (gbo->keyfunc == Py_None) {
243 newkey = newvalue;
244 Py_INCREF(newvalue);
245 } else {
246 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
247 newvalue, NULL);
248 if (newkey == NULL) {
249 Py_DECREF(newvalue);
250 return NULL;
251 }
252 }
253
254 assert(gbo->currkey == NULL);
255 gbo->currkey = newkey;
256 gbo->currvalue = newvalue;
257 }
258
259 assert(gbo->currkey != NULL);
260 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
261 if (rcmp <= 0)
262 /* got any error or current group is end */
263 return NULL;
264
265 r = gbo->currvalue;
266 gbo->currvalue = NULL;
Raymond Hettinger75ccea32004-09-01 07:02:44 +0000267 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000268
269 return r;
270}
271
272static PyTypeObject _grouper_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000273 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000274 "itertools._grouper", /* tp_name */
275 sizeof(_grouperobject), /* tp_basicsize */
276 0, /* tp_itemsize */
277 /* methods */
278 (destructor)_grouper_dealloc, /* tp_dealloc */
279 0, /* tp_print */
280 0, /* tp_getattr */
281 0, /* tp_setattr */
282 0, /* tp_compare */
283 0, /* tp_repr */
284 0, /* tp_as_number */
285 0, /* tp_as_sequence */
286 0, /* tp_as_mapping */
287 0, /* tp_hash */
288 0, /* tp_call */
289 0, /* tp_str */
290 PyObject_GenericGetAttr, /* tp_getattro */
291 0, /* tp_setattro */
292 0, /* tp_as_buffer */
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000293 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000294 0, /* tp_doc */
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000295 (traverseproc)_grouper_traverse,/* tp_traverse */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000296 0, /* tp_clear */
297 0, /* tp_richcompare */
298 0, /* tp_weaklistoffset */
299 PyObject_SelfIter, /* tp_iter */
300 (iternextfunc)_grouper_next, /* tp_iternext */
301 0, /* tp_methods */
302 0, /* tp_members */
303 0, /* tp_getset */
304 0, /* tp_base */
305 0, /* tp_dict */
306 0, /* tp_descr_get */
307 0, /* tp_descr_set */
308 0, /* tp_dictoffset */
309 0, /* tp_init */
310 0, /* tp_alloc */
311 0, /* tp_new */
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000312 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000313};
314
315
316
Raymond Hettingerad983e72003-11-12 14:32:26 +0000317/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000318
Raymond Hettingerad983e72003-11-12 14:32:26 +0000319/* The teedataobject pre-allocates space for LINKCELLS number of objects.
320 To help the object fit neatly inside cache lines (space for 16 to 32
321 pointers), the value should be a multiple of 16 minus space for
322 the other structure members including PyHEAD overhead. The larger the
323 value, the less memory overhead per object and the less time spent
324 allocating/deallocating new links. The smaller the number, the less
325 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000326*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000327#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000328
329typedef struct {
330 PyObject_HEAD
331 PyObject *it;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000332 int numread;
333 PyObject *nextlink;
334 PyObject *(values[LINKCELLS]);
335} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000336
337typedef struct {
338 PyObject_HEAD
Raymond Hettingerad983e72003-11-12 14:32:26 +0000339 teedataobject *dataobj;
340 int index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000341 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000342} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000343
Raymond Hettingerad983e72003-11-12 14:32:26 +0000344static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000345
346static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000347teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000348{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000349 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000350
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000351 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000352 if (tdo == NULL)
Raymond Hettinger45143692003-10-25 06:37:47 +0000353 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000354
355 tdo->numread = 0;
356 tdo->nextlink = NULL;
357 Py_INCREF(it);
358 tdo->it = it;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000359 PyObject_GC_Track(tdo);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000360 return (PyObject *)tdo;
361}
362
363static PyObject *
364teedataobject_jumplink(teedataobject *tdo)
365{
366 if (tdo->nextlink == NULL)
367 tdo->nextlink = teedataobject_new(tdo->it);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000368 Py_XINCREF(tdo->nextlink);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000369 return tdo->nextlink;
370}
371
372static PyObject *
373teedataobject_getitem(teedataobject *tdo, int i)
374{
375 PyObject *value;
376
377 assert(i < LINKCELLS);
378 if (i < tdo->numread)
379 value = tdo->values[i];
380 else {
381 /* this is the lead iterator, so fetch more data */
382 assert(i == tdo->numread);
383 value = PyIter_Next(tdo->it);
384 if (value == NULL)
385 return NULL;
386 tdo->numread++;
387 tdo->values[i] = value;
Raymond Hettinger45143692003-10-25 06:37:47 +0000388 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000389 Py_INCREF(value);
390 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000391}
392
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000393static int
394teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
395{
396 int i;
397 Py_VISIT(tdo->it);
398 for (i = 0; i < tdo->numread; i++)
399 Py_VISIT(tdo->values[i]);
400 Py_VISIT(tdo->nextlink);
401 return 0;
402}
403
404static int
405teedataobject_clear(teedataobject *tdo)
406{
407 int i;
408 Py_CLEAR(tdo->it);
409 for (i=0 ; i<tdo->numread ; i++)
410 Py_CLEAR(tdo->values[i]);
411 Py_CLEAR(tdo->nextlink);
412 return 0;
413}
414
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000416teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417{
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000418 PyObject_GC_UnTrack(tdo);
419 teedataobject_clear(tdo);
420 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000421}
422
Raymond Hettingerad983e72003-11-12 14:32:26 +0000423PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000424
Raymond Hettingerad983e72003-11-12 14:32:26 +0000425static PyTypeObject teedataobject_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000426 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000427 "itertools.tee_dataobject", /* tp_name */
428 sizeof(teedataobject), /* tp_basicsize */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000429 0, /* tp_itemsize */
430 /* methods */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000431 (destructor)teedataobject_dealloc, /* tp_dealloc */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000432 0, /* tp_print */
433 0, /* tp_getattr */
434 0, /* tp_setattr */
435 0, /* tp_compare */
436 0, /* tp_repr */
437 0, /* tp_as_number */
438 0, /* tp_as_sequence */
439 0, /* tp_as_mapping */
440 0, /* tp_hash */
441 0, /* tp_call */
442 0, /* tp_str */
443 PyObject_GenericGetAttr, /* tp_getattro */
444 0, /* tp_setattro */
445 0, /* tp_as_buffer */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000446 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000447 teedataobject_doc, /* tp_doc */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000448 (traverseproc)teedataobject_traverse, /* tp_traverse */
449 (inquiry)teedataobject_clear, /* tp_clear */
450 0, /* tp_richcompare */
451 0, /* tp_weaklistoffset */
452 0, /* tp_iter */
453 0, /* tp_iternext */
454 0, /* tp_methods */
455 0, /* tp_members */
456 0, /* tp_getset */
457 0, /* tp_base */
458 0, /* tp_dict */
459 0, /* tp_descr_get */
460 0, /* tp_descr_set */
461 0, /* tp_dictoffset */
462 0, /* tp_init */
463 0, /* tp_alloc */
464 0, /* tp_new */
465 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000466};
467
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000468
469static PyTypeObject tee_type;
470
471static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000472tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000473{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000474 PyObject *value, *link;
475
476 if (to->index >= LINKCELLS) {
477 link = teedataobject_jumplink(to->dataobj);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000478 Py_DECREF(to->dataobj);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000479 to->dataobj = (teedataobject *)link;
480 to->index = 0;
481 }
482 value = teedataobject_getitem(to->dataobj, to->index);
483 if (value == NULL)
484 return NULL;
485 to->index++;
486 return value;
487}
488
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000489static int
490tee_traverse(teeobject *to, visitproc visit, void *arg)
491{
492 Py_VISIT((PyObject *)to->dataobj);
493 return 0;
494}
495
Raymond Hettingerad983e72003-11-12 14:32:26 +0000496static PyObject *
497tee_copy(teeobject *to)
498{
499 teeobject *newto;
500
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000501 newto = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000502 if (newto == NULL)
503 return NULL;
504 Py_INCREF(to->dataobj);
505 newto->dataobj = to->dataobj;
506 newto->index = to->index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000507 newto->weakreflist = NULL;
Thomas Woutersb3deb942006-04-15 22:33:13 +0000508 PyObject_GC_Track(newto);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000509 return (PyObject *)newto;
510}
511
512PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
513
514static PyObject *
515tee_fromiterable(PyObject *iterable)
516{
517 teeobject *to;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000518 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000519
520 it = PyObject_GetIter(iterable);
521 if (it == NULL)
522 return NULL;
523 if (PyObject_TypeCheck(it, &tee_type)) {
524 to = (teeobject *)tee_copy((teeobject *)it);
525 goto done;
526 }
527
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000528 to = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000529 if (to == NULL)
530 goto done;
531 to->dataobj = (teedataobject *)teedataobject_new(it);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000532 if (!to->dataobj) {
533 PyObject_GC_Del(to);
534 to = NULL;
535 goto done;
536 }
537
Raymond Hettingerad983e72003-11-12 14:32:26 +0000538 to->index = 0;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000539 to->weakreflist = NULL;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000540 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000541done:
542 Py_XDECREF(it);
543 return (PyObject *)to;
544}
545
546static PyObject *
547tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
548{
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000549 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000550
551 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
552 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000553 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000554}
555
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000556static int
557tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000558{
Raymond Hettingera9f60922004-10-17 16:40:14 +0000559 if (to->weakreflist != NULL)
560 PyObject_ClearWeakRefs((PyObject *) to);
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000561 Py_CLEAR(to->dataobj);
562 return 0;
563}
564
565static void
566tee_dealloc(teeobject *to)
567{
568 PyObject_GC_UnTrack(to);
569 tee_clear(to);
570 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000571}
572
Raymond Hettingerad983e72003-11-12 14:32:26 +0000573PyDoc_STRVAR(teeobject_doc,
574"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000575
Raymond Hettingerad983e72003-11-12 14:32:26 +0000576static PyMethodDef tee_methods[] = {
577 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
578 {NULL, NULL} /* sentinel */
579};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000580
581static PyTypeObject tee_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000582 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000583 "itertools.tee", /* tp_name */
584 sizeof(teeobject), /* tp_basicsize */
585 0, /* tp_itemsize */
586 /* methods */
587 (destructor)tee_dealloc, /* tp_dealloc */
588 0, /* tp_print */
589 0, /* tp_getattr */
590 0, /* tp_setattr */
591 0, /* tp_compare */
592 0, /* tp_repr */
593 0, /* tp_as_number */
594 0, /* tp_as_sequence */
595 0, /* tp_as_mapping */
596 0, /* tp_hash */
597 0, /* tp_call */
598 0, /* tp_str */
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000599 0, /* tp_getattro */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000600 0, /* tp_setattro */
601 0, /* tp_as_buffer */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000602 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000603 teeobject_doc, /* tp_doc */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000604 (traverseproc)tee_traverse, /* tp_traverse */
605 (inquiry)tee_clear, /* tp_clear */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000606 0, /* tp_richcompare */
Raymond Hettingera9f60922004-10-17 16:40:14 +0000607 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000608 PyObject_SelfIter, /* tp_iter */
609 (iternextfunc)tee_next, /* tp_iternext */
610 tee_methods, /* tp_methods */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000611 0, /* tp_members */
612 0, /* tp_getset */
613 0, /* tp_base */
614 0, /* tp_dict */
615 0, /* tp_descr_get */
616 0, /* tp_descr_set */
617 0, /* tp_dictoffset */
618 0, /* tp_init */
619 0, /* tp_alloc */
620 tee_new, /* tp_new */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000621 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000622};
623
Raymond Hettingerad983e72003-11-12 14:32:26 +0000624static PyObject *
625tee(PyObject *self, PyObject *args)
626{
Neal Norwitz69e88972006-09-02 02:58:13 +0000627 Py_ssize_t i, n=2;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000628 PyObject *it, *iterable, *copyable, *result;
629
Neal Norwitz69e88972006-09-02 02:58:13 +0000630 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000631 return NULL;
Neal Norwitz69e88972006-09-02 02:58:13 +0000632 if (n < 0) {
633 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
634 return NULL;
635 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000636 result = PyTuple_New(n);
637 if (result == NULL)
638 return NULL;
639 if (n == 0)
640 return result;
641 it = PyObject_GetIter(iterable);
642 if (it == NULL) {
643 Py_DECREF(result);
644 return NULL;
645 }
646 if (!PyObject_HasAttrString(it, "__copy__")) {
647 copyable = tee_fromiterable(it);
648 Py_DECREF(it);
649 if (copyable == NULL) {
650 Py_DECREF(result);
651 return NULL;
652 }
653 } else
654 copyable = it;
655 PyTuple_SET_ITEM(result, 0, copyable);
656 for (i=1 ; i<n ; i++) {
657 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
658 if (copyable == NULL) {
659 Py_DECREF(result);
660 return NULL;
661 }
662 PyTuple_SET_ITEM(result, i, copyable);
663 }
664 return result;
665}
666
667PyDoc_STRVAR(tee_doc,
668"tee(iterable, n=2) --> tuple of n independent iterators.");
669
670
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000671/* cycle object **********************************************************/
672
673typedef struct {
674 PyObject_HEAD
675 PyObject *it;
676 PyObject *saved;
677 int firstpass;
678} cycleobject;
679
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000680static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000681
682static PyObject *
683cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
684{
685 PyObject *it;
686 PyObject *iterable;
687 PyObject *saved;
688 cycleobject *lz;
689
Georg Brandlb84c1372007-01-21 10:28:43 +0000690 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000691 return NULL;
692
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000693 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
694 return NULL;
695
696 /* Get iterator. */
697 it = PyObject_GetIter(iterable);
698 if (it == NULL)
699 return NULL;
700
701 saved = PyList_New(0);
702 if (saved == NULL) {
703 Py_DECREF(it);
704 return NULL;
705 }
706
707 /* create cycleobject structure */
708 lz = (cycleobject *)type->tp_alloc(type, 0);
709 if (lz == NULL) {
710 Py_DECREF(it);
711 Py_DECREF(saved);
712 return NULL;
713 }
714 lz->it = it;
715 lz->saved = saved;
716 lz->firstpass = 0;
717
718 return (PyObject *)lz;
719}
720
721static void
722cycle_dealloc(cycleobject *lz)
723{
724 PyObject_GC_UnTrack(lz);
725 Py_XDECREF(lz->saved);
726 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +0000727 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000728}
729
730static int
731cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
732{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000733 Py_VISIT(lz->it);
734 Py_VISIT(lz->saved);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000735 return 0;
736}
737
738static PyObject *
739cycle_next(cycleobject *lz)
740{
741 PyObject *item;
742 PyObject *it;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000743 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000744
745 while (1) {
746 item = PyIter_Next(lz->it);
747 if (item != NULL) {
748 if (!lz->firstpass)
749 PyList_Append(lz->saved, item);
750 return item;
751 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000752 if (PyErr_Occurred()) {
753 if (PyErr_ExceptionMatches(PyExc_StopIteration))
754 PyErr_Clear();
755 else
756 return NULL;
757 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000758 if (PyList_Size(lz->saved) == 0)
759 return NULL;
760 it = PyObject_GetIter(lz->saved);
761 if (it == NULL)
762 return NULL;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000763 tmp = lz->it;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000764 lz->it = it;
765 lz->firstpass = 1;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000766 Py_DECREF(tmp);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000767 }
768}
769
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000770PyDoc_STRVAR(cycle_doc,
771"cycle(iterable) --> cycle object\n\
772\n\
773Return elements from the iterable until it is exhausted.\n\
774Then repeat the sequence indefinitely.");
775
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000776static PyTypeObject cycle_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000777 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000778 "itertools.cycle", /* tp_name */
779 sizeof(cycleobject), /* tp_basicsize */
780 0, /* tp_itemsize */
781 /* methods */
782 (destructor)cycle_dealloc, /* tp_dealloc */
783 0, /* tp_print */
784 0, /* tp_getattr */
785 0, /* tp_setattr */
786 0, /* tp_compare */
787 0, /* tp_repr */
788 0, /* tp_as_number */
789 0, /* tp_as_sequence */
790 0, /* tp_as_mapping */
791 0, /* tp_hash */
792 0, /* tp_call */
793 0, /* tp_str */
794 PyObject_GenericGetAttr, /* tp_getattro */
795 0, /* tp_setattro */
796 0, /* tp_as_buffer */
797 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
798 Py_TPFLAGS_BASETYPE, /* tp_flags */
799 cycle_doc, /* tp_doc */
800 (traverseproc)cycle_traverse, /* tp_traverse */
801 0, /* tp_clear */
802 0, /* tp_richcompare */
803 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000804 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000805 (iternextfunc)cycle_next, /* tp_iternext */
806 0, /* tp_methods */
807 0, /* tp_members */
808 0, /* tp_getset */
809 0, /* tp_base */
810 0, /* tp_dict */
811 0, /* tp_descr_get */
812 0, /* tp_descr_set */
813 0, /* tp_dictoffset */
814 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000815 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000816 cycle_new, /* tp_new */
817 PyObject_GC_Del, /* tp_free */
818};
819
820
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000821/* dropwhile object **********************************************************/
822
823typedef struct {
824 PyObject_HEAD
825 PyObject *func;
826 PyObject *it;
827 long start;
828} dropwhileobject;
829
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000830static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000831
832static PyObject *
833dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
834{
835 PyObject *func, *seq;
836 PyObject *it;
837 dropwhileobject *lz;
838
Georg Brandlb84c1372007-01-21 10:28:43 +0000839 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000840 return NULL;
841
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000842 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
843 return NULL;
844
845 /* Get iterator. */
846 it = PyObject_GetIter(seq);
847 if (it == NULL)
848 return NULL;
849
850 /* create dropwhileobject structure */
851 lz = (dropwhileobject *)type->tp_alloc(type, 0);
852 if (lz == NULL) {
853 Py_DECREF(it);
854 return NULL;
855 }
856 Py_INCREF(func);
857 lz->func = func;
858 lz->it = it;
859 lz->start = 0;
860
861 return (PyObject *)lz;
862}
863
864static void
865dropwhile_dealloc(dropwhileobject *lz)
866{
867 PyObject_GC_UnTrack(lz);
868 Py_XDECREF(lz->func);
869 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +0000870 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000871}
872
873static int
874dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
875{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000876 Py_VISIT(lz->it);
877 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000878 return 0;
879}
880
881static PyObject *
882dropwhile_next(dropwhileobject *lz)
883{
884 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000885 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000886 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000887 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000888
Christian Heimese93237d2007-12-19 02:37:44 +0000889 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000890 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000891 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000892 if (item == NULL)
893 return NULL;
894 if (lz->start == 1)
895 return item;
896
897 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
898 if (good == NULL) {
899 Py_DECREF(item);
900 return NULL;
901 }
902 ok = PyObject_IsTrue(good);
903 Py_DECREF(good);
904 if (!ok) {
905 lz->start = 1;
906 return item;
907 }
908 Py_DECREF(item);
909 }
910}
911
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000912PyDoc_STRVAR(dropwhile_doc,
913"dropwhile(predicate, iterable) --> dropwhile object\n\
914\n\
915Drop items from the iterable while predicate(item) is true.\n\
916Afterwards, return every element until the iterable is exhausted.");
917
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000918static PyTypeObject dropwhile_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000919 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000920 "itertools.dropwhile", /* tp_name */
921 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000922 0, /* tp_itemsize */
923 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000924 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000925 0, /* tp_print */
926 0, /* tp_getattr */
927 0, /* tp_setattr */
928 0, /* tp_compare */
929 0, /* tp_repr */
930 0, /* tp_as_number */
931 0, /* tp_as_sequence */
932 0, /* tp_as_mapping */
933 0, /* tp_hash */
934 0, /* tp_call */
935 0, /* tp_str */
936 PyObject_GenericGetAttr, /* tp_getattro */
937 0, /* tp_setattro */
938 0, /* tp_as_buffer */
939 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
940 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000941 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000942 (traverseproc)dropwhile_traverse, /* tp_traverse */
943 0, /* tp_clear */
944 0, /* tp_richcompare */
945 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000946 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000947 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000948 0, /* tp_methods */
949 0, /* tp_members */
950 0, /* tp_getset */
951 0, /* tp_base */
952 0, /* tp_dict */
953 0, /* tp_descr_get */
954 0, /* tp_descr_set */
955 0, /* tp_dictoffset */
956 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000957 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000958 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000959 PyObject_GC_Del, /* tp_free */
960};
961
962
963/* takewhile object **********************************************************/
964
965typedef struct {
966 PyObject_HEAD
967 PyObject *func;
968 PyObject *it;
969 long stop;
970} takewhileobject;
971
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000972static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000973
974static PyObject *
975takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
976{
977 PyObject *func, *seq;
978 PyObject *it;
979 takewhileobject *lz;
980
Georg Brandlb84c1372007-01-21 10:28:43 +0000981 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000982 return NULL;
983
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000984 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
985 return NULL;
986
987 /* Get iterator. */
988 it = PyObject_GetIter(seq);
989 if (it == NULL)
990 return NULL;
991
992 /* create takewhileobject structure */
993 lz = (takewhileobject *)type->tp_alloc(type, 0);
994 if (lz == NULL) {
995 Py_DECREF(it);
996 return NULL;
997 }
998 Py_INCREF(func);
999 lz->func = func;
1000 lz->it = it;
1001 lz->stop = 0;
1002
1003 return (PyObject *)lz;
1004}
1005
1006static void
1007takewhile_dealloc(takewhileobject *lz)
1008{
1009 PyObject_GC_UnTrack(lz);
1010 Py_XDECREF(lz->func);
1011 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00001012 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001013}
1014
1015static int
1016takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1017{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001018 Py_VISIT(lz->it);
1019 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001020 return 0;
1021}
1022
1023static PyObject *
1024takewhile_next(takewhileobject *lz)
1025{
1026 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001027 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001028 long ok;
1029
1030 if (lz->stop == 1)
1031 return NULL;
1032
Christian Heimese93237d2007-12-19 02:37:44 +00001033 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001034 if (item == NULL)
1035 return NULL;
1036
1037 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1038 if (good == NULL) {
1039 Py_DECREF(item);
1040 return NULL;
1041 }
1042 ok = PyObject_IsTrue(good);
1043 Py_DECREF(good);
1044 if (ok)
1045 return item;
1046 Py_DECREF(item);
1047 lz->stop = 1;
1048 return NULL;
1049}
1050
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001051PyDoc_STRVAR(takewhile_doc,
1052"takewhile(predicate, iterable) --> takewhile object\n\
1053\n\
1054Return successive entries from an iterable as long as the \n\
1055predicate evaluates to true for each entry.");
1056
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001057static PyTypeObject takewhile_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001058 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001059 "itertools.takewhile", /* tp_name */
1060 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001061 0, /* tp_itemsize */
1062 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001063 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001064 0, /* tp_print */
1065 0, /* tp_getattr */
1066 0, /* tp_setattr */
1067 0, /* tp_compare */
1068 0, /* tp_repr */
1069 0, /* tp_as_number */
1070 0, /* tp_as_sequence */
1071 0, /* tp_as_mapping */
1072 0, /* tp_hash */
1073 0, /* tp_call */
1074 0, /* tp_str */
1075 PyObject_GenericGetAttr, /* tp_getattro */
1076 0, /* tp_setattro */
1077 0, /* tp_as_buffer */
1078 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1079 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001080 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001081 (traverseproc)takewhile_traverse, /* tp_traverse */
1082 0, /* tp_clear */
1083 0, /* tp_richcompare */
1084 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001085 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001086 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001087 0, /* tp_methods */
1088 0, /* tp_members */
1089 0, /* tp_getset */
1090 0, /* tp_base */
1091 0, /* tp_dict */
1092 0, /* tp_descr_get */
1093 0, /* tp_descr_set */
1094 0, /* tp_dictoffset */
1095 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001096 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001097 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001098 PyObject_GC_Del, /* tp_free */
1099};
1100
1101
1102/* islice object ************************************************************/
1103
1104typedef struct {
1105 PyObject_HEAD
1106 PyObject *it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001107 Py_ssize_t next;
1108 Py_ssize_t stop;
1109 Py_ssize_t step;
1110 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001111} isliceobject;
1112
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001113static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001114
1115static PyObject *
1116islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1117{
1118 PyObject *seq;
Jack Diederich6c433a92006-05-26 11:15:17 +00001119 Py_ssize_t start=0, stop=-1, step=1;
Raymond Hettingerb2594052004-12-05 09:25:51 +00001120 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001121 Py_ssize_t numargs;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001122 isliceobject *lz;
1123
Georg Brandlb84c1372007-01-21 10:28:43 +00001124 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001125 return NULL;
1126
Raymond Hettingerb2594052004-12-05 09:25:51 +00001127 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001128 return NULL;
1129
Raymond Hettingerb2594052004-12-05 09:25:51 +00001130 numargs = PyTuple_Size(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001131 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001132 if (a1 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001133 stop = PyInt_AsSsize_t(a1);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001134 if (stop == -1) {
1135 if (PyErr_Occurred())
1136 PyErr_Clear();
1137 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001138 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001139 return NULL;
1140 }
1141 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001142 } else {
Raymond Hettingerb2594052004-12-05 09:25:51 +00001143 if (a1 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001144 start = PyInt_AsSsize_t(a1);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001145 if (start == -1 && PyErr_Occurred())
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001146 PyErr_Clear();
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001147 if (a2 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001148 stop = PyInt_AsSsize_t(a2);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001149 if (stop == -1) {
1150 if (PyErr_Occurred())
1151 PyErr_Clear();
1152 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001153 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001154 return NULL;
1155 }
1156 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001157 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001158 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001159 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001160 "Indices for islice() must be non-negative integers or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001161 return NULL;
1162 }
1163
Raymond Hettingerb2594052004-12-05 09:25:51 +00001164 if (a3 != NULL) {
1165 if (a3 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001166 step = PyInt_AsSsize_t(a3);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001167 if (step == -1 && PyErr_Occurred())
1168 PyErr_Clear();
1169 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001170 if (step<1) {
1171 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001172 "Step for islice() must be a positive integer or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001173 return NULL;
1174 }
1175
1176 /* Get iterator. */
1177 it = PyObject_GetIter(seq);
1178 if (it == NULL)
1179 return NULL;
1180
1181 /* create isliceobject structure */
1182 lz = (isliceobject *)type->tp_alloc(type, 0);
1183 if (lz == NULL) {
1184 Py_DECREF(it);
1185 return NULL;
1186 }
1187 lz->it = it;
1188 lz->next = start;
1189 lz->stop = stop;
1190 lz->step = step;
1191 lz->cnt = 0L;
1192
1193 return (PyObject *)lz;
1194}
1195
1196static void
1197islice_dealloc(isliceobject *lz)
1198{
1199 PyObject_GC_UnTrack(lz);
1200 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00001201 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001202}
1203
1204static int
1205islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1206{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001207 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001208 return 0;
1209}
1210
1211static PyObject *
1212islice_next(isliceobject *lz)
1213{
1214 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001215 PyObject *it = lz->it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001216 Py_ssize_t oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001217 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001218
Christian Heimese93237d2007-12-19 02:37:44 +00001219 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001220 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001221 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001222 if (item == NULL)
1223 return NULL;
1224 Py_DECREF(item);
1225 lz->cnt++;
1226 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001227 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001228 return NULL;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001229 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001230 if (item == NULL)
1231 return NULL;
1232 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001233 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001234 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001235 if (lz->next < oldnext) /* Check for overflow */
1236 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001237 return item;
1238}
1239
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001240PyDoc_STRVAR(islice_doc,
1241"islice(iterable, [start,] stop [, step]) --> islice object\n\
1242\n\
1243Return an iterator whose next() method returns selected values from an\n\
1244iterable. If start is specified, will skip all preceding elements;\n\
1245otherwise, start defaults to zero. Step defaults to one. If\n\
1246specified as another value, step determines how many values are \n\
1247skipped between successive calls. Works like a slice() on a list\n\
1248but returns an iterator.");
1249
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001250static PyTypeObject islice_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001251 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001252 "itertools.islice", /* tp_name */
1253 sizeof(isliceobject), /* tp_basicsize */
1254 0, /* tp_itemsize */
1255 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001256 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001257 0, /* tp_print */
1258 0, /* tp_getattr */
1259 0, /* tp_setattr */
1260 0, /* tp_compare */
1261 0, /* tp_repr */
1262 0, /* tp_as_number */
1263 0, /* tp_as_sequence */
1264 0, /* tp_as_mapping */
1265 0, /* tp_hash */
1266 0, /* tp_call */
1267 0, /* tp_str */
1268 PyObject_GenericGetAttr, /* tp_getattro */
1269 0, /* tp_setattro */
1270 0, /* tp_as_buffer */
1271 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1272 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001273 islice_doc, /* tp_doc */
1274 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001275 0, /* tp_clear */
1276 0, /* tp_richcompare */
1277 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001278 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001279 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001280 0, /* tp_methods */
1281 0, /* tp_members */
1282 0, /* tp_getset */
1283 0, /* tp_base */
1284 0, /* tp_dict */
1285 0, /* tp_descr_get */
1286 0, /* tp_descr_set */
1287 0, /* tp_dictoffset */
1288 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001289 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001290 islice_new, /* tp_new */
1291 PyObject_GC_Del, /* tp_free */
1292};
1293
1294
1295/* starmap object ************************************************************/
1296
1297typedef struct {
1298 PyObject_HEAD
1299 PyObject *func;
1300 PyObject *it;
1301} starmapobject;
1302
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001303static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001304
1305static PyObject *
1306starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1307{
1308 PyObject *func, *seq;
1309 PyObject *it;
1310 starmapobject *lz;
1311
Georg Brandlb84c1372007-01-21 10:28:43 +00001312 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001313 return NULL;
1314
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001315 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1316 return NULL;
1317
1318 /* Get iterator. */
1319 it = PyObject_GetIter(seq);
1320 if (it == NULL)
1321 return NULL;
1322
1323 /* create starmapobject structure */
1324 lz = (starmapobject *)type->tp_alloc(type, 0);
1325 if (lz == NULL) {
1326 Py_DECREF(it);
1327 return NULL;
1328 }
1329 Py_INCREF(func);
1330 lz->func = func;
1331 lz->it = it;
1332
1333 return (PyObject *)lz;
1334}
1335
1336static void
1337starmap_dealloc(starmapobject *lz)
1338{
1339 PyObject_GC_UnTrack(lz);
1340 Py_XDECREF(lz->func);
1341 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00001342 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001343}
1344
1345static int
1346starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1347{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001348 Py_VISIT(lz->it);
1349 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001350 return 0;
1351}
1352
1353static PyObject *
1354starmap_next(starmapobject *lz)
1355{
1356 PyObject *args;
1357 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001358 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001359
Christian Heimese93237d2007-12-19 02:37:44 +00001360 args = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001361 if (args == NULL)
1362 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001363 if (!PyTuple_CheckExact(args)) {
Raymond Hettinger47317092008-01-17 03:02:14 +00001364 PyObject *newargs = PySequence_Tuple(args);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001365 Py_DECREF(args);
Raymond Hettinger47317092008-01-17 03:02:14 +00001366 if (newargs == NULL)
1367 return NULL;
1368 args = newargs;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001369 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001370 result = PyObject_Call(lz->func, args, NULL);
1371 Py_DECREF(args);
1372 return result;
1373}
1374
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001375PyDoc_STRVAR(starmap_doc,
1376"starmap(function, sequence) --> starmap object\n\
1377\n\
1378Return an iterator whose values are returned from the function evaluated\n\
1379with a argument tuple taken from the given sequence.");
1380
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001381static PyTypeObject starmap_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001382 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001383 "itertools.starmap", /* tp_name */
1384 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001385 0, /* tp_itemsize */
1386 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001387 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001388 0, /* tp_print */
1389 0, /* tp_getattr */
1390 0, /* tp_setattr */
1391 0, /* tp_compare */
1392 0, /* tp_repr */
1393 0, /* tp_as_number */
1394 0, /* tp_as_sequence */
1395 0, /* tp_as_mapping */
1396 0, /* tp_hash */
1397 0, /* tp_call */
1398 0, /* tp_str */
1399 PyObject_GenericGetAttr, /* tp_getattro */
1400 0, /* tp_setattro */
1401 0, /* tp_as_buffer */
1402 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1403 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001404 starmap_doc, /* tp_doc */
1405 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001406 0, /* tp_clear */
1407 0, /* tp_richcompare */
1408 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001409 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001410 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001411 0, /* tp_methods */
1412 0, /* tp_members */
1413 0, /* tp_getset */
1414 0, /* tp_base */
1415 0, /* tp_dict */
1416 0, /* tp_descr_get */
1417 0, /* tp_descr_set */
1418 0, /* tp_dictoffset */
1419 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001420 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001421 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001422 PyObject_GC_Del, /* tp_free */
1423};
1424
1425
1426/* imap object ************************************************************/
1427
1428typedef struct {
1429 PyObject_HEAD
1430 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001431 PyObject *func;
1432} imapobject;
1433
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001434static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001435
1436static PyObject *
1437imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1438{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001439 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001440 imapobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001441 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001442
Georg Brandlb84c1372007-01-21 10:28:43 +00001443 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001444 return NULL;
1445
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001446 numargs = PyTuple_Size(args);
1447 if (numargs < 2) {
1448 PyErr_SetString(PyExc_TypeError,
1449 "imap() must have at least two arguments.");
1450 return NULL;
1451 }
1452
1453 iters = PyTuple_New(numargs-1);
1454 if (iters == NULL)
1455 return NULL;
1456
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457 for (i=1 ; i<numargs ; i++) {
1458 /* Get iterator. */
1459 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1460 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001461 Py_DECREF(iters);
1462 return NULL;
1463 }
1464 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001465 }
1466
1467 /* create imapobject structure */
1468 lz = (imapobject *)type->tp_alloc(type, 0);
1469 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001470 Py_DECREF(iters);
1471 return NULL;
1472 }
1473 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001474 func = PyTuple_GET_ITEM(args, 0);
1475 Py_INCREF(func);
1476 lz->func = func;
1477
1478 return (PyObject *)lz;
1479}
1480
1481static void
1482imap_dealloc(imapobject *lz)
1483{
1484 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001485 Py_XDECREF(lz->iters);
1486 Py_XDECREF(lz->func);
Christian Heimese93237d2007-12-19 02:37:44 +00001487 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001488}
1489
1490static int
1491imap_traverse(imapobject *lz, visitproc visit, void *arg)
1492{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001493 Py_VISIT(lz->iters);
1494 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001495 return 0;
1496}
1497
Raymond Hettinger2012f172003-02-07 05:32:58 +00001498/*
1499imap() is an iterator version of __builtins__.map() except that it does
1500not have the None fill-in feature. That was intentionally left out for
1501the following reasons:
1502
1503 1) Itertools are designed to be easily combined and chained together.
1504 Having all tools stop with the shortest input is a unifying principle
1505 that makes it easier to combine finite iterators (supplying data) with
1506 infinite iterators like count() and repeat() (for supplying sequential
1507 or constant arguments to a function).
1508
1509 2) In typical use cases for combining itertools, having one finite data
1510 supplier run out before another is likely to be an error condition which
1511 should not pass silently by automatically supplying None.
1512
1513 3) The use cases for automatic None fill-in are rare -- not many functions
1514 do something useful when a parameter suddenly switches type and becomes
1515 None.
1516
1517 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001518 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001519
1520 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1521*/
1522
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001523static PyObject *
1524imap_next(imapobject *lz)
1525{
1526 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001527 PyObject *argtuple;
1528 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001529 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001530
1531 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001532 argtuple = PyTuple_New(numargs);
1533 if (argtuple == NULL)
1534 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001535
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001536 for (i=0 ; i<numargs ; i++) {
1537 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1538 if (val == NULL) {
1539 Py_DECREF(argtuple);
1540 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001541 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001542 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001543 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001544 if (lz->func == Py_None)
1545 return argtuple;
1546 result = PyObject_Call(lz->func, argtuple, NULL);
1547 Py_DECREF(argtuple);
1548 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001549}
1550
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001551PyDoc_STRVAR(imap_doc,
1552"imap(func, *iterables) --> imap object\n\
1553\n\
1554Make an iterator that computes the function using arguments from\n\
1555each of the iterables. Like map() except that it returns\n\
1556an iterator instead of a list and that it stops when the shortest\n\
1557iterable is exhausted instead of filling in None for shorter\n\
1558iterables.");
1559
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001560static PyTypeObject imap_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001561 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001562 "itertools.imap", /* tp_name */
1563 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001564 0, /* tp_itemsize */
1565 /* methods */
1566 (destructor)imap_dealloc, /* tp_dealloc */
1567 0, /* tp_print */
1568 0, /* tp_getattr */
1569 0, /* tp_setattr */
1570 0, /* tp_compare */
1571 0, /* tp_repr */
1572 0, /* tp_as_number */
1573 0, /* tp_as_sequence */
1574 0, /* tp_as_mapping */
1575 0, /* tp_hash */
1576 0, /* tp_call */
1577 0, /* tp_str */
1578 PyObject_GenericGetAttr, /* tp_getattro */
1579 0, /* tp_setattro */
1580 0, /* tp_as_buffer */
1581 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1582 Py_TPFLAGS_BASETYPE, /* tp_flags */
1583 imap_doc, /* tp_doc */
1584 (traverseproc)imap_traverse, /* tp_traverse */
1585 0, /* tp_clear */
1586 0, /* tp_richcompare */
1587 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001588 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001589 (iternextfunc)imap_next, /* tp_iternext */
1590 0, /* tp_methods */
1591 0, /* tp_members */
1592 0, /* tp_getset */
1593 0, /* tp_base */
1594 0, /* tp_dict */
1595 0, /* tp_descr_get */
1596 0, /* tp_descr_set */
1597 0, /* tp_dictoffset */
1598 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001599 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001600 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001601 PyObject_GC_Del, /* tp_free */
1602};
1603
1604
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001605/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001606
1607typedef struct {
1608 PyObject_HEAD
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001609 PyObject *source; /* Iterator over input iterables */
1610 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001611} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001612
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001613static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001614
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001615static PyObject *
1616chain_new_internal(PyTypeObject *type, PyObject *source)
1617{
1618 chainobject *lz;
1619
1620 lz = (chainobject *)type->tp_alloc(type, 0);
1621 if (lz == NULL) {
1622 Py_DECREF(source);
1623 return NULL;
1624 }
1625
1626 lz->source = source;
1627 lz->active = NULL;
1628 return (PyObject *)lz;
1629}
1630
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001631static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001632chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001633{
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001634 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001635
Georg Brandlb84c1372007-01-21 10:28:43 +00001636 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001637 return NULL;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001638
1639 source = PyObject_GetIter(args);
1640 if (source == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001641 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001642
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001643 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001644}
1645
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001646static PyObject *
1647chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1648{
1649 PyObject *source;
1650
1651 source = PyObject_GetIter(arg);
1652 if (source == NULL)
1653 return NULL;
1654
1655 return chain_new_internal(type, source);
1656}
1657
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001658static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001659chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001660{
1661 PyObject_GC_UnTrack(lz);
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001662 Py_XDECREF(lz->active);
1663 Py_XDECREF(lz->source);
Christian Heimese93237d2007-12-19 02:37:44 +00001664 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001665}
1666
Raymond Hettinger2012f172003-02-07 05:32:58 +00001667static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001668chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001669{
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001670 Py_VISIT(lz->source);
1671 Py_VISIT(lz->active);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001672 return 0;
1673}
1674
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001675static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001676chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001677{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001678 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001679
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001680 if (lz->source == NULL)
1681 return NULL; /* already stopped */
1682
1683 if (lz->active == NULL) {
1684 PyObject *iterable = PyIter_Next(lz->source);
1685 if (iterable == NULL) {
1686 Py_CLEAR(lz->source);
1687 return NULL; /* no more input sources */
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001688 }
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001689 lz->active = PyObject_GetIter(iterable);
Raymond Hettingerf1cca2b2008-03-04 22:29:44 +00001690 Py_DECREF(iterable);
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001691 if (lz->active == NULL) {
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001692 Py_CLEAR(lz->source);
1693 return NULL; /* input not iterable */
1694 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001695 }
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001696 item = PyIter_Next(lz->active);
1697 if (item != NULL)
1698 return item;
1699 if (PyErr_Occurred()) {
1700 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1701 PyErr_Clear();
1702 else
1703 return NULL; /* input raised an exception */
1704 }
1705 Py_CLEAR(lz->active);
1706 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001707}
1708
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001709PyDoc_STRVAR(chain_doc,
1710"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001711\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001712Return a chain object whose .next() method returns elements from the\n\
1713first iterable until it is exhausted, then elements from the next\n\
1714iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001715
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001716PyDoc_STRVAR(chain_from_iterable_doc,
1717"chain.from_iterable(iterable) --> chain object\n\
1718\n\
1719Alternate chain() contructor taking a single iterable argument\n\
1720that evaluates lazily.");
1721
1722static PyMethodDef chain_methods[] = {
1723 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1724 chain_from_iterable_doc},
1725 {NULL, NULL} /* sentinel */
1726};
1727
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001728static PyTypeObject chain_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001729 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001730 "itertools.chain", /* tp_name */
1731 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001732 0, /* tp_itemsize */
1733 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001734 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001735 0, /* tp_print */
1736 0, /* tp_getattr */
1737 0, /* tp_setattr */
1738 0, /* tp_compare */
1739 0, /* tp_repr */
1740 0, /* tp_as_number */
1741 0, /* tp_as_sequence */
1742 0, /* tp_as_mapping */
1743 0, /* tp_hash */
1744 0, /* tp_call */
1745 0, /* tp_str */
1746 PyObject_GenericGetAttr, /* tp_getattro */
1747 0, /* tp_setattro */
1748 0, /* tp_as_buffer */
1749 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1750 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001751 chain_doc, /* tp_doc */
1752 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001753 0, /* tp_clear */
1754 0, /* tp_richcompare */
1755 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001756 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001757 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001758 chain_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001759 0, /* tp_members */
1760 0, /* tp_getset */
1761 0, /* tp_base */
1762 0, /* tp_dict */
1763 0, /* tp_descr_get */
1764 0, /* tp_descr_set */
1765 0, /* tp_dictoffset */
1766 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001767 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001768 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001769 PyObject_GC_Del, /* tp_free */
1770};
1771
1772
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001773/* product object ************************************************************/
1774
1775typedef struct {
1776 PyObject_HEAD
1777 PyObject *pools; /* tuple of pool tuples */
Raymond Hettinger532316d2008-02-23 04:03:50 +00001778 Py_ssize_t *indices; /* one index per pool */
1779 PyObject *result; /* most recently returned result tuple */
1780 int stopped; /* set to 1 when the product iterator is exhausted */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001781} productobject;
1782
1783static PyTypeObject product_type;
1784
1785static PyObject *
1786product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1787{
1788 productobject *lz;
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001789 Py_ssize_t nargs, npools, repeat=1;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001790 PyObject *pools = NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001791 Py_ssize_t *indices = NULL;
1792 Py_ssize_t i;
1793
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001794 if (kwds != NULL) {
1795 char *kwlist[] = {"repeat", 0};
1796 PyObject *tmpargs = PyTuple_New(0);
1797 if (tmpargs == NULL)
1798 return NULL;
1799 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1800 Py_DECREF(tmpargs);
1801 return NULL;
1802 }
1803 Py_DECREF(tmpargs);
1804 if (repeat < 0) {
1805 PyErr_SetString(PyExc_ValueError,
1806 "repeat argument cannot be negative");
1807 return NULL;
1808 }
1809 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001810
1811 assert(PyTuple_Check(args));
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001812 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1813 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001814
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001815 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
Raymond Hettinger61024b92008-03-02 11:57:16 +00001816 if (indices == NULL) {
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001817 PyErr_NoMemory();
1818 goto error;
1819 }
1820
1821 pools = PyTuple_New(npools);
1822 if (pools == NULL)
1823 goto error;
1824
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001825 for (i=0; i < nargs ; ++i) {
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001826 PyObject *item = PyTuple_GET_ITEM(args, i);
1827 PyObject *pool = PySequence_Tuple(item);
1828 if (pool == NULL)
1829 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001830 PyTuple_SET_ITEM(pools, i, pool);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001831 indices[i] = 0;
1832 }
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001833 for ( ; i < npools; ++i) {
1834 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1835 Py_INCREF(pool);
1836 PyTuple_SET_ITEM(pools, i, pool);
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001837 indices[i] = 0;
1838 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001839
1840 /* create productobject structure */
1841 lz = (productobject *)type->tp_alloc(type, 0);
Raymond Hettinger3bd77122008-02-27 01:08:04 +00001842 if (lz == NULL)
Raymond Hettinger73d79632008-02-23 02:20:41 +00001843 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001844
1845 lz->pools = pools;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001846 lz->indices = indices;
1847 lz->result = NULL;
1848 lz->stopped = 0;
1849
1850 return (PyObject *)lz;
1851
1852error:
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001853 if (indices != NULL)
1854 PyMem_Free(indices);
1855 Py_XDECREF(pools);
1856 return NULL;
1857}
1858
1859static void
1860product_dealloc(productobject *lz)
1861{
1862 PyObject_GC_UnTrack(lz);
1863 Py_XDECREF(lz->pools);
1864 Py_XDECREF(lz->result);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00001865 if (lz->indices != NULL)
1866 PyMem_Free(lz->indices);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001867 Py_TYPE(lz)->tp_free(lz);
1868}
1869
1870static int
1871product_traverse(productobject *lz, visitproc visit, void *arg)
1872{
1873 Py_VISIT(lz->pools);
1874 Py_VISIT(lz->result);
1875 return 0;
1876}
1877
1878static PyObject *
1879product_next(productobject *lz)
1880{
1881 PyObject *pool;
1882 PyObject *elem;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001883 PyObject *oldelem;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001884 PyObject *pools = lz->pools;
1885 PyObject *result = lz->result;
1886 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1887 Py_ssize_t i;
1888
1889 if (lz->stopped)
1890 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001891
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001892 if (result == NULL) {
Raymond Hettinger73d79632008-02-23 02:20:41 +00001893 /* On the first pass, return an initial tuple filled with the
Raymond Hettingerd553d852008-03-04 04:17:08 +00001894 first element from each pool. */
Raymond Hettinger73d79632008-02-23 02:20:41 +00001895 result = PyTuple_New(npools);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001896 if (result == NULL)
1897 goto empty;
1898 lz->result = result;
1899 for (i=0; i < npools; i++) {
1900 pool = PyTuple_GET_ITEM(pools, i);
1901 if (PyTuple_GET_SIZE(pool) == 0)
1902 goto empty;
1903 elem = PyTuple_GET_ITEM(pool, 0);
1904 Py_INCREF(elem);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001905 PyTuple_SET_ITEM(result, i, elem);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001906 }
1907 } else {
1908 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001909
1910 /* Copy the previous result tuple or re-use it if available */
1911 if (Py_REFCNT(result) > 1) {
1912 PyObject *old_result = result;
1913 result = PyTuple_New(npools);
1914 if (result == NULL)
1915 goto empty;
1916 lz->result = result;
1917 for (i=0; i < npools; i++) {
1918 elem = PyTuple_GET_ITEM(old_result, i);
1919 Py_INCREF(elem);
1920 PyTuple_SET_ITEM(result, i, elem);
1921 }
1922 Py_DECREF(old_result);
1923 }
1924 /* Now, we've got the only copy so we can update it in-place */
Raymond Hettingere3fabd12008-03-02 12:02:19 +00001925 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001926
1927 /* Update the pool indices right-to-left. Only advance to the
1928 next pool when the previous one rolls-over */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001929 for (i=npools-1 ; i >= 0 ; i--) {
1930 pool = PyTuple_GET_ITEM(pools, i);
1931 indices[i]++;
Raymond Hettinger61024b92008-03-02 11:57:16 +00001932 if (indices[i] == PyTuple_GET_SIZE(pool)) {
Raymond Hettinger73d79632008-02-23 02:20:41 +00001933 /* Roll-over and advance to next pool */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001934 indices[i] = 0;
1935 elem = PyTuple_GET_ITEM(pool, 0);
1936 Py_INCREF(elem);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001937 oldelem = PyTuple_GET_ITEM(result, i);
1938 PyTuple_SET_ITEM(result, i, elem);
1939 Py_DECREF(oldelem);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001940 } else {
Raymond Hettinger73d79632008-02-23 02:20:41 +00001941 /* No rollover. Just increment and stop here. */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001942 elem = PyTuple_GET_ITEM(pool, indices[i]);
1943 Py_INCREF(elem);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001944 oldelem = PyTuple_GET_ITEM(result, i);
1945 PyTuple_SET_ITEM(result, i, elem);
1946 Py_DECREF(oldelem);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001947 break;
1948 }
1949 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00001950
1951 /* If i is negative, then the indices have all rolled-over
1952 and we're done. */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001953 if (i < 0)
Raymond Hettinger73d79632008-02-23 02:20:41 +00001954 goto empty;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001955 }
1956
Raymond Hettinger73d79632008-02-23 02:20:41 +00001957 Py_INCREF(result);
1958 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001959
1960empty:
1961 lz->stopped = 1;
1962 return NULL;
1963}
1964
1965PyDoc_STRVAR(product_doc,
1966"product(*iterables) --> product object\n\
1967\n\
Raymond Hettinger73d79632008-02-23 02:20:41 +00001968Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001969For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1970The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1971cycle in a manner similar to an odometer (with the rightmost element changing\n\
1972on every iteration).\n\n\
Raymond Hettinger94034ea2009-02-09 18:39:41 +00001973To compute the product of an iterable with itself, specify the number\n\
1974of repetitions with the optional repeat keyword argument. For example,\n\
1975product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001976product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1977product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1978
1979static PyTypeObject product_type = {
1980 PyVarObject_HEAD_INIT(NULL, 0)
1981 "itertools.product", /* tp_name */
1982 sizeof(productobject), /* tp_basicsize */
1983 0, /* tp_itemsize */
1984 /* methods */
1985 (destructor)product_dealloc, /* tp_dealloc */
1986 0, /* tp_print */
1987 0, /* tp_getattr */
1988 0, /* tp_setattr */
1989 0, /* tp_compare */
1990 0, /* tp_repr */
1991 0, /* tp_as_number */
1992 0, /* tp_as_sequence */
1993 0, /* tp_as_mapping */
1994 0, /* tp_hash */
1995 0, /* tp_call */
1996 0, /* tp_str */
1997 PyObject_GenericGetAttr, /* tp_getattro */
1998 0, /* tp_setattro */
1999 0, /* tp_as_buffer */
2000 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2001 Py_TPFLAGS_BASETYPE, /* tp_flags */
2002 product_doc, /* tp_doc */
2003 (traverseproc)product_traverse, /* tp_traverse */
2004 0, /* tp_clear */
2005 0, /* tp_richcompare */
2006 0, /* tp_weaklistoffset */
2007 PyObject_SelfIter, /* tp_iter */
2008 (iternextfunc)product_next, /* tp_iternext */
2009 0, /* tp_methods */
2010 0, /* tp_members */
2011 0, /* tp_getset */
2012 0, /* tp_base */
2013 0, /* tp_dict */
2014 0, /* tp_descr_get */
2015 0, /* tp_descr_set */
2016 0, /* tp_dictoffset */
2017 0, /* tp_init */
2018 0, /* tp_alloc */
2019 product_new, /* tp_new */
2020 PyObject_GC_Del, /* tp_free */
2021};
2022
2023
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002024/* combinations object ************************************************************/
2025
2026typedef struct {
2027 PyObject_HEAD
2028 PyObject *pool; /* input converted to a tuple */
2029 Py_ssize_t *indices; /* one index per result element */
2030 PyObject *result; /* most recently returned result tuple */
2031 Py_ssize_t r; /* size of result tuple */
2032 int stopped; /* set to 1 when the combinations iterator is exhausted */
2033} combinationsobject;
2034
2035static PyTypeObject combinations_type;
2036
2037static PyObject *
2038combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2039{
2040 combinationsobject *co;
2041 Py_ssize_t n;
2042 Py_ssize_t r;
2043 PyObject *pool = NULL;
2044 PyObject *iterable = NULL;
2045 Py_ssize_t *indices = NULL;
2046 Py_ssize_t i;
2047 static char *kwargs[] = {"iterable", "r", NULL};
2048
2049 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2050 &iterable, &r))
2051 return NULL;
2052
2053 pool = PySequence_Tuple(iterable);
2054 if (pool == NULL)
2055 goto error;
2056 n = PyTuple_GET_SIZE(pool);
2057 if (r < 0) {
2058 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2059 goto error;
2060 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002061
2062 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2063 if (indices == NULL) {
2064 PyErr_NoMemory();
2065 goto error;
2066 }
2067
2068 for (i=0 ; i<r ; i++)
2069 indices[i] = i;
2070
2071 /* create combinationsobject structure */
2072 co = (combinationsobject *)type->tp_alloc(type, 0);
2073 if (co == NULL)
2074 goto error;
2075
2076 co->pool = pool;
2077 co->indices = indices;
2078 co->result = NULL;
2079 co->r = r;
Raymond Hettinger5b913e32009-01-08 06:39:04 +00002080 co->stopped = r > n ? 1 : 0;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002081
2082 return (PyObject *)co;
2083
2084error:
2085 if (indices != NULL)
2086 PyMem_Free(indices);
2087 Py_XDECREF(pool);
2088 return NULL;
2089}
2090
2091static void
2092combinations_dealloc(combinationsobject *co)
2093{
2094 PyObject_GC_UnTrack(co);
2095 Py_XDECREF(co->pool);
2096 Py_XDECREF(co->result);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002097 if (co->indices != NULL)
2098 PyMem_Free(co->indices);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002099 Py_TYPE(co)->tp_free(co);
2100}
2101
2102static int
2103combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2104{
2105 Py_VISIT(co->pool);
2106 Py_VISIT(co->result);
2107 return 0;
2108}
2109
2110static PyObject *
2111combinations_next(combinationsobject *co)
2112{
2113 PyObject *elem;
2114 PyObject *oldelem;
2115 PyObject *pool = co->pool;
2116 Py_ssize_t *indices = co->indices;
2117 PyObject *result = co->result;
2118 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2119 Py_ssize_t r = co->r;
2120 Py_ssize_t i, j, index;
2121
2122 if (co->stopped)
2123 return NULL;
2124
2125 if (result == NULL) {
2126 /* On the first pass, initialize result tuple using the indices */
2127 result = PyTuple_New(r);
2128 if (result == NULL)
2129 goto empty;
2130 co->result = result;
2131 for (i=0; i<r ; i++) {
2132 index = indices[i];
2133 elem = PyTuple_GET_ITEM(pool, index);
2134 Py_INCREF(elem);
2135 PyTuple_SET_ITEM(result, i, elem);
2136 }
2137 } else {
2138 /* Copy the previous result tuple or re-use it if available */
2139 if (Py_REFCNT(result) > 1) {
2140 PyObject *old_result = result;
2141 result = PyTuple_New(r);
2142 if (result == NULL)
2143 goto empty;
2144 co->result = result;
2145 for (i=0; i<r ; i++) {
2146 elem = PyTuple_GET_ITEM(old_result, i);
2147 Py_INCREF(elem);
2148 PyTuple_SET_ITEM(result, i, elem);
2149 }
2150 Py_DECREF(old_result);
2151 }
Christian Heimescdddf182008-02-28 11:18:49 +00002152 /* Now, we've got the only copy so we can update it in-place
2153 * CPython's empty tuple is a singleton and cached in
2154 * PyTuple's freelist.
2155 */
2156 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002157
2158 /* Scan indices right-to-left until finding one that is not
2159 at its maximum (i + n - r). */
2160 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2161 ;
2162
2163 /* If i is negative, then the indices are all at
2164 their maximum value and we're done. */
2165 if (i < 0)
2166 goto empty;
2167
2168 /* Increment the current index which we know is not at its
2169 maximum. Then move back to the right setting each index
2170 to its lowest possible value (one higher than the index
2171 to its left -- this maintains the sort order invariant). */
2172 indices[i]++;
2173 for (j=i+1 ; j<r ; j++)
2174 indices[j] = indices[j-1] + 1;
2175
2176 /* Update the result tuple for the new indices
2177 starting with i, the leftmost index that changed */
2178 for ( ; i<r ; i++) {
2179 index = indices[i];
2180 elem = PyTuple_GET_ITEM(pool, index);
2181 Py_INCREF(elem);
2182 oldelem = PyTuple_GET_ITEM(result, i);
2183 PyTuple_SET_ITEM(result, i, elem);
2184 Py_DECREF(oldelem);
2185 }
2186 }
2187
2188 Py_INCREF(result);
2189 return result;
2190
2191empty:
2192 co->stopped = 1;
2193 return NULL;
2194}
2195
2196PyDoc_STRVAR(combinations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +00002197"combinations(iterable[, r]) --> combinations object\n\
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002198\n\
2199Return successive r-length combinations of elements in the iterable.\n\n\
2200combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2201
2202static PyTypeObject combinations_type = {
2203 PyVarObject_HEAD_INIT(NULL, 0)
2204 "itertools.combinations", /* tp_name */
2205 sizeof(combinationsobject), /* tp_basicsize */
2206 0, /* tp_itemsize */
2207 /* methods */
2208 (destructor)combinations_dealloc, /* tp_dealloc */
2209 0, /* tp_print */
2210 0, /* tp_getattr */
2211 0, /* tp_setattr */
2212 0, /* tp_compare */
2213 0, /* tp_repr */
2214 0, /* tp_as_number */
2215 0, /* tp_as_sequence */
2216 0, /* tp_as_mapping */
2217 0, /* tp_hash */
2218 0, /* tp_call */
2219 0, /* tp_str */
2220 PyObject_GenericGetAttr, /* tp_getattro */
2221 0, /* tp_setattro */
2222 0, /* tp_as_buffer */
2223 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2224 Py_TPFLAGS_BASETYPE, /* tp_flags */
2225 combinations_doc, /* tp_doc */
2226 (traverseproc)combinations_traverse, /* tp_traverse */
2227 0, /* tp_clear */
2228 0, /* tp_richcompare */
2229 0, /* tp_weaklistoffset */
2230 PyObject_SelfIter, /* tp_iter */
2231 (iternextfunc)combinations_next, /* tp_iternext */
2232 0, /* tp_methods */
2233 0, /* tp_members */
2234 0, /* tp_getset */
2235 0, /* tp_base */
2236 0, /* tp_dict */
2237 0, /* tp_descr_get */
2238 0, /* tp_descr_set */
2239 0, /* tp_dictoffset */
2240 0, /* tp_init */
2241 0, /* tp_alloc */
2242 combinations_new, /* tp_new */
2243 PyObject_GC_Del, /* tp_free */
2244};
2245
2246
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002247/* combinations with replacement object *******************************************/
2248
2249/* Equivalent to:
2250
2251 def combinations_with_replacement(iterable, r):
2252 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2253 # number items returned: (n+r-1)! / r! / (n-1)!
2254 pool = tuple(iterable)
2255 n = len(pool)
2256 indices = [0] * r
2257 yield tuple(pool[i] for i in indices)
2258 while 1:
2259 for i in reversed(range(r)):
2260 if indices[i] != n - 1:
2261 break
2262 else:
2263 return
2264 indices[i:] = [indices[i] + 1] * (r - i)
2265 yield tuple(pool[i] for i in indices)
2266
2267 def combinations_with_replacement2(iterable, r):
2268 'Alternate version that filters from product()'
2269 pool = tuple(iterable)
2270 n = len(pool)
2271 for indices in product(range(n), repeat=r):
2272 if sorted(indices) == list(indices):
2273 yield tuple(pool[i] for i in indices)
2274*/
2275typedef struct {
2276 PyObject_HEAD
2277 PyObject *pool; /* input converted to a tuple */
2278 Py_ssize_t *indices; /* one index per result element */
2279 PyObject *result; /* most recently returned result tuple */
2280 Py_ssize_t r; /* size of result tuple */
2281 int stopped; /* set to 1 when the cwr iterator is exhausted */
2282} cwrobject;
2283
2284static PyTypeObject cwr_type;
2285
2286static PyObject *
2287cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2288{
2289 cwrobject *co;
2290 Py_ssize_t n;
2291 Py_ssize_t r;
2292 PyObject *pool = NULL;
2293 PyObject *iterable = NULL;
2294 Py_ssize_t *indices = NULL;
2295 Py_ssize_t i;
2296 static char *kwargs[] = {"iterable", "r", NULL};
2297
2298 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2299 &iterable, &r))
2300 return NULL;
2301
2302 pool = PySequence_Tuple(iterable);
2303 if (pool == NULL)
2304 goto error;
2305 n = PyTuple_GET_SIZE(pool);
2306 if (r < 0) {
2307 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2308 goto error;
2309 }
2310
2311 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2312 if (indices == NULL) {
2313 PyErr_NoMemory();
2314 goto error;
2315 }
2316
2317 for (i=0 ; i<r ; i++)
2318 indices[i] = 0;
2319
2320 /* create cwrobject structure */
2321 co = (cwrobject *)type->tp_alloc(type, 0);
2322 if (co == NULL)
2323 goto error;
2324
2325 co->pool = pool;
2326 co->indices = indices;
2327 co->result = NULL;
2328 co->r = r;
2329 co->stopped = !n && r;
2330
2331 return (PyObject *)co;
2332
2333error:
2334 if (indices != NULL)
2335 PyMem_Free(indices);
2336 Py_XDECREF(pool);
2337 return NULL;
2338}
2339
2340static void
2341cwr_dealloc(cwrobject *co)
2342{
2343 PyObject_GC_UnTrack(co);
2344 Py_XDECREF(co->pool);
2345 Py_XDECREF(co->result);
2346 if (co->indices != NULL)
2347 PyMem_Free(co->indices);
2348 Py_TYPE(co)->tp_free(co);
2349}
2350
2351static int
2352cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2353{
2354 Py_VISIT(co->pool);
2355 Py_VISIT(co->result);
2356 return 0;
2357}
2358
2359static PyObject *
2360cwr_next(cwrobject *co)
2361{
2362 PyObject *elem;
2363 PyObject *oldelem;
2364 PyObject *pool = co->pool;
2365 Py_ssize_t *indices = co->indices;
2366 PyObject *result = co->result;
2367 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2368 Py_ssize_t r = co->r;
2369 Py_ssize_t i, j, index;
2370
2371 if (co->stopped)
2372 return NULL;
2373
2374 if (result == NULL) {
2375 /* On the first pass, initialize result tuple using the indices */
2376 result = PyTuple_New(r);
2377 if (result == NULL)
2378 goto empty;
2379 co->result = result;
2380 for (i=0; i<r ; i++) {
2381 index = indices[i];
2382 elem = PyTuple_GET_ITEM(pool, index);
2383 Py_INCREF(elem);
2384 PyTuple_SET_ITEM(result, i, elem);
2385 }
2386 } else {
2387 /* Copy the previous result tuple or re-use it if available */
2388 if (Py_REFCNT(result) > 1) {
2389 PyObject *old_result = result;
2390 result = PyTuple_New(r);
2391 if (result == NULL)
2392 goto empty;
2393 co->result = result;
2394 for (i=0; i<r ; i++) {
2395 elem = PyTuple_GET_ITEM(old_result, i);
2396 Py_INCREF(elem);
2397 PyTuple_SET_ITEM(result, i, elem);
2398 }
2399 Py_DECREF(old_result);
2400 }
2401 /* Now, we've got the only copy so we can update it in-place CPython's
2402 empty tuple is a singleton and cached in PyTuple's freelist. */
2403 assert(r == 0 || Py_REFCNT(result) == 1);
2404
2405 /* Scan indices right-to-left until finding one that is not
2406 * at its maximum (n-1). */
2407 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2408 ;
2409
2410 /* If i is negative, then the indices are all at
2411 their maximum value and we're done. */
2412 if (i < 0)
2413 goto empty;
2414
2415 /* Increment the current index which we know is not at its
2416 maximum. Then set all to the right to the same value. */
2417 indices[i]++;
2418 for (j=i+1 ; j<r ; j++)
2419 indices[j] = indices[j-1];
2420
2421 /* Update the result tuple for the new indices
2422 starting with i, the leftmost index that changed */
2423 for ( ; i<r ; i++) {
2424 index = indices[i];
2425 elem = PyTuple_GET_ITEM(pool, index);
2426 Py_INCREF(elem);
2427 oldelem = PyTuple_GET_ITEM(result, i);
2428 PyTuple_SET_ITEM(result, i, elem);
2429 Py_DECREF(oldelem);
2430 }
2431 }
2432
2433 Py_INCREF(result);
2434 return result;
2435
2436empty:
2437 co->stopped = 1;
2438 return NULL;
2439}
2440
2441PyDoc_STRVAR(cwr_doc,
2442"combinations_with_replacement(iterable[, r]) --> combinations_with_replacement object\n\
2443\n\
2444Return successive r-length combinations of elements in the iterable\n\
2445allowing individual elements to have successive repeats.\n\
2446combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2447
2448static PyTypeObject cwr_type = {
2449 PyVarObject_HEAD_INIT(NULL, 0)
2450 "itertools.combinations_with_replacement", /* tp_name */
2451 sizeof(cwrobject), /* tp_basicsize */
2452 0, /* tp_itemsize */
2453 /* methods */
2454 (destructor)cwr_dealloc, /* tp_dealloc */
2455 0, /* tp_print */
2456 0, /* tp_getattr */
2457 0, /* tp_setattr */
2458 0, /* tp_compare */
2459 0, /* tp_repr */
2460 0, /* tp_as_number */
2461 0, /* tp_as_sequence */
2462 0, /* tp_as_mapping */
2463 0, /* tp_hash */
2464 0, /* tp_call */
2465 0, /* tp_str */
2466 PyObject_GenericGetAttr, /* tp_getattro */
2467 0, /* tp_setattro */
2468 0, /* tp_as_buffer */
2469 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2470 Py_TPFLAGS_BASETYPE, /* tp_flags */
2471 cwr_doc, /* tp_doc */
2472 (traverseproc)cwr_traverse, /* tp_traverse */
2473 0, /* tp_clear */
2474 0, /* tp_richcompare */
2475 0, /* tp_weaklistoffset */
2476 PyObject_SelfIter, /* tp_iter */
2477 (iternextfunc)cwr_next, /* tp_iternext */
2478 0, /* tp_methods */
2479 0, /* tp_members */
2480 0, /* tp_getset */
2481 0, /* tp_base */
2482 0, /* tp_dict */
2483 0, /* tp_descr_get */
2484 0, /* tp_descr_set */
2485 0, /* tp_dictoffset */
2486 0, /* tp_init */
2487 0, /* tp_alloc */
2488 cwr_new, /* tp_new */
2489 PyObject_GC_Del, /* tp_free */
2490};
2491
2492
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002493/* permutations object ************************************************************
2494
2495def permutations(iterable, r=None):
2496 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2497 pool = tuple(iterable)
2498 n = len(pool)
2499 r = n if r is None else r
2500 indices = range(n)
2501 cycles = range(n-r+1, n+1)[::-1]
2502 yield tuple(pool[i] for i in indices[:r])
2503 while n:
2504 for i in reversed(range(r)):
2505 cycles[i] -= 1
2506 if cycles[i] == 0:
2507 indices[i:] = indices[i+1:] + indices[i:i+1]
2508 cycles[i] = n - i
2509 else:
2510 j = cycles[i]
2511 indices[i], indices[-j] = indices[-j], indices[i]
2512 yield tuple(pool[i] for i in indices[:r])
2513 break
2514 else:
2515 return
2516*/
2517
2518typedef struct {
2519 PyObject_HEAD
2520 PyObject *pool; /* input converted to a tuple */
2521 Py_ssize_t *indices; /* one index per element in the pool */
2522 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2523 PyObject *result; /* most recently returned result tuple */
2524 Py_ssize_t r; /* size of result tuple */
2525 int stopped; /* set to 1 when the permutations iterator is exhausted */
2526} permutationsobject;
2527
2528static PyTypeObject permutations_type;
2529
2530static PyObject *
2531permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2532{
2533 permutationsobject *po;
2534 Py_ssize_t n;
2535 Py_ssize_t r;
2536 PyObject *robj = Py_None;
2537 PyObject *pool = NULL;
2538 PyObject *iterable = NULL;
2539 Py_ssize_t *indices = NULL;
2540 Py_ssize_t *cycles = NULL;
2541 Py_ssize_t i;
2542 static char *kwargs[] = {"iterable", "r", NULL};
2543
2544 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2545 &iterable, &robj))
2546 return NULL;
2547
2548 pool = PySequence_Tuple(iterable);
2549 if (pool == NULL)
2550 goto error;
2551 n = PyTuple_GET_SIZE(pool);
2552
2553 r = n;
2554 if (robj != Py_None) {
2555 r = PyInt_AsSsize_t(robj);
2556 if (r == -1 && PyErr_Occurred())
2557 goto error;
2558 }
2559 if (r < 0) {
2560 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2561 goto error;
2562 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002563
2564 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2565 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2566 if (indices == NULL || cycles == NULL) {
2567 PyErr_NoMemory();
2568 goto error;
2569 }
2570
2571 for (i=0 ; i<n ; i++)
2572 indices[i] = i;
2573 for (i=0 ; i<r ; i++)
2574 cycles[i] = n - i;
2575
2576 /* create permutationsobject structure */
2577 po = (permutationsobject *)type->tp_alloc(type, 0);
2578 if (po == NULL)
2579 goto error;
2580
2581 po->pool = pool;
2582 po->indices = indices;
2583 po->cycles = cycles;
2584 po->result = NULL;
2585 po->r = r;
Raymond Hettinger5b913e32009-01-08 06:39:04 +00002586 po->stopped = r > n ? 1 : 0;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002587
2588 return (PyObject *)po;
2589
2590error:
2591 if (indices != NULL)
2592 PyMem_Free(indices);
2593 if (cycles != NULL)
2594 PyMem_Free(cycles);
2595 Py_XDECREF(pool);
2596 return NULL;
2597}
2598
2599static void
2600permutations_dealloc(permutationsobject *po)
2601{
2602 PyObject_GC_UnTrack(po);
2603 Py_XDECREF(po->pool);
2604 Py_XDECREF(po->result);
2605 PyMem_Free(po->indices);
2606 PyMem_Free(po->cycles);
2607 Py_TYPE(po)->tp_free(po);
2608}
2609
2610static int
2611permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2612{
Raymond Hettinger6e3e4152008-03-05 21:04:32 +00002613 Py_VISIT(po->pool);
2614 Py_VISIT(po->result);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002615 return 0;
2616}
2617
2618static PyObject *
2619permutations_next(permutationsobject *po)
2620{
2621 PyObject *elem;
2622 PyObject *oldelem;
2623 PyObject *pool = po->pool;
2624 Py_ssize_t *indices = po->indices;
2625 Py_ssize_t *cycles = po->cycles;
2626 PyObject *result = po->result;
2627 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2628 Py_ssize_t r = po->r;
2629 Py_ssize_t i, j, k, index;
2630
2631 if (po->stopped)
2632 return NULL;
2633
2634 if (result == NULL) {
2635 /* On the first pass, initialize result tuple using the indices */
2636 result = PyTuple_New(r);
2637 if (result == NULL)
2638 goto empty;
2639 po->result = result;
2640 for (i=0; i<r ; i++) {
2641 index = indices[i];
2642 elem = PyTuple_GET_ITEM(pool, index);
2643 Py_INCREF(elem);
2644 PyTuple_SET_ITEM(result, i, elem);
2645 }
2646 } else {
2647 if (n == 0)
2648 goto empty;
2649
2650 /* Copy the previous result tuple or re-use it if available */
2651 if (Py_REFCNT(result) > 1) {
2652 PyObject *old_result = result;
2653 result = PyTuple_New(r);
2654 if (result == NULL)
2655 goto empty;
2656 po->result = result;
2657 for (i=0; i<r ; i++) {
2658 elem = PyTuple_GET_ITEM(old_result, i);
2659 Py_INCREF(elem);
2660 PyTuple_SET_ITEM(result, i, elem);
2661 }
2662 Py_DECREF(old_result);
2663 }
2664 /* Now, we've got the only copy so we can update it in-place */
2665 assert(r == 0 || Py_REFCNT(result) == 1);
2666
2667 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2668 for (i=r-1 ; i>=0 ; i--) {
2669 cycles[i] -= 1;
2670 if (cycles[i] == 0) {
2671 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2672 index = indices[i];
2673 for (j=i ; j<n-1 ; j++)
2674 indices[j] = indices[j+1];
2675 indices[n-1] = index;
2676 cycles[i] = n - i;
2677 } else {
2678 j = cycles[i];
2679 index = indices[i];
2680 indices[i] = indices[n-j];
2681 indices[n-j] = index;
2682
2683 for (k=i; k<r ; k++) {
2684 /* start with i, the leftmost element that changed */
2685 /* yield tuple(pool[k] for k in indices[:r]) */
2686 index = indices[k];
2687 elem = PyTuple_GET_ITEM(pool, index);
2688 Py_INCREF(elem);
2689 oldelem = PyTuple_GET_ITEM(result, k);
2690 PyTuple_SET_ITEM(result, k, elem);
2691 Py_DECREF(oldelem);
2692 }
2693 break;
2694 }
2695 }
2696 /* If i is negative, then the cycles have all
2697 rolled-over and we're done. */
2698 if (i < 0)
2699 goto empty;
2700 }
2701 Py_INCREF(result);
2702 return result;
2703
2704empty:
2705 po->stopped = 1;
2706 return NULL;
2707}
2708
2709PyDoc_STRVAR(permutations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +00002710"permutations(iterable[, r]) --> permutations object\n\
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002711\n\
2712Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl8550d692008-06-10 12:46:39 +00002713permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002714
2715static PyTypeObject permutations_type = {
2716 PyVarObject_HEAD_INIT(NULL, 0)
2717 "itertools.permutations", /* tp_name */
2718 sizeof(permutationsobject), /* tp_basicsize */
2719 0, /* tp_itemsize */
2720 /* methods */
2721 (destructor)permutations_dealloc, /* tp_dealloc */
2722 0, /* tp_print */
2723 0, /* tp_getattr */
2724 0, /* tp_setattr */
2725 0, /* tp_compare */
2726 0, /* tp_repr */
2727 0, /* tp_as_number */
2728 0, /* tp_as_sequence */
2729 0, /* tp_as_mapping */
2730 0, /* tp_hash */
2731 0, /* tp_call */
2732 0, /* tp_str */
2733 PyObject_GenericGetAttr, /* tp_getattro */
2734 0, /* tp_setattro */
2735 0, /* tp_as_buffer */
2736 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2737 Py_TPFLAGS_BASETYPE, /* tp_flags */
2738 permutations_doc, /* tp_doc */
2739 (traverseproc)permutations_traverse, /* tp_traverse */
2740 0, /* tp_clear */
2741 0, /* tp_richcompare */
2742 0, /* tp_weaklistoffset */
2743 PyObject_SelfIter, /* tp_iter */
2744 (iternextfunc)permutations_next, /* tp_iternext */
2745 0, /* tp_methods */
2746 0, /* tp_members */
2747 0, /* tp_getset */
2748 0, /* tp_base */
2749 0, /* tp_dict */
2750 0, /* tp_descr_get */
2751 0, /* tp_descr_set */
2752 0, /* tp_dictoffset */
2753 0, /* tp_init */
2754 0, /* tp_alloc */
2755 permutations_new, /* tp_new */
2756 PyObject_GC_Del, /* tp_free */
2757};
2758
2759
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002760/* compress object ************************************************************/
2761
2762/* Equivalent to:
2763
2764 def compress(data, selectors):
2765 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2766 return (d for d, s in izip(data, selectors) if s)
2767*/
2768
2769typedef struct {
2770 PyObject_HEAD
2771 PyObject *data;
2772 PyObject *selectors;
2773} compressobject;
2774
2775static PyTypeObject compress_type;
2776
2777static PyObject *
2778compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2779{
2780 PyObject *seq1, *seq2;
2781 PyObject *data=NULL, *selectors=NULL;
2782 compressobject *lz;
2783
2784 if (type == &compress_type && !_PyArg_NoKeywords("compress()", kwds))
2785 return NULL;
2786
2787 if (!PyArg_UnpackTuple(args, "compress", 2, 2, &seq1, &seq2))
2788 return NULL;
2789
2790 data = PyObject_GetIter(seq1);
2791 if (data == NULL)
2792 goto fail;
2793 selectors = PyObject_GetIter(seq2);
2794 if (selectors == NULL)
2795 goto fail;
2796
2797 /* create compressobject structure */
2798 lz = (compressobject *)type->tp_alloc(type, 0);
2799 if (lz == NULL)
2800 goto fail;
2801 lz->data = data;
2802 lz->selectors = selectors;
2803 return (PyObject *)lz;
2804
2805fail:
2806 Py_XDECREF(data);
2807 Py_XDECREF(selectors);
2808 return NULL;
2809}
2810
2811static void
2812compress_dealloc(compressobject *lz)
2813{
2814 PyObject_GC_UnTrack(lz);
2815 Py_XDECREF(lz->data);
2816 Py_XDECREF(lz->selectors);
2817 Py_TYPE(lz)->tp_free(lz);
2818}
2819
2820static int
2821compress_traverse(compressobject *lz, visitproc visit, void *arg)
2822{
2823 Py_VISIT(lz->data);
2824 Py_VISIT(lz->selectors);
2825 return 0;
2826}
2827
2828static PyObject *
2829compress_next(compressobject *lz)
2830{
2831 PyObject *data = lz->data, *selectors = lz->selectors;
2832 PyObject *datum, *selector;
2833 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2834 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2835 int ok;
2836
2837 while (1) {
2838 /* Steps: get datum, get selector, evaluate selector.
2839 Order is important (to match the pure python version
2840 in terms of which input gets a chance to raise an
2841 exception first).
2842 */
2843
2844 datum = datanext(data);
2845 if (datum == NULL)
2846 return NULL;
2847
2848 selector = selectornext(selectors);
2849 if (selector == NULL) {
2850 Py_DECREF(datum);
2851 return NULL;
2852 }
2853
2854 ok = PyObject_IsTrue(selector);
2855 Py_DECREF(selector);
2856 if (ok == 1)
2857 return datum;
2858 Py_DECREF(datum);
2859 if (ok == -1)
2860 return NULL;
2861 }
2862}
2863
2864PyDoc_STRVAR(compress_doc,
2865"compress(data sequence, selector sequence) --> iterator over selected data\n\
2866\n\
2867Return data elements corresponding to true selector elements.\n\
2868Forms a shorter iterator from selected data elements using the\n\
2869selectors to choose the data elements.");
2870
2871static PyTypeObject compress_type = {
2872 PyVarObject_HEAD_INIT(NULL, 0)
2873 "itertools.compress", /* tp_name */
2874 sizeof(compressobject), /* tp_basicsize */
2875 0, /* tp_itemsize */
2876 /* methods */
2877 (destructor)compress_dealloc, /* tp_dealloc */
2878 0, /* tp_print */
2879 0, /* tp_getattr */
2880 0, /* tp_setattr */
2881 0, /* tp_compare */
2882 0, /* tp_repr */
2883 0, /* tp_as_number */
2884 0, /* tp_as_sequence */
2885 0, /* tp_as_mapping */
2886 0, /* tp_hash */
2887 0, /* tp_call */
2888 0, /* tp_str */
2889 PyObject_GenericGetAttr, /* tp_getattro */
2890 0, /* tp_setattro */
2891 0, /* tp_as_buffer */
2892 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2893 Py_TPFLAGS_BASETYPE, /* tp_flags */
2894 compress_doc, /* tp_doc */
2895 (traverseproc)compress_traverse, /* tp_traverse */
2896 0, /* tp_clear */
2897 0, /* tp_richcompare */
2898 0, /* tp_weaklistoffset */
2899 PyObject_SelfIter, /* tp_iter */
2900 (iternextfunc)compress_next, /* tp_iternext */
2901 0, /* tp_methods */
2902 0, /* tp_members */
2903 0, /* tp_getset */
2904 0, /* tp_base */
2905 0, /* tp_dict */
2906 0, /* tp_descr_get */
2907 0, /* tp_descr_set */
2908 0, /* tp_dictoffset */
2909 0, /* tp_init */
2910 0, /* tp_alloc */
2911 compress_new, /* tp_new */
2912 PyObject_GC_Del, /* tp_free */
2913};
2914
2915
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002916/* ifilter object ************************************************************/
2917
2918typedef struct {
2919 PyObject_HEAD
2920 PyObject *func;
2921 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002922} ifilterobject;
2923
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002924static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002925
2926static PyObject *
2927ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2928{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002929 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002930 PyObject *it;
2931 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002932
Georg Brandlb84c1372007-01-21 10:28:43 +00002933 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002934 return NULL;
2935
Raymond Hettinger60eca932003-02-09 06:40:58 +00002936 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002937 return NULL;
2938
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002939 /* Get iterator. */
2940 it = PyObject_GetIter(seq);
2941 if (it == NULL)
2942 return NULL;
2943
2944 /* create ifilterobject structure */
2945 lz = (ifilterobject *)type->tp_alloc(type, 0);
2946 if (lz == NULL) {
2947 Py_DECREF(it);
2948 return NULL;
2949 }
2950 Py_INCREF(func);
2951 lz->func = func;
2952 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002953
2954 return (PyObject *)lz;
2955}
2956
2957static void
2958ifilter_dealloc(ifilterobject *lz)
2959{
2960 PyObject_GC_UnTrack(lz);
2961 Py_XDECREF(lz->func);
2962 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00002963 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002964}
2965
2966static int
2967ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
2968{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002969 Py_VISIT(lz->it);
2970 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002971 return 0;
2972}
2973
2974static PyObject *
2975ifilter_next(ifilterobject *lz)
2976{
2977 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002978 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002979 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002980 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002981
Christian Heimese93237d2007-12-19 02:37:44 +00002982 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002983 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002984 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002985 if (item == NULL)
2986 return NULL;
2987
Facundo Batistab1d70e22008-02-25 23:46:02 +00002988 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002989 ok = PyObject_IsTrue(item);
2990 } else {
2991 PyObject *good;
2992 good = PyObject_CallFunctionObjArgs(lz->func,
2993 item, NULL);
2994 if (good == NULL) {
2995 Py_DECREF(item);
2996 return NULL;
2997 }
2998 ok = PyObject_IsTrue(good);
2999 Py_DECREF(good);
3000 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003001 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003002 return item;
3003 Py_DECREF(item);
3004 }
3005}
3006
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003007PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003008"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003009\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003010Return those items of sequence for which function(item) is true.\n\
3011If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003012
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003013static PyTypeObject ifilter_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003014 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003015 "itertools.ifilter", /* tp_name */
3016 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003017 0, /* tp_itemsize */
3018 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003019 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003020 0, /* tp_print */
3021 0, /* tp_getattr */
3022 0, /* tp_setattr */
3023 0, /* tp_compare */
3024 0, /* tp_repr */
3025 0, /* tp_as_number */
3026 0, /* tp_as_sequence */
3027 0, /* tp_as_mapping */
3028 0, /* tp_hash */
3029 0, /* tp_call */
3030 0, /* tp_str */
3031 PyObject_GenericGetAttr, /* tp_getattro */
3032 0, /* tp_setattro */
3033 0, /* tp_as_buffer */
3034 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3035 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003036 ifilter_doc, /* tp_doc */
3037 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003038 0, /* tp_clear */
3039 0, /* tp_richcompare */
3040 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003041 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003042 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003043 0, /* tp_methods */
3044 0, /* tp_members */
3045 0, /* tp_getset */
3046 0, /* tp_base */
3047 0, /* tp_dict */
3048 0, /* tp_descr_get */
3049 0, /* tp_descr_set */
3050 0, /* tp_dictoffset */
3051 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003052 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003053 ifilter_new, /* tp_new */
3054 PyObject_GC_Del, /* tp_free */
3055};
3056
3057
3058/* ifilterfalse object ************************************************************/
3059
3060typedef struct {
3061 PyObject_HEAD
3062 PyObject *func;
3063 PyObject *it;
3064} ifilterfalseobject;
3065
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003066static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003067
3068static PyObject *
3069ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3070{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00003071 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003072 PyObject *it;
3073 ifilterfalseobject *lz;
3074
Georg Brandlb84c1372007-01-21 10:28:43 +00003075 if (type == &ifilterfalse_type &&
3076 !_PyArg_NoKeywords("ifilterfalse()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00003077 return NULL;
3078
Raymond Hettinger60eca932003-02-09 06:40:58 +00003079 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
3080 return NULL;
3081
3082 /* Get iterator. */
3083 it = PyObject_GetIter(seq);
3084 if (it == NULL)
3085 return NULL;
3086
3087 /* create ifilterfalseobject structure */
3088 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
3089 if (lz == NULL) {
3090 Py_DECREF(it);
3091 return NULL;
3092 }
3093 Py_INCREF(func);
3094 lz->func = func;
3095 lz->it = it;
3096
3097 return (PyObject *)lz;
3098}
3099
3100static void
3101ifilterfalse_dealloc(ifilterfalseobject *lz)
3102{
3103 PyObject_GC_UnTrack(lz);
3104 Py_XDECREF(lz->func);
3105 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00003106 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003107}
3108
3109static int
3110ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
3111{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00003112 Py_VISIT(lz->it);
3113 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003114 return 0;
3115}
3116
3117static PyObject *
3118ifilterfalse_next(ifilterfalseobject *lz)
3119{
3120 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00003121 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003122 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00003123 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003124
Christian Heimese93237d2007-12-19 02:37:44 +00003125 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003126 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00003127 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003128 if (item == NULL)
3129 return NULL;
3130
Facundo Batistab1d70e22008-02-25 23:46:02 +00003131 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
Raymond Hettinger60eca932003-02-09 06:40:58 +00003132 ok = PyObject_IsTrue(item);
3133 } else {
3134 PyObject *good;
3135 good = PyObject_CallFunctionObjArgs(lz->func,
3136 item, NULL);
3137 if (good == NULL) {
3138 Py_DECREF(item);
3139 return NULL;
3140 }
3141 ok = PyObject_IsTrue(good);
3142 Py_DECREF(good);
3143 }
3144 if (!ok)
3145 return item;
3146 Py_DECREF(item);
3147 }
3148}
3149
Raymond Hettinger60eca932003-02-09 06:40:58 +00003150PyDoc_STRVAR(ifilterfalse_doc,
3151"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
3152\n\
3153Return those items of sequence for which function(item) is false.\n\
3154If function is None, return the items that are false.");
3155
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003156static PyTypeObject ifilterfalse_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003157 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003158 "itertools.ifilterfalse", /* tp_name */
3159 sizeof(ifilterfalseobject), /* tp_basicsize */
3160 0, /* tp_itemsize */
3161 /* methods */
3162 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
3163 0, /* tp_print */
3164 0, /* tp_getattr */
3165 0, /* tp_setattr */
3166 0, /* tp_compare */
3167 0, /* tp_repr */
3168 0, /* tp_as_number */
3169 0, /* tp_as_sequence */
3170 0, /* tp_as_mapping */
3171 0, /* tp_hash */
3172 0, /* tp_call */
3173 0, /* tp_str */
3174 PyObject_GenericGetAttr, /* tp_getattro */
3175 0, /* tp_setattro */
3176 0, /* tp_as_buffer */
3177 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3178 Py_TPFLAGS_BASETYPE, /* tp_flags */
3179 ifilterfalse_doc, /* tp_doc */
3180 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
3181 0, /* tp_clear */
3182 0, /* tp_richcompare */
3183 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003184 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003185 (iternextfunc)ifilterfalse_next, /* tp_iternext */
3186 0, /* tp_methods */
3187 0, /* tp_members */
3188 0, /* tp_getset */
3189 0, /* tp_base */
3190 0, /* tp_dict */
3191 0, /* tp_descr_get */
3192 0, /* tp_descr_set */
3193 0, /* tp_dictoffset */
3194 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003195 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003196 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003197 PyObject_GC_Del, /* tp_free */
3198};
3199
3200
3201/* count object ************************************************************/
3202
3203typedef struct {
3204 PyObject_HEAD
Jack Diederich6c433a92006-05-26 11:15:17 +00003205 Py_ssize_t cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003206 PyObject *long_cnt; /* Arbitrarily large count when cnt >= PY_SSIZE_T_MAX */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003207} countobject;
3208
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003209static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003210
3211static PyObject *
3212count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3213{
3214 countobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00003215 Py_ssize_t cnt = 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003216 PyObject *cnt_arg = NULL;
3217 PyObject *long_cnt = NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003218
Georg Brandlb84c1372007-01-21 10:28:43 +00003219 if (type == &count_type && !_PyArg_NoKeywords("count()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00003220 return NULL;
3221
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003222 if (!PyArg_UnpackTuple(args, "count", 0, 1, &cnt_arg))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003223 return NULL;
3224
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003225 if (cnt_arg != NULL) {
3226 cnt = PyInt_AsSsize_t(cnt_arg);
3227 if (cnt == -1 && PyErr_Occurred()) {
3228 PyErr_Clear();
3229 if (!PyLong_Check(cnt_arg)) {
3230 PyErr_SetString(PyExc_TypeError, "an integer is required");
3231 return NULL;
3232 }
3233 long_cnt = cnt_arg;
3234 Py_INCREF(long_cnt);
3235 cnt = PY_SSIZE_T_MAX;
3236 }
3237 }
3238
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003239 /* create countobject structure */
3240 lz = (countobject *)PyObject_New(countobject, &count_type);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003241 if (lz == NULL) {
3242 Py_XDECREF(long_cnt);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003243 return NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003244 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003245 lz->cnt = cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003246 lz->long_cnt = long_cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003247
3248 return (PyObject *)lz;
3249}
3250
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003251static void
3252count_dealloc(countobject *lz)
3253{
3254 Py_XDECREF(lz->long_cnt);
3255 PyObject_Del(lz);
3256}
3257
3258static PyObject *
3259count_nextlong(countobject *lz)
3260{
3261 static PyObject *one = NULL;
3262 PyObject *cnt;
3263 PyObject *stepped_up;
3264
3265 if (lz->long_cnt == NULL) {
3266 lz->long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
3267 if (lz->long_cnt == NULL)
3268 return NULL;
3269 }
3270 if (one == NULL) {
3271 one = PyInt_FromLong(1);
3272 if (one == NULL)
3273 return NULL;
3274 }
3275 cnt = lz->long_cnt;
3276 assert(cnt != NULL);
3277 stepped_up = PyNumber_Add(cnt, one);
3278 if (stepped_up == NULL)
3279 return NULL;
3280 lz->long_cnt = stepped_up;
3281 return cnt;
3282}
3283
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003284static PyObject *
3285count_next(countobject *lz)
3286{
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003287 if (lz->cnt == PY_SSIZE_T_MAX)
3288 return count_nextlong(lz);
Jack Diederich36234e82006-09-21 17:50:26 +00003289 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003290}
3291
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003292static PyObject *
3293count_repr(countobject *lz)
3294{
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003295 PyObject *cnt_repr;
3296 PyObject *result;
3297
3298 if (lz->cnt != PY_SSIZE_T_MAX)
Gregory P. Smithdd96db62008-06-09 04:58:54 +00003299 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003300
3301 cnt_repr = PyObject_Repr(lz->long_cnt);
3302 if (cnt_repr == NULL)
3303 return NULL;
Gregory P. Smithdd96db62008-06-09 04:58:54 +00003304 result = PyString_FromFormat("count(%s)", PyString_AS_STRING(cnt_repr));
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003305 Py_DECREF(cnt_repr);
3306 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003307}
3308
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003309PyDoc_STRVAR(count_doc,
3310"count([firstval]) --> count object\n\
3311\n\
3312Return a count object whose .next() method returns consecutive\n\
3313integers starting from zero or, if specified, from firstval.");
3314
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003315static PyTypeObject count_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003316 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003317 "itertools.count", /* tp_name */
3318 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003319 0, /* tp_itemsize */
3320 /* methods */
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003321 (destructor)count_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003322 0, /* tp_print */
3323 0, /* tp_getattr */
3324 0, /* tp_setattr */
3325 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003326 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003327 0, /* tp_as_number */
3328 0, /* tp_as_sequence */
3329 0, /* tp_as_mapping */
3330 0, /* tp_hash */
3331 0, /* tp_call */
3332 0, /* tp_str */
3333 PyObject_GenericGetAttr, /* tp_getattro */
3334 0, /* tp_setattro */
3335 0, /* tp_as_buffer */
3336 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003337 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003338 0, /* tp_traverse */
3339 0, /* tp_clear */
3340 0, /* tp_richcompare */
3341 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003342 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003343 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003344 0, /* tp_methods */
3345 0, /* tp_members */
3346 0, /* tp_getset */
3347 0, /* tp_base */
3348 0, /* tp_dict */
3349 0, /* tp_descr_get */
3350 0, /* tp_descr_set */
3351 0, /* tp_dictoffset */
3352 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003353 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003354 count_new, /* tp_new */
3355};
3356
3357
3358/* izip object ************************************************************/
3359
3360#include "Python.h"
3361
3362typedef struct {
3363 PyObject_HEAD
Martin v. Löwisad0a4622006-02-16 14:30:23 +00003364 Py_ssize_t tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003365 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00003366 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003367} izipobject;
3368
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003369static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003370
3371static PyObject *
3372izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3373{
3374 izipobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00003375 Py_ssize_t i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003376 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00003377 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00003378 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003379
Georg Brandlb84c1372007-01-21 10:28:43 +00003380 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00003381 return NULL;
3382
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003383 /* args must be a tuple */
3384 assert(PyTuple_Check(args));
3385
3386 /* obtain iterators */
3387 ittuple = PyTuple_New(tuplesize);
Neal Norwitz2f3136b2006-05-27 05:18:57 +00003388 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003389 return NULL;
3390 for (i=0; i < tuplesize; ++i) {
3391 PyObject *item = PyTuple_GET_ITEM(args, i);
3392 PyObject *it = PyObject_GetIter(item);
3393 if (it == NULL) {
3394 if (PyErr_ExceptionMatches(PyExc_TypeError))
3395 PyErr_Format(PyExc_TypeError,
Neal Norwitz2f3136b2006-05-27 05:18:57 +00003396 "izip argument #%zd must support iteration",
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003397 i+1);
3398 Py_DECREF(ittuple);
3399 return NULL;
3400 }
3401 PyTuple_SET_ITEM(ittuple, i, it);
3402 }
3403
Raymond Hettinger2012f172003-02-07 05:32:58 +00003404 /* create a result holder */
3405 result = PyTuple_New(tuplesize);
3406 if (result == NULL) {
3407 Py_DECREF(ittuple);
3408 return NULL;
3409 }
3410 for (i=0 ; i < tuplesize ; i++) {
3411 Py_INCREF(Py_None);
3412 PyTuple_SET_ITEM(result, i, Py_None);
3413 }
3414
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003415 /* create izipobject structure */
3416 lz = (izipobject *)type->tp_alloc(type, 0);
3417 if (lz == NULL) {
3418 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003419 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003420 return NULL;
3421 }
3422 lz->ittuple = ittuple;
3423 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00003424 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003425
3426 return (PyObject *)lz;
3427}
3428
3429static void
3430izip_dealloc(izipobject *lz)
3431{
3432 PyObject_GC_UnTrack(lz);
3433 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003434 Py_XDECREF(lz->result);
Christian Heimese93237d2007-12-19 02:37:44 +00003435 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003436}
3437
3438static int
3439izip_traverse(izipobject *lz, visitproc visit, void *arg)
3440{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00003441 Py_VISIT(lz->ittuple);
3442 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003443 return 0;
3444}
3445
3446static PyObject *
3447izip_next(izipobject *lz)
3448{
Jack Diederich6c433a92006-05-26 11:15:17 +00003449 Py_ssize_t i;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00003450 Py_ssize_t tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00003451 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003452 PyObject *it;
3453 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00003454 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003455
Raymond Hettingerb5a42082003-08-08 05:10:41 +00003456 if (tuplesize == 0)
3457 return NULL;
Christian Heimese93237d2007-12-19 02:37:44 +00003458 if (Py_REFCNT(result) == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00003459 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003460 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003461 it = PyTuple_GET_ITEM(lz->ittuple, i);
Christian Heimese93237d2007-12-19 02:37:44 +00003462 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00003463 if (item == NULL) {
3464 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003465 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00003466 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00003467 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003468 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00003469 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003470 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003471 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00003472 result = PyTuple_New(tuplesize);
3473 if (result == NULL)
3474 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003475 for (i=0 ; i < tuplesize ; i++) {
3476 it = PyTuple_GET_ITEM(lz->ittuple, i);
Christian Heimese93237d2007-12-19 02:37:44 +00003477 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003478 if (item == NULL) {
3479 Py_DECREF(result);
3480 return NULL;
3481 }
3482 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003483 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003484 }
3485 return result;
3486}
3487
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003488PyDoc_STRVAR(izip_doc,
3489"izip(iter1 [,iter2 [...]]) --> izip object\n\
3490\n\
3491Return a izip object whose .next() method returns a tuple where\n\
3492the i-th element comes from the i-th iterable argument. The .next()\n\
3493method continues until the shortest iterable in the argument sequence\n\
3494is exhausted and then it raises StopIteration. Works like the zip()\n\
3495function but consumes less memory by returning an iterator instead of\n\
3496a list.");
3497
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003498static PyTypeObject izip_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003499 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003500 "itertools.izip", /* tp_name */
3501 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003502 0, /* tp_itemsize */
3503 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003504 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003505 0, /* tp_print */
3506 0, /* tp_getattr */
3507 0, /* tp_setattr */
3508 0, /* tp_compare */
3509 0, /* tp_repr */
3510 0, /* tp_as_number */
3511 0, /* tp_as_sequence */
3512 0, /* tp_as_mapping */
3513 0, /* tp_hash */
3514 0, /* tp_call */
3515 0, /* tp_str */
3516 PyObject_GenericGetAttr, /* tp_getattro */
3517 0, /* tp_setattro */
3518 0, /* tp_as_buffer */
3519 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3520 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003521 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003522 (traverseproc)izip_traverse, /* tp_traverse */
3523 0, /* tp_clear */
3524 0, /* tp_richcompare */
3525 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003526 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003527 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003528 0, /* tp_methods */
3529 0, /* tp_members */
3530 0, /* tp_getset */
3531 0, /* tp_base */
3532 0, /* tp_dict */
3533 0, /* tp_descr_get */
3534 0, /* tp_descr_set */
3535 0, /* tp_dictoffset */
3536 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003537 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003538 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003539 PyObject_GC_Del, /* tp_free */
3540};
3541
3542
3543/* repeat object ************************************************************/
3544
3545typedef struct {
3546 PyObject_HEAD
3547 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00003548 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003549} repeatobject;
3550
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003551static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003552
3553static PyObject *
3554repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3555{
3556 repeatobject *ro;
3557 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00003558 Py_ssize_t cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003559
Georg Brandlb84c1372007-01-21 10:28:43 +00003560 if (type == &repeat_type && !_PyArg_NoKeywords("repeat()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00003561 return NULL;
3562
Jack Diederich6c433a92006-05-26 11:15:17 +00003563 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003564 return NULL;
3565
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003566 if (PyTuple_Size(args) == 2 && cnt < 0)
3567 cnt = 0;
3568
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003569 ro = (repeatobject *)type->tp_alloc(type, 0);
3570 if (ro == NULL)
3571 return NULL;
3572 Py_INCREF(element);
3573 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003574 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003575 return (PyObject *)ro;
3576}
3577
3578static void
3579repeat_dealloc(repeatobject *ro)
3580{
3581 PyObject_GC_UnTrack(ro);
3582 Py_XDECREF(ro->element);
Christian Heimese93237d2007-12-19 02:37:44 +00003583 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003584}
3585
3586static int
3587repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3588{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00003589 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003590 return 0;
3591}
3592
3593static PyObject *
3594repeat_next(repeatobject *ro)
3595{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003596 if (ro->cnt == 0)
3597 return NULL;
3598 if (ro->cnt > 0)
3599 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003600 Py_INCREF(ro->element);
3601 return ro->element;
3602}
3603
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003604static PyObject *
3605repeat_repr(repeatobject *ro)
3606{
3607 PyObject *result, *objrepr;
3608
3609 objrepr = PyObject_Repr(ro->element);
3610 if (objrepr == NULL)
3611 return NULL;
3612
3613 if (ro->cnt == -1)
Gregory P. Smithdd96db62008-06-09 04:58:54 +00003614 result = PyString_FromFormat("repeat(%s)",
3615 PyString_AS_STRING(objrepr));
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003616 else
Gregory P. Smithdd96db62008-06-09 04:58:54 +00003617 result = PyString_FromFormat("repeat(%s, %zd)",
3618 PyString_AS_STRING(objrepr), ro->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003619 Py_DECREF(objrepr);
3620 return result;
3621}
3622
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003623static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003624repeat_len(repeatobject *ro)
3625{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003626 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003627 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003628 return NULL;
3629 }
Jack Diederich6c433a92006-05-26 11:15:17 +00003630 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003631}
3632
Armin Rigof5b3e362006-02-11 21:32:43 +00003633PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003634
3635static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00003636 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003637 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003638};
3639
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003640PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003641"repeat(element [,times]) -> create an iterator which returns the element\n\
3642for the specified number of times. If not specified, returns the element\n\
3643endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003644
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003645static PyTypeObject repeat_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003646 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003647 "itertools.repeat", /* tp_name */
3648 sizeof(repeatobject), /* tp_basicsize */
3649 0, /* tp_itemsize */
3650 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003651 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003652 0, /* tp_print */
3653 0, /* tp_getattr */
3654 0, /* tp_setattr */
3655 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003656 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003657 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003658 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003659 0, /* tp_as_mapping */
3660 0, /* tp_hash */
3661 0, /* tp_call */
3662 0, /* tp_str */
3663 PyObject_GenericGetAttr, /* tp_getattro */
3664 0, /* tp_setattro */
3665 0, /* tp_as_buffer */
3666 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3667 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003668 repeat_doc, /* tp_doc */
3669 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003670 0, /* tp_clear */
3671 0, /* tp_richcompare */
3672 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003673 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003674 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003675 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003676 0, /* tp_members */
3677 0, /* tp_getset */
3678 0, /* tp_base */
3679 0, /* tp_dict */
3680 0, /* tp_descr_get */
3681 0, /* tp_descr_set */
3682 0, /* tp_dictoffset */
3683 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003684 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003685 repeat_new, /* tp_new */
3686 PyObject_GC_Del, /* tp_free */
3687};
3688
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003689/* iziplongest object ************************************************************/
3690
3691#include "Python.h"
3692
3693typedef struct {
3694 PyObject_HEAD
3695 Py_ssize_t tuplesize;
3696 Py_ssize_t numactive;
3697 PyObject *ittuple; /* tuple of iterators */
3698 PyObject *result;
3699 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003700} iziplongestobject;
3701
3702static PyTypeObject iziplongest_type;
3703
3704static PyObject *
3705izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3706{
3707 iziplongestobject *lz;
3708 Py_ssize_t i;
3709 PyObject *ittuple; /* tuple of iterators */
3710 PyObject *result;
3711 PyObject *fillvalue = Py_None;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003712 Py_ssize_t tuplesize = PySequence_Length(args);
3713
3714 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3715 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3716 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3717 PyErr_SetString(PyExc_TypeError,
3718 "izip_longest() got an unexpected keyword argument");
3719 return NULL;
3720 }
3721 }
3722
3723 /* args must be a tuple */
3724 assert(PyTuple_Check(args));
3725
3726 /* obtain iterators */
3727 ittuple = PyTuple_New(tuplesize);
3728 if (ittuple == NULL)
3729 return NULL;
3730 for (i=0; i < tuplesize; ++i) {
3731 PyObject *item = PyTuple_GET_ITEM(args, i);
3732 PyObject *it = PyObject_GetIter(item);
3733 if (it == NULL) {
3734 if (PyErr_ExceptionMatches(PyExc_TypeError))
3735 PyErr_Format(PyExc_TypeError,
3736 "izip_longest argument #%zd must support iteration",
3737 i+1);
3738 Py_DECREF(ittuple);
3739 return NULL;
3740 }
3741 PyTuple_SET_ITEM(ittuple, i, it);
3742 }
3743
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003744 /* create a result holder */
3745 result = PyTuple_New(tuplesize);
3746 if (result == NULL) {
3747 Py_DECREF(ittuple);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003748 return NULL;
3749 }
3750 for (i=0 ; i < tuplesize ; i++) {
3751 Py_INCREF(Py_None);
3752 PyTuple_SET_ITEM(result, i, Py_None);
3753 }
3754
3755 /* create iziplongestobject structure */
3756 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3757 if (lz == NULL) {
3758 Py_DECREF(ittuple);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003759 Py_DECREF(result);
3760 return NULL;
3761 }
3762 lz->ittuple = ittuple;
3763 lz->tuplesize = tuplesize;
3764 lz->numactive = tuplesize;
3765 lz->result = result;
3766 Py_INCREF(fillvalue);
3767 lz->fillvalue = fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003768 return (PyObject *)lz;
3769}
3770
3771static void
3772izip_longest_dealloc(iziplongestobject *lz)
3773{
3774 PyObject_GC_UnTrack(lz);
3775 Py_XDECREF(lz->ittuple);
3776 Py_XDECREF(lz->result);
3777 Py_XDECREF(lz->fillvalue);
Christian Heimese93237d2007-12-19 02:37:44 +00003778 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003779}
3780
3781static int
3782izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3783{
3784 Py_VISIT(lz->ittuple);
3785 Py_VISIT(lz->result);
3786 Py_VISIT(lz->fillvalue);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003787 return 0;
3788}
3789
3790static PyObject *
3791izip_longest_next(iziplongestobject *lz)
3792{
3793 Py_ssize_t i;
3794 Py_ssize_t tuplesize = lz->tuplesize;
3795 PyObject *result = lz->result;
3796 PyObject *it;
3797 PyObject *item;
3798 PyObject *olditem;
3799
3800 if (tuplesize == 0)
3801 return NULL;
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003802 if (lz->numactive == 0)
3803 return NULL;
Christian Heimese93237d2007-12-19 02:37:44 +00003804 if (Py_REFCNT(result) == 1) {
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003805 Py_INCREF(result);
3806 for (i=0 ; i < tuplesize ; i++) {
3807 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003808 if (it == NULL) {
3809 Py_INCREF(lz->fillvalue);
3810 item = lz->fillvalue;
3811 } else {
Christian Heimese93237d2007-12-19 02:37:44 +00003812 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003813 if (item == NULL) {
3814 lz->numactive -= 1;
3815 if (lz->numactive == 0) {
3816 Py_DECREF(result);
3817 return NULL;
3818 } else {
3819 Py_INCREF(lz->fillvalue);
3820 item = lz->fillvalue;
3821 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3822 Py_DECREF(it);
3823 }
3824 }
3825 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003826 olditem = PyTuple_GET_ITEM(result, i);
3827 PyTuple_SET_ITEM(result, i, item);
3828 Py_DECREF(olditem);
3829 }
3830 } else {
3831 result = PyTuple_New(tuplesize);
3832 if (result == NULL)
3833 return NULL;
3834 for (i=0 ; i < tuplesize ; i++) {
3835 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003836 if (it == NULL) {
3837 Py_INCREF(lz->fillvalue);
3838 item = lz->fillvalue;
3839 } else {
Christian Heimese93237d2007-12-19 02:37:44 +00003840 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003841 if (item == NULL) {
3842 lz->numactive -= 1;
3843 if (lz->numactive == 0) {
3844 Py_DECREF(result);
3845 return NULL;
3846 } else {
3847 Py_INCREF(lz->fillvalue);
3848 item = lz->fillvalue;
3849 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3850 Py_DECREF(it);
3851 }
3852 }
3853 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003854 PyTuple_SET_ITEM(result, i, item);
3855 }
3856 }
3857 return result;
3858}
3859
3860PyDoc_STRVAR(izip_longest_doc,
3861"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3862\n\
3863Return an izip_longest object whose .next() method returns a tuple where\n\
3864the i-th element comes from the i-th iterable argument. The .next()\n\
3865method continues until the longest iterable in the argument sequence\n\
3866is exhausted and then it raises StopIteration. When the shorter iterables\n\
3867are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3868defaults to None or can be specified by a keyword argument.\n\
3869");
3870
3871static PyTypeObject iziplongest_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003872 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003873 "itertools.izip_longest", /* tp_name */
3874 sizeof(iziplongestobject), /* tp_basicsize */
3875 0, /* tp_itemsize */
3876 /* methods */
3877 (destructor)izip_longest_dealloc, /* tp_dealloc */
3878 0, /* tp_print */
3879 0, /* tp_getattr */
3880 0, /* tp_setattr */
3881 0, /* tp_compare */
3882 0, /* tp_repr */
3883 0, /* tp_as_number */
3884 0, /* tp_as_sequence */
3885 0, /* tp_as_mapping */
3886 0, /* tp_hash */
3887 0, /* tp_call */
3888 0, /* tp_str */
3889 PyObject_GenericGetAttr, /* tp_getattro */
3890 0, /* tp_setattro */
3891 0, /* tp_as_buffer */
3892 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3893 Py_TPFLAGS_BASETYPE, /* tp_flags */
3894 izip_longest_doc, /* tp_doc */
3895 (traverseproc)izip_longest_traverse, /* tp_traverse */
3896 0, /* tp_clear */
3897 0, /* tp_richcompare */
3898 0, /* tp_weaklistoffset */
3899 PyObject_SelfIter, /* tp_iter */
3900 (iternextfunc)izip_longest_next, /* tp_iternext */
3901 0, /* tp_methods */
3902 0, /* tp_members */
3903 0, /* tp_getset */
3904 0, /* tp_base */
3905 0, /* tp_dict */
3906 0, /* tp_descr_get */
3907 0, /* tp_descr_set */
3908 0, /* tp_dictoffset */
3909 0, /* tp_init */
3910 0, /* tp_alloc */
3911 izip_longest_new, /* tp_new */
3912 PyObject_GC_Del, /* tp_free */
3913};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003914
3915/* module level code ********************************************************/
3916
3917PyDoc_STRVAR(module_doc,
3918"Functional tools for creating and using iterators.\n\
3919\n\
3920Infinite iterators:\n\
3921count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003922cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003923repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003924\n\
3925Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00003926chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettingere76816b2009-01-29 03:43:44 +00003927compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00003928dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
3929groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003930ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
3931ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003932islice(seq, [start,] stop [, step]) --> elements from\n\
3933 seq[start:stop:step]\n\
3934imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
3935starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003936tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003937takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00003938izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
3939izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
3940\n\
3941Combinatoric generators:\n\
3942product(p, q, ... [repeat=1]) --> cartesian product\n\
3943permutations(p[, r])\n\
3944combinations(p[, r])\n\
3945combinations_with_replacement(p[, r])\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003946");
3947
3948
Raymond Hettingerad983e72003-11-12 14:32:26 +00003949static PyMethodDef module_methods[] = {
3950 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3951 {NULL, NULL} /* sentinel */
3952};
3953
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003954PyMODINIT_FUNC
3955inititertools(void)
3956{
Raymond Hettinger60eca932003-02-09 06:40:58 +00003957 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003958 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003959 char *name;
3960 PyTypeObject *typelist[] = {
Raymond Hettinger93e804d2008-02-26 23:40:50 +00003961 &combinations_type,
Raymond Hettingerd081abc2009-01-27 02:58:49 +00003962 &cwr_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003963 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003964 &dropwhile_type,
3965 &takewhile_type,
3966 &islice_type,
3967 &starmap_type,
3968 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003969 &chain_type,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00003970 &compress_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003971 &ifilter_type,
3972 &ifilterfalse_type,
3973 &count_type,
3974 &izip_type,
Raymond Hettinger50986cc2008-02-22 03:16:42 +00003975 &iziplongest_type,
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00003976 &permutations_type,
David Wolever2724ab92008-03-19 02:35:45 +00003977 &product_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003978 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003979 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003980 NULL
3981 };
3982
Christian Heimese93237d2007-12-19 02:37:44 +00003983 Py_TYPE(&teedataobject_type) = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00003984 m = Py_InitModule3("itertools", module_methods, module_doc);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00003985 if (m == NULL)
3986 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003987
Raymond Hettinger60eca932003-02-09 06:40:58 +00003988 for (i=0 ; typelist[i] != NULL ; i++) {
3989 if (PyType_Ready(typelist[i]) < 0)
3990 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003991 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003992 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003993 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003994 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003995 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003996
3997 if (PyType_Ready(&teedataobject_type) < 0)
3998 return;
3999 if (PyType_Ready(&tee_type) < 0)
4000 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00004001 if (PyType_Ready(&_grouper_type) < 0)
4002 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004003}