blob: 4aa02b62012fba181e716e4bbbebeb12311eaf60 [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) {
Benjamin Petersona7b0c122009-02-16 21:23:04 +0000748 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
749 Py_DECREF(item);
750 return NULL;
751 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000752 return item;
753 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000754 if (PyErr_Occurred()) {
755 if (PyErr_ExceptionMatches(PyExc_StopIteration))
756 PyErr_Clear();
757 else
758 return NULL;
759 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000760 if (PyList_Size(lz->saved) == 0)
761 return NULL;
762 it = PyObject_GetIter(lz->saved);
763 if (it == NULL)
764 return NULL;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000765 tmp = lz->it;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000766 lz->it = it;
767 lz->firstpass = 1;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000768 Py_DECREF(tmp);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000769 }
770}
771
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000772PyDoc_STRVAR(cycle_doc,
773"cycle(iterable) --> cycle object\n\
774\n\
775Return elements from the iterable until it is exhausted.\n\
776Then repeat the sequence indefinitely.");
777
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000778static PyTypeObject cycle_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000779 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000780 "itertools.cycle", /* tp_name */
781 sizeof(cycleobject), /* tp_basicsize */
782 0, /* tp_itemsize */
783 /* methods */
784 (destructor)cycle_dealloc, /* tp_dealloc */
785 0, /* tp_print */
786 0, /* tp_getattr */
787 0, /* tp_setattr */
788 0, /* tp_compare */
789 0, /* tp_repr */
790 0, /* tp_as_number */
791 0, /* tp_as_sequence */
792 0, /* tp_as_mapping */
793 0, /* tp_hash */
794 0, /* tp_call */
795 0, /* tp_str */
796 PyObject_GenericGetAttr, /* tp_getattro */
797 0, /* tp_setattro */
798 0, /* tp_as_buffer */
799 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
800 Py_TPFLAGS_BASETYPE, /* tp_flags */
801 cycle_doc, /* tp_doc */
802 (traverseproc)cycle_traverse, /* tp_traverse */
803 0, /* tp_clear */
804 0, /* tp_richcompare */
805 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000806 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000807 (iternextfunc)cycle_next, /* tp_iternext */
808 0, /* tp_methods */
809 0, /* tp_members */
810 0, /* tp_getset */
811 0, /* tp_base */
812 0, /* tp_dict */
813 0, /* tp_descr_get */
814 0, /* tp_descr_set */
815 0, /* tp_dictoffset */
816 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000817 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000818 cycle_new, /* tp_new */
819 PyObject_GC_Del, /* tp_free */
820};
821
822
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000823/* dropwhile object **********************************************************/
824
825typedef struct {
826 PyObject_HEAD
827 PyObject *func;
828 PyObject *it;
829 long start;
830} dropwhileobject;
831
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000832static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000833
834static PyObject *
835dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
836{
837 PyObject *func, *seq;
838 PyObject *it;
839 dropwhileobject *lz;
840
Georg Brandlb84c1372007-01-21 10:28:43 +0000841 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000842 return NULL;
843
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000844 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
845 return NULL;
846
847 /* Get iterator. */
848 it = PyObject_GetIter(seq);
849 if (it == NULL)
850 return NULL;
851
852 /* create dropwhileobject structure */
853 lz = (dropwhileobject *)type->tp_alloc(type, 0);
854 if (lz == NULL) {
855 Py_DECREF(it);
856 return NULL;
857 }
858 Py_INCREF(func);
859 lz->func = func;
860 lz->it = it;
861 lz->start = 0;
862
863 return (PyObject *)lz;
864}
865
866static void
867dropwhile_dealloc(dropwhileobject *lz)
868{
869 PyObject_GC_UnTrack(lz);
870 Py_XDECREF(lz->func);
871 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +0000872 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000873}
874
875static int
876dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
877{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000878 Py_VISIT(lz->it);
879 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000880 return 0;
881}
882
883static PyObject *
884dropwhile_next(dropwhileobject *lz)
885{
886 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000887 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000888 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000889 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000890
Christian Heimese93237d2007-12-19 02:37:44 +0000891 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000892 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000893 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000894 if (item == NULL)
895 return NULL;
896 if (lz->start == 1)
897 return item;
898
899 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
900 if (good == NULL) {
901 Py_DECREF(item);
902 return NULL;
903 }
904 ok = PyObject_IsTrue(good);
905 Py_DECREF(good);
906 if (!ok) {
907 lz->start = 1;
908 return item;
909 }
910 Py_DECREF(item);
911 }
912}
913
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000914PyDoc_STRVAR(dropwhile_doc,
915"dropwhile(predicate, iterable) --> dropwhile object\n\
916\n\
917Drop items from the iterable while predicate(item) is true.\n\
918Afterwards, return every element until the iterable is exhausted.");
919
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000920static PyTypeObject dropwhile_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000921 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000922 "itertools.dropwhile", /* tp_name */
923 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000924 0, /* tp_itemsize */
925 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000926 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000927 0, /* tp_print */
928 0, /* tp_getattr */
929 0, /* tp_setattr */
930 0, /* tp_compare */
931 0, /* tp_repr */
932 0, /* tp_as_number */
933 0, /* tp_as_sequence */
934 0, /* tp_as_mapping */
935 0, /* tp_hash */
936 0, /* tp_call */
937 0, /* tp_str */
938 PyObject_GenericGetAttr, /* tp_getattro */
939 0, /* tp_setattro */
940 0, /* tp_as_buffer */
941 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
942 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000943 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000944 (traverseproc)dropwhile_traverse, /* tp_traverse */
945 0, /* tp_clear */
946 0, /* tp_richcompare */
947 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000948 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000949 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000950 0, /* tp_methods */
951 0, /* tp_members */
952 0, /* tp_getset */
953 0, /* tp_base */
954 0, /* tp_dict */
955 0, /* tp_descr_get */
956 0, /* tp_descr_set */
957 0, /* tp_dictoffset */
958 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000959 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000960 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000961 PyObject_GC_Del, /* tp_free */
962};
963
964
965/* takewhile object **********************************************************/
966
967typedef struct {
968 PyObject_HEAD
969 PyObject *func;
970 PyObject *it;
971 long stop;
972} takewhileobject;
973
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000974static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000975
976static PyObject *
977takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
978{
979 PyObject *func, *seq;
980 PyObject *it;
981 takewhileobject *lz;
982
Georg Brandlb84c1372007-01-21 10:28:43 +0000983 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000984 return NULL;
985
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000986 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
987 return NULL;
988
989 /* Get iterator. */
990 it = PyObject_GetIter(seq);
991 if (it == NULL)
992 return NULL;
993
994 /* create takewhileobject structure */
995 lz = (takewhileobject *)type->tp_alloc(type, 0);
996 if (lz == NULL) {
997 Py_DECREF(it);
998 return NULL;
999 }
1000 Py_INCREF(func);
1001 lz->func = func;
1002 lz->it = it;
1003 lz->stop = 0;
1004
1005 return (PyObject *)lz;
1006}
1007
1008static void
1009takewhile_dealloc(takewhileobject *lz)
1010{
1011 PyObject_GC_UnTrack(lz);
1012 Py_XDECREF(lz->func);
1013 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00001014 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001015}
1016
1017static int
1018takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1019{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001020 Py_VISIT(lz->it);
1021 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001022 return 0;
1023}
1024
1025static PyObject *
1026takewhile_next(takewhileobject *lz)
1027{
1028 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001029 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001030 long ok;
1031
1032 if (lz->stop == 1)
1033 return NULL;
1034
Christian Heimese93237d2007-12-19 02:37:44 +00001035 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001036 if (item == NULL)
1037 return NULL;
1038
1039 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1040 if (good == NULL) {
1041 Py_DECREF(item);
1042 return NULL;
1043 }
1044 ok = PyObject_IsTrue(good);
1045 Py_DECREF(good);
1046 if (ok)
1047 return item;
1048 Py_DECREF(item);
1049 lz->stop = 1;
1050 return NULL;
1051}
1052
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053PyDoc_STRVAR(takewhile_doc,
1054"takewhile(predicate, iterable) --> takewhile object\n\
1055\n\
1056Return successive entries from an iterable as long as the \n\
1057predicate evaluates to true for each entry.");
1058
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001059static PyTypeObject takewhile_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001060 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001061 "itertools.takewhile", /* tp_name */
1062 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001063 0, /* tp_itemsize */
1064 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001065 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001066 0, /* tp_print */
1067 0, /* tp_getattr */
1068 0, /* tp_setattr */
1069 0, /* tp_compare */
1070 0, /* tp_repr */
1071 0, /* tp_as_number */
1072 0, /* tp_as_sequence */
1073 0, /* tp_as_mapping */
1074 0, /* tp_hash */
1075 0, /* tp_call */
1076 0, /* tp_str */
1077 PyObject_GenericGetAttr, /* tp_getattro */
1078 0, /* tp_setattro */
1079 0, /* tp_as_buffer */
1080 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1081 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001082 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001083 (traverseproc)takewhile_traverse, /* tp_traverse */
1084 0, /* tp_clear */
1085 0, /* tp_richcompare */
1086 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001087 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001088 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001089 0, /* tp_methods */
1090 0, /* tp_members */
1091 0, /* tp_getset */
1092 0, /* tp_base */
1093 0, /* tp_dict */
1094 0, /* tp_descr_get */
1095 0, /* tp_descr_set */
1096 0, /* tp_dictoffset */
1097 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001098 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001099 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001100 PyObject_GC_Del, /* tp_free */
1101};
1102
1103
1104/* islice object ************************************************************/
1105
1106typedef struct {
1107 PyObject_HEAD
1108 PyObject *it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001109 Py_ssize_t next;
1110 Py_ssize_t stop;
1111 Py_ssize_t step;
1112 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001113} isliceobject;
1114
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001115static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116
1117static PyObject *
1118islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1119{
1120 PyObject *seq;
Jack Diederich6c433a92006-05-26 11:15:17 +00001121 Py_ssize_t start=0, stop=-1, step=1;
Raymond Hettingerb2594052004-12-05 09:25:51 +00001122 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001123 Py_ssize_t numargs;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001124 isliceobject *lz;
1125
Georg Brandlb84c1372007-01-21 10:28:43 +00001126 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001127 return NULL;
1128
Raymond Hettingerb2594052004-12-05 09:25:51 +00001129 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001130 return NULL;
1131
Raymond Hettingerb2594052004-12-05 09:25:51 +00001132 numargs = PyTuple_Size(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001133 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001134 if (a1 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001135 stop = PyInt_AsSsize_t(a1);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001136 if (stop == -1) {
1137 if (PyErr_Occurred())
1138 PyErr_Clear();
1139 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001140 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001141 return NULL;
1142 }
1143 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001144 } else {
Raymond Hettingerb2594052004-12-05 09:25:51 +00001145 if (a1 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001146 start = PyInt_AsSsize_t(a1);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001147 if (start == -1 && PyErr_Occurred())
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001148 PyErr_Clear();
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001149 if (a2 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001150 stop = PyInt_AsSsize_t(a2);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001151 if (stop == -1) {
1152 if (PyErr_Occurred())
1153 PyErr_Clear();
1154 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001155 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001156 return NULL;
1157 }
1158 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001159 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001160 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001161 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001162 "Indices for islice() must be non-negative integers or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001163 return NULL;
1164 }
1165
Raymond Hettingerb2594052004-12-05 09:25:51 +00001166 if (a3 != NULL) {
1167 if (a3 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001168 step = PyInt_AsSsize_t(a3);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001169 if (step == -1 && PyErr_Occurred())
1170 PyErr_Clear();
1171 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001172 if (step<1) {
1173 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001174 "Step for islice() must be a positive integer or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001175 return NULL;
1176 }
1177
1178 /* Get iterator. */
1179 it = PyObject_GetIter(seq);
1180 if (it == NULL)
1181 return NULL;
1182
1183 /* create isliceobject structure */
1184 lz = (isliceobject *)type->tp_alloc(type, 0);
1185 if (lz == NULL) {
1186 Py_DECREF(it);
1187 return NULL;
1188 }
1189 lz->it = it;
1190 lz->next = start;
1191 lz->stop = stop;
1192 lz->step = step;
1193 lz->cnt = 0L;
1194
1195 return (PyObject *)lz;
1196}
1197
1198static void
1199islice_dealloc(isliceobject *lz)
1200{
1201 PyObject_GC_UnTrack(lz);
1202 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00001203 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204}
1205
1206static int
1207islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1208{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001209 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001210 return 0;
1211}
1212
1213static PyObject *
1214islice_next(isliceobject *lz)
1215{
1216 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001217 PyObject *it = lz->it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001218 Py_ssize_t oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001219 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001220
Christian Heimese93237d2007-12-19 02:37:44 +00001221 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001222 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001223 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001224 if (item == NULL)
1225 return NULL;
1226 Py_DECREF(item);
1227 lz->cnt++;
1228 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001229 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001230 return NULL;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001231 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001232 if (item == NULL)
1233 return NULL;
1234 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001235 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001236 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001237 if (lz->next < oldnext) /* Check for overflow */
1238 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001239 return item;
1240}
1241
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001242PyDoc_STRVAR(islice_doc,
1243"islice(iterable, [start,] stop [, step]) --> islice object\n\
1244\n\
1245Return an iterator whose next() method returns selected values from an\n\
1246iterable. If start is specified, will skip all preceding elements;\n\
1247otherwise, start defaults to zero. Step defaults to one. If\n\
1248specified as another value, step determines how many values are \n\
1249skipped between successive calls. Works like a slice() on a list\n\
1250but returns an iterator.");
1251
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001252static PyTypeObject islice_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001253 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001254 "itertools.islice", /* tp_name */
1255 sizeof(isliceobject), /* tp_basicsize */
1256 0, /* tp_itemsize */
1257 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001258 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001259 0, /* tp_print */
1260 0, /* tp_getattr */
1261 0, /* tp_setattr */
1262 0, /* tp_compare */
1263 0, /* tp_repr */
1264 0, /* tp_as_number */
1265 0, /* tp_as_sequence */
1266 0, /* tp_as_mapping */
1267 0, /* tp_hash */
1268 0, /* tp_call */
1269 0, /* tp_str */
1270 PyObject_GenericGetAttr, /* tp_getattro */
1271 0, /* tp_setattro */
1272 0, /* tp_as_buffer */
1273 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1274 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001275 islice_doc, /* tp_doc */
1276 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001277 0, /* tp_clear */
1278 0, /* tp_richcompare */
1279 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001280 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001281 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001282 0, /* tp_methods */
1283 0, /* tp_members */
1284 0, /* tp_getset */
1285 0, /* tp_base */
1286 0, /* tp_dict */
1287 0, /* tp_descr_get */
1288 0, /* tp_descr_set */
1289 0, /* tp_dictoffset */
1290 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001291 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001292 islice_new, /* tp_new */
1293 PyObject_GC_Del, /* tp_free */
1294};
1295
1296
1297/* starmap object ************************************************************/
1298
1299typedef struct {
1300 PyObject_HEAD
1301 PyObject *func;
1302 PyObject *it;
1303} starmapobject;
1304
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001305static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001306
1307static PyObject *
1308starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1309{
1310 PyObject *func, *seq;
1311 PyObject *it;
1312 starmapobject *lz;
1313
Georg Brandlb84c1372007-01-21 10:28:43 +00001314 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001315 return NULL;
1316
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001317 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1318 return NULL;
1319
1320 /* Get iterator. */
1321 it = PyObject_GetIter(seq);
1322 if (it == NULL)
1323 return NULL;
1324
1325 /* create starmapobject structure */
1326 lz = (starmapobject *)type->tp_alloc(type, 0);
1327 if (lz == NULL) {
1328 Py_DECREF(it);
1329 return NULL;
1330 }
1331 Py_INCREF(func);
1332 lz->func = func;
1333 lz->it = it;
1334
1335 return (PyObject *)lz;
1336}
1337
1338static void
1339starmap_dealloc(starmapobject *lz)
1340{
1341 PyObject_GC_UnTrack(lz);
1342 Py_XDECREF(lz->func);
1343 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00001344 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001345}
1346
1347static int
1348starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1349{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001350 Py_VISIT(lz->it);
1351 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001352 return 0;
1353}
1354
1355static PyObject *
1356starmap_next(starmapobject *lz)
1357{
1358 PyObject *args;
1359 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001360 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001361
Christian Heimese93237d2007-12-19 02:37:44 +00001362 args = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001363 if (args == NULL)
1364 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001365 if (!PyTuple_CheckExact(args)) {
Raymond Hettinger47317092008-01-17 03:02:14 +00001366 PyObject *newargs = PySequence_Tuple(args);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001367 Py_DECREF(args);
Raymond Hettinger47317092008-01-17 03:02:14 +00001368 if (newargs == NULL)
1369 return NULL;
1370 args = newargs;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001371 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001372 result = PyObject_Call(lz->func, args, NULL);
1373 Py_DECREF(args);
1374 return result;
1375}
1376
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001377PyDoc_STRVAR(starmap_doc,
1378"starmap(function, sequence) --> starmap object\n\
1379\n\
1380Return an iterator whose values are returned from the function evaluated\n\
1381with a argument tuple taken from the given sequence.");
1382
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001383static PyTypeObject starmap_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001384 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001385 "itertools.starmap", /* tp_name */
1386 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001387 0, /* tp_itemsize */
1388 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001389 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001390 0, /* tp_print */
1391 0, /* tp_getattr */
1392 0, /* tp_setattr */
1393 0, /* tp_compare */
1394 0, /* tp_repr */
1395 0, /* tp_as_number */
1396 0, /* tp_as_sequence */
1397 0, /* tp_as_mapping */
1398 0, /* tp_hash */
1399 0, /* tp_call */
1400 0, /* tp_str */
1401 PyObject_GenericGetAttr, /* tp_getattro */
1402 0, /* tp_setattro */
1403 0, /* tp_as_buffer */
1404 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1405 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001406 starmap_doc, /* tp_doc */
1407 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001408 0, /* tp_clear */
1409 0, /* tp_richcompare */
1410 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001411 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001412 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001413 0, /* tp_methods */
1414 0, /* tp_members */
1415 0, /* tp_getset */
1416 0, /* tp_base */
1417 0, /* tp_dict */
1418 0, /* tp_descr_get */
1419 0, /* tp_descr_set */
1420 0, /* tp_dictoffset */
1421 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001422 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001423 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001424 PyObject_GC_Del, /* tp_free */
1425};
1426
1427
1428/* imap object ************************************************************/
1429
1430typedef struct {
1431 PyObject_HEAD
1432 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001433 PyObject *func;
1434} imapobject;
1435
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001436static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001437
1438static PyObject *
1439imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1440{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001441 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001442 imapobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001443 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001444
Georg Brandlb84c1372007-01-21 10:28:43 +00001445 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001446 return NULL;
1447
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001448 numargs = PyTuple_Size(args);
1449 if (numargs < 2) {
1450 PyErr_SetString(PyExc_TypeError,
1451 "imap() must have at least two arguments.");
1452 return NULL;
1453 }
1454
1455 iters = PyTuple_New(numargs-1);
1456 if (iters == NULL)
1457 return NULL;
1458
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459 for (i=1 ; i<numargs ; i++) {
1460 /* Get iterator. */
1461 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1462 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001463 Py_DECREF(iters);
1464 return NULL;
1465 }
1466 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001467 }
1468
1469 /* create imapobject structure */
1470 lz = (imapobject *)type->tp_alloc(type, 0);
1471 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001472 Py_DECREF(iters);
1473 return NULL;
1474 }
1475 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001476 func = PyTuple_GET_ITEM(args, 0);
1477 Py_INCREF(func);
1478 lz->func = func;
1479
1480 return (PyObject *)lz;
1481}
1482
1483static void
1484imap_dealloc(imapobject *lz)
1485{
1486 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001487 Py_XDECREF(lz->iters);
1488 Py_XDECREF(lz->func);
Christian Heimese93237d2007-12-19 02:37:44 +00001489 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001490}
1491
1492static int
1493imap_traverse(imapobject *lz, visitproc visit, void *arg)
1494{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001495 Py_VISIT(lz->iters);
1496 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001497 return 0;
1498}
1499
Raymond Hettinger2012f172003-02-07 05:32:58 +00001500/*
1501imap() is an iterator version of __builtins__.map() except that it does
1502not have the None fill-in feature. That was intentionally left out for
1503the following reasons:
1504
1505 1) Itertools are designed to be easily combined and chained together.
1506 Having all tools stop with the shortest input is a unifying principle
1507 that makes it easier to combine finite iterators (supplying data) with
1508 infinite iterators like count() and repeat() (for supplying sequential
1509 or constant arguments to a function).
1510
1511 2) In typical use cases for combining itertools, having one finite data
1512 supplier run out before another is likely to be an error condition which
1513 should not pass silently by automatically supplying None.
1514
1515 3) The use cases for automatic None fill-in are rare -- not many functions
1516 do something useful when a parameter suddenly switches type and becomes
1517 None.
1518
1519 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001520 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001521
1522 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1523*/
1524
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001525static PyObject *
1526imap_next(imapobject *lz)
1527{
1528 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001529 PyObject *argtuple;
1530 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001531 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001532
1533 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001534 argtuple = PyTuple_New(numargs);
1535 if (argtuple == NULL)
1536 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001537
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001538 for (i=0 ; i<numargs ; i++) {
1539 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1540 if (val == NULL) {
1541 Py_DECREF(argtuple);
1542 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001543 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001544 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001545 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001546 if (lz->func == Py_None)
1547 return argtuple;
1548 result = PyObject_Call(lz->func, argtuple, NULL);
1549 Py_DECREF(argtuple);
1550 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001551}
1552
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001553PyDoc_STRVAR(imap_doc,
1554"imap(func, *iterables) --> imap object\n\
1555\n\
1556Make an iterator that computes the function using arguments from\n\
1557each of the iterables. Like map() except that it returns\n\
1558an iterator instead of a list and that it stops when the shortest\n\
1559iterable is exhausted instead of filling in None for shorter\n\
1560iterables.");
1561
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001562static PyTypeObject imap_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001563 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001564 "itertools.imap", /* tp_name */
1565 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001566 0, /* tp_itemsize */
1567 /* methods */
1568 (destructor)imap_dealloc, /* tp_dealloc */
1569 0, /* tp_print */
1570 0, /* tp_getattr */
1571 0, /* tp_setattr */
1572 0, /* tp_compare */
1573 0, /* tp_repr */
1574 0, /* tp_as_number */
1575 0, /* tp_as_sequence */
1576 0, /* tp_as_mapping */
1577 0, /* tp_hash */
1578 0, /* tp_call */
1579 0, /* tp_str */
1580 PyObject_GenericGetAttr, /* tp_getattro */
1581 0, /* tp_setattro */
1582 0, /* tp_as_buffer */
1583 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1584 Py_TPFLAGS_BASETYPE, /* tp_flags */
1585 imap_doc, /* tp_doc */
1586 (traverseproc)imap_traverse, /* tp_traverse */
1587 0, /* tp_clear */
1588 0, /* tp_richcompare */
1589 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001590 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001591 (iternextfunc)imap_next, /* tp_iternext */
1592 0, /* tp_methods */
1593 0, /* tp_members */
1594 0, /* tp_getset */
1595 0, /* tp_base */
1596 0, /* tp_dict */
1597 0, /* tp_descr_get */
1598 0, /* tp_descr_set */
1599 0, /* tp_dictoffset */
1600 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001601 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001602 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001603 PyObject_GC_Del, /* tp_free */
1604};
1605
1606
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001607/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001608
1609typedef struct {
1610 PyObject_HEAD
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001611 PyObject *source; /* Iterator over input iterables */
1612 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001613} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001614
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001615static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001616
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001617static PyObject *
1618chain_new_internal(PyTypeObject *type, PyObject *source)
1619{
1620 chainobject *lz;
1621
1622 lz = (chainobject *)type->tp_alloc(type, 0);
1623 if (lz == NULL) {
1624 Py_DECREF(source);
1625 return NULL;
1626 }
1627
1628 lz->source = source;
1629 lz->active = NULL;
1630 return (PyObject *)lz;
1631}
1632
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001633static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001634chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001635{
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001636 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001637
Georg Brandlb84c1372007-01-21 10:28:43 +00001638 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001639 return NULL;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001640
1641 source = PyObject_GetIter(args);
1642 if (source == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001643 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001644
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001645 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001646}
1647
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001648static PyObject *
1649chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1650{
1651 PyObject *source;
1652
1653 source = PyObject_GetIter(arg);
1654 if (source == NULL)
1655 return NULL;
1656
1657 return chain_new_internal(type, source);
1658}
1659
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001660static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001661chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001662{
1663 PyObject_GC_UnTrack(lz);
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001664 Py_XDECREF(lz->active);
1665 Py_XDECREF(lz->source);
Christian Heimese93237d2007-12-19 02:37:44 +00001666 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001667}
1668
Raymond Hettinger2012f172003-02-07 05:32:58 +00001669static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001670chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001671{
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001672 Py_VISIT(lz->source);
1673 Py_VISIT(lz->active);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001674 return 0;
1675}
1676
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001677static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001678chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001679{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001680 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001681
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001682 if (lz->source == NULL)
1683 return NULL; /* already stopped */
1684
1685 if (lz->active == NULL) {
1686 PyObject *iterable = PyIter_Next(lz->source);
1687 if (iterable == NULL) {
1688 Py_CLEAR(lz->source);
1689 return NULL; /* no more input sources */
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001690 }
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001691 lz->active = PyObject_GetIter(iterable);
Raymond Hettingerf1cca2b2008-03-04 22:29:44 +00001692 Py_DECREF(iterable);
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001693 if (lz->active == NULL) {
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001694 Py_CLEAR(lz->source);
1695 return NULL; /* input not iterable */
1696 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001697 }
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001698 item = PyIter_Next(lz->active);
1699 if (item != NULL)
1700 return item;
1701 if (PyErr_Occurred()) {
1702 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1703 PyErr_Clear();
1704 else
1705 return NULL; /* input raised an exception */
1706 }
1707 Py_CLEAR(lz->active);
1708 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001709}
1710
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001711PyDoc_STRVAR(chain_doc,
1712"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001713\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001714Return a chain object whose .next() method returns elements from the\n\
1715first iterable until it is exhausted, then elements from the next\n\
1716iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001717
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001718PyDoc_STRVAR(chain_from_iterable_doc,
1719"chain.from_iterable(iterable) --> chain object\n\
1720\n\
1721Alternate chain() contructor taking a single iterable argument\n\
1722that evaluates lazily.");
1723
1724static PyMethodDef chain_methods[] = {
1725 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1726 chain_from_iterable_doc},
1727 {NULL, NULL} /* sentinel */
1728};
1729
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001730static PyTypeObject chain_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001731 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001732 "itertools.chain", /* tp_name */
1733 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001734 0, /* tp_itemsize */
1735 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001736 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001737 0, /* tp_print */
1738 0, /* tp_getattr */
1739 0, /* tp_setattr */
1740 0, /* tp_compare */
1741 0, /* tp_repr */
1742 0, /* tp_as_number */
1743 0, /* tp_as_sequence */
1744 0, /* tp_as_mapping */
1745 0, /* tp_hash */
1746 0, /* tp_call */
1747 0, /* tp_str */
1748 PyObject_GenericGetAttr, /* tp_getattro */
1749 0, /* tp_setattro */
1750 0, /* tp_as_buffer */
1751 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1752 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001753 chain_doc, /* tp_doc */
1754 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001755 0, /* tp_clear */
1756 0, /* tp_richcompare */
1757 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001758 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001759 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001760 chain_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001761 0, /* tp_members */
1762 0, /* tp_getset */
1763 0, /* tp_base */
1764 0, /* tp_dict */
1765 0, /* tp_descr_get */
1766 0, /* tp_descr_set */
1767 0, /* tp_dictoffset */
1768 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001769 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001770 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001771 PyObject_GC_Del, /* tp_free */
1772};
1773
1774
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001775/* product object ************************************************************/
1776
1777typedef struct {
1778 PyObject_HEAD
1779 PyObject *pools; /* tuple of pool tuples */
Raymond Hettinger532316d2008-02-23 04:03:50 +00001780 Py_ssize_t *indices; /* one index per pool */
1781 PyObject *result; /* most recently returned result tuple */
1782 int stopped; /* set to 1 when the product iterator is exhausted */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001783} productobject;
1784
1785static PyTypeObject product_type;
1786
1787static PyObject *
1788product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1789{
1790 productobject *lz;
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001791 Py_ssize_t nargs, npools, repeat=1;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001792 PyObject *pools = NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001793 Py_ssize_t *indices = NULL;
1794 Py_ssize_t i;
1795
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001796 if (kwds != NULL) {
1797 char *kwlist[] = {"repeat", 0};
1798 PyObject *tmpargs = PyTuple_New(0);
1799 if (tmpargs == NULL)
1800 return NULL;
1801 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1802 Py_DECREF(tmpargs);
1803 return NULL;
1804 }
1805 Py_DECREF(tmpargs);
1806 if (repeat < 0) {
1807 PyErr_SetString(PyExc_ValueError,
1808 "repeat argument cannot be negative");
1809 return NULL;
1810 }
1811 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001812
1813 assert(PyTuple_Check(args));
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001814 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1815 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001816
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001817 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
Raymond Hettinger61024b92008-03-02 11:57:16 +00001818 if (indices == NULL) {
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001819 PyErr_NoMemory();
1820 goto error;
1821 }
1822
1823 pools = PyTuple_New(npools);
1824 if (pools == NULL)
1825 goto error;
1826
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001827 for (i=0; i < nargs ; ++i) {
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001828 PyObject *item = PyTuple_GET_ITEM(args, i);
1829 PyObject *pool = PySequence_Tuple(item);
1830 if (pool == NULL)
1831 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001832 PyTuple_SET_ITEM(pools, i, pool);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001833 indices[i] = 0;
1834 }
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001835 for ( ; i < npools; ++i) {
1836 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1837 Py_INCREF(pool);
1838 PyTuple_SET_ITEM(pools, i, pool);
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001839 indices[i] = 0;
1840 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001841
1842 /* create productobject structure */
1843 lz = (productobject *)type->tp_alloc(type, 0);
Raymond Hettinger3bd77122008-02-27 01:08:04 +00001844 if (lz == NULL)
Raymond Hettinger73d79632008-02-23 02:20:41 +00001845 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001846
1847 lz->pools = pools;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001848 lz->indices = indices;
1849 lz->result = NULL;
1850 lz->stopped = 0;
1851
1852 return (PyObject *)lz;
1853
1854error:
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001855 if (indices != NULL)
1856 PyMem_Free(indices);
1857 Py_XDECREF(pools);
1858 return NULL;
1859}
1860
1861static void
1862product_dealloc(productobject *lz)
1863{
1864 PyObject_GC_UnTrack(lz);
1865 Py_XDECREF(lz->pools);
1866 Py_XDECREF(lz->result);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00001867 if (lz->indices != NULL)
1868 PyMem_Free(lz->indices);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001869 Py_TYPE(lz)->tp_free(lz);
1870}
1871
1872static int
1873product_traverse(productobject *lz, visitproc visit, void *arg)
1874{
1875 Py_VISIT(lz->pools);
1876 Py_VISIT(lz->result);
1877 return 0;
1878}
1879
1880static PyObject *
1881product_next(productobject *lz)
1882{
1883 PyObject *pool;
1884 PyObject *elem;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001885 PyObject *oldelem;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001886 PyObject *pools = lz->pools;
1887 PyObject *result = lz->result;
1888 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1889 Py_ssize_t i;
1890
1891 if (lz->stopped)
1892 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001893
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001894 if (result == NULL) {
Raymond Hettinger73d79632008-02-23 02:20:41 +00001895 /* On the first pass, return an initial tuple filled with the
Raymond Hettingerd553d852008-03-04 04:17:08 +00001896 first element from each pool. */
Raymond Hettinger73d79632008-02-23 02:20:41 +00001897 result = PyTuple_New(npools);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001898 if (result == NULL)
1899 goto empty;
1900 lz->result = result;
1901 for (i=0; i < npools; i++) {
1902 pool = PyTuple_GET_ITEM(pools, i);
1903 if (PyTuple_GET_SIZE(pool) == 0)
1904 goto empty;
1905 elem = PyTuple_GET_ITEM(pool, 0);
1906 Py_INCREF(elem);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001907 PyTuple_SET_ITEM(result, i, elem);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001908 }
1909 } else {
1910 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001911
1912 /* Copy the previous result tuple or re-use it if available */
1913 if (Py_REFCNT(result) > 1) {
1914 PyObject *old_result = result;
1915 result = PyTuple_New(npools);
1916 if (result == NULL)
1917 goto empty;
1918 lz->result = result;
1919 for (i=0; i < npools; i++) {
1920 elem = PyTuple_GET_ITEM(old_result, i);
1921 Py_INCREF(elem);
1922 PyTuple_SET_ITEM(result, i, elem);
1923 }
1924 Py_DECREF(old_result);
1925 }
1926 /* Now, we've got the only copy so we can update it in-place */
Raymond Hettingere3fabd12008-03-02 12:02:19 +00001927 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001928
1929 /* Update the pool indices right-to-left. Only advance to the
1930 next pool when the previous one rolls-over */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001931 for (i=npools-1 ; i >= 0 ; i--) {
1932 pool = PyTuple_GET_ITEM(pools, i);
1933 indices[i]++;
Raymond Hettinger61024b92008-03-02 11:57:16 +00001934 if (indices[i] == PyTuple_GET_SIZE(pool)) {
Raymond Hettinger73d79632008-02-23 02:20:41 +00001935 /* Roll-over and advance to next pool */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001936 indices[i] = 0;
1937 elem = PyTuple_GET_ITEM(pool, 0);
1938 Py_INCREF(elem);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001939 oldelem = PyTuple_GET_ITEM(result, i);
1940 PyTuple_SET_ITEM(result, i, elem);
1941 Py_DECREF(oldelem);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001942 } else {
Raymond Hettinger73d79632008-02-23 02:20:41 +00001943 /* No rollover. Just increment and stop here. */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001944 elem = PyTuple_GET_ITEM(pool, indices[i]);
1945 Py_INCREF(elem);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001946 oldelem = PyTuple_GET_ITEM(result, i);
1947 PyTuple_SET_ITEM(result, i, elem);
1948 Py_DECREF(oldelem);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001949 break;
1950 }
1951 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00001952
1953 /* If i is negative, then the indices have all rolled-over
1954 and we're done. */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001955 if (i < 0)
Raymond Hettinger73d79632008-02-23 02:20:41 +00001956 goto empty;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001957 }
1958
Raymond Hettinger73d79632008-02-23 02:20:41 +00001959 Py_INCREF(result);
1960 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001961
1962empty:
1963 lz->stopped = 1;
1964 return NULL;
1965}
1966
1967PyDoc_STRVAR(product_doc,
1968"product(*iterables) --> product object\n\
1969\n\
Raymond Hettinger73d79632008-02-23 02:20:41 +00001970Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001971For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1972The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1973cycle in a manner similar to an odometer (with the rightmost element changing\n\
1974on every iteration).\n\n\
Raymond Hettinger94034ea2009-02-09 18:39:41 +00001975To compute the product of an iterable with itself, specify the number\n\
1976of repetitions with the optional repeat keyword argument. For example,\n\
1977product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001978product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1979product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1980
1981static PyTypeObject product_type = {
1982 PyVarObject_HEAD_INIT(NULL, 0)
1983 "itertools.product", /* tp_name */
1984 sizeof(productobject), /* tp_basicsize */
1985 0, /* tp_itemsize */
1986 /* methods */
1987 (destructor)product_dealloc, /* tp_dealloc */
1988 0, /* tp_print */
1989 0, /* tp_getattr */
1990 0, /* tp_setattr */
1991 0, /* tp_compare */
1992 0, /* tp_repr */
1993 0, /* tp_as_number */
1994 0, /* tp_as_sequence */
1995 0, /* tp_as_mapping */
1996 0, /* tp_hash */
1997 0, /* tp_call */
1998 0, /* tp_str */
1999 PyObject_GenericGetAttr, /* tp_getattro */
2000 0, /* tp_setattro */
2001 0, /* tp_as_buffer */
2002 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2003 Py_TPFLAGS_BASETYPE, /* tp_flags */
2004 product_doc, /* tp_doc */
2005 (traverseproc)product_traverse, /* tp_traverse */
2006 0, /* tp_clear */
2007 0, /* tp_richcompare */
2008 0, /* tp_weaklistoffset */
2009 PyObject_SelfIter, /* tp_iter */
2010 (iternextfunc)product_next, /* tp_iternext */
2011 0, /* tp_methods */
2012 0, /* tp_members */
2013 0, /* tp_getset */
2014 0, /* tp_base */
2015 0, /* tp_dict */
2016 0, /* tp_descr_get */
2017 0, /* tp_descr_set */
2018 0, /* tp_dictoffset */
2019 0, /* tp_init */
2020 0, /* tp_alloc */
2021 product_new, /* tp_new */
2022 PyObject_GC_Del, /* tp_free */
2023};
2024
2025
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002026/* combinations object ************************************************************/
2027
2028typedef struct {
2029 PyObject_HEAD
2030 PyObject *pool; /* input converted to a tuple */
2031 Py_ssize_t *indices; /* one index per result element */
2032 PyObject *result; /* most recently returned result tuple */
2033 Py_ssize_t r; /* size of result tuple */
2034 int stopped; /* set to 1 when the combinations iterator is exhausted */
2035} combinationsobject;
2036
2037static PyTypeObject combinations_type;
2038
2039static PyObject *
2040combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2041{
2042 combinationsobject *co;
2043 Py_ssize_t n;
2044 Py_ssize_t r;
2045 PyObject *pool = NULL;
2046 PyObject *iterable = NULL;
2047 Py_ssize_t *indices = NULL;
2048 Py_ssize_t i;
2049 static char *kwargs[] = {"iterable", "r", NULL};
2050
2051 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2052 &iterable, &r))
2053 return NULL;
2054
2055 pool = PySequence_Tuple(iterable);
2056 if (pool == NULL)
2057 goto error;
2058 n = PyTuple_GET_SIZE(pool);
2059 if (r < 0) {
2060 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2061 goto error;
2062 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002063
2064 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2065 if (indices == NULL) {
2066 PyErr_NoMemory();
2067 goto error;
2068 }
2069
2070 for (i=0 ; i<r ; i++)
2071 indices[i] = i;
2072
2073 /* create combinationsobject structure */
2074 co = (combinationsobject *)type->tp_alloc(type, 0);
2075 if (co == NULL)
2076 goto error;
2077
2078 co->pool = pool;
2079 co->indices = indices;
2080 co->result = NULL;
2081 co->r = r;
Raymond Hettinger5b913e32009-01-08 06:39:04 +00002082 co->stopped = r > n ? 1 : 0;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002083
2084 return (PyObject *)co;
2085
2086error:
2087 if (indices != NULL)
2088 PyMem_Free(indices);
2089 Py_XDECREF(pool);
2090 return NULL;
2091}
2092
2093static void
2094combinations_dealloc(combinationsobject *co)
2095{
2096 PyObject_GC_UnTrack(co);
2097 Py_XDECREF(co->pool);
2098 Py_XDECREF(co->result);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002099 if (co->indices != NULL)
2100 PyMem_Free(co->indices);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002101 Py_TYPE(co)->tp_free(co);
2102}
2103
2104static int
2105combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2106{
2107 Py_VISIT(co->pool);
2108 Py_VISIT(co->result);
2109 return 0;
2110}
2111
2112static PyObject *
2113combinations_next(combinationsobject *co)
2114{
2115 PyObject *elem;
2116 PyObject *oldelem;
2117 PyObject *pool = co->pool;
2118 Py_ssize_t *indices = co->indices;
2119 PyObject *result = co->result;
2120 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2121 Py_ssize_t r = co->r;
2122 Py_ssize_t i, j, index;
2123
2124 if (co->stopped)
2125 return NULL;
2126
2127 if (result == NULL) {
2128 /* On the first pass, initialize result tuple using the indices */
2129 result = PyTuple_New(r);
2130 if (result == NULL)
2131 goto empty;
2132 co->result = result;
2133 for (i=0; i<r ; i++) {
2134 index = indices[i];
2135 elem = PyTuple_GET_ITEM(pool, index);
2136 Py_INCREF(elem);
2137 PyTuple_SET_ITEM(result, i, elem);
2138 }
2139 } else {
2140 /* Copy the previous result tuple or re-use it if available */
2141 if (Py_REFCNT(result) > 1) {
2142 PyObject *old_result = result;
2143 result = PyTuple_New(r);
2144 if (result == NULL)
2145 goto empty;
2146 co->result = result;
2147 for (i=0; i<r ; i++) {
2148 elem = PyTuple_GET_ITEM(old_result, i);
2149 Py_INCREF(elem);
2150 PyTuple_SET_ITEM(result, i, elem);
2151 }
2152 Py_DECREF(old_result);
2153 }
Christian Heimescdddf182008-02-28 11:18:49 +00002154 /* Now, we've got the only copy so we can update it in-place
2155 * CPython's empty tuple is a singleton and cached in
2156 * PyTuple's freelist.
2157 */
2158 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002159
2160 /* Scan indices right-to-left until finding one that is not
2161 at its maximum (i + n - r). */
2162 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2163 ;
2164
2165 /* If i is negative, then the indices are all at
2166 their maximum value and we're done. */
2167 if (i < 0)
2168 goto empty;
2169
2170 /* Increment the current index which we know is not at its
2171 maximum. Then move back to the right setting each index
2172 to its lowest possible value (one higher than the index
2173 to its left -- this maintains the sort order invariant). */
2174 indices[i]++;
2175 for (j=i+1 ; j<r ; j++)
2176 indices[j] = indices[j-1] + 1;
2177
2178 /* Update the result tuple for the new indices
2179 starting with i, the leftmost index that changed */
2180 for ( ; i<r ; i++) {
2181 index = indices[i];
2182 elem = PyTuple_GET_ITEM(pool, index);
2183 Py_INCREF(elem);
2184 oldelem = PyTuple_GET_ITEM(result, i);
2185 PyTuple_SET_ITEM(result, i, elem);
2186 Py_DECREF(oldelem);
2187 }
2188 }
2189
2190 Py_INCREF(result);
2191 return result;
2192
2193empty:
2194 co->stopped = 1;
2195 return NULL;
2196}
2197
2198PyDoc_STRVAR(combinations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +00002199"combinations(iterable[, r]) --> combinations object\n\
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002200\n\
2201Return successive r-length combinations of elements in the iterable.\n\n\
2202combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2203
2204static PyTypeObject combinations_type = {
2205 PyVarObject_HEAD_INIT(NULL, 0)
2206 "itertools.combinations", /* tp_name */
2207 sizeof(combinationsobject), /* tp_basicsize */
2208 0, /* tp_itemsize */
2209 /* methods */
2210 (destructor)combinations_dealloc, /* tp_dealloc */
2211 0, /* tp_print */
2212 0, /* tp_getattr */
2213 0, /* tp_setattr */
2214 0, /* tp_compare */
2215 0, /* tp_repr */
2216 0, /* tp_as_number */
2217 0, /* tp_as_sequence */
2218 0, /* tp_as_mapping */
2219 0, /* tp_hash */
2220 0, /* tp_call */
2221 0, /* tp_str */
2222 PyObject_GenericGetAttr, /* tp_getattro */
2223 0, /* tp_setattro */
2224 0, /* tp_as_buffer */
2225 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2226 Py_TPFLAGS_BASETYPE, /* tp_flags */
2227 combinations_doc, /* tp_doc */
2228 (traverseproc)combinations_traverse, /* tp_traverse */
2229 0, /* tp_clear */
2230 0, /* tp_richcompare */
2231 0, /* tp_weaklistoffset */
2232 PyObject_SelfIter, /* tp_iter */
2233 (iternextfunc)combinations_next, /* tp_iternext */
2234 0, /* tp_methods */
2235 0, /* tp_members */
2236 0, /* tp_getset */
2237 0, /* tp_base */
2238 0, /* tp_dict */
2239 0, /* tp_descr_get */
2240 0, /* tp_descr_set */
2241 0, /* tp_dictoffset */
2242 0, /* tp_init */
2243 0, /* tp_alloc */
2244 combinations_new, /* tp_new */
2245 PyObject_GC_Del, /* tp_free */
2246};
2247
2248
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002249/* combinations with replacement object *******************************************/
2250
2251/* Equivalent to:
2252
2253 def combinations_with_replacement(iterable, r):
2254 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2255 # number items returned: (n+r-1)! / r! / (n-1)!
2256 pool = tuple(iterable)
2257 n = len(pool)
2258 indices = [0] * r
2259 yield tuple(pool[i] for i in indices)
2260 while 1:
2261 for i in reversed(range(r)):
2262 if indices[i] != n - 1:
2263 break
2264 else:
2265 return
2266 indices[i:] = [indices[i] + 1] * (r - i)
2267 yield tuple(pool[i] for i in indices)
2268
2269 def combinations_with_replacement2(iterable, r):
2270 'Alternate version that filters from product()'
2271 pool = tuple(iterable)
2272 n = len(pool)
2273 for indices in product(range(n), repeat=r):
2274 if sorted(indices) == list(indices):
2275 yield tuple(pool[i] for i in indices)
2276*/
2277typedef struct {
2278 PyObject_HEAD
2279 PyObject *pool; /* input converted to a tuple */
2280 Py_ssize_t *indices; /* one index per result element */
2281 PyObject *result; /* most recently returned result tuple */
2282 Py_ssize_t r; /* size of result tuple */
2283 int stopped; /* set to 1 when the cwr iterator is exhausted */
2284} cwrobject;
2285
2286static PyTypeObject cwr_type;
2287
2288static PyObject *
2289cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2290{
2291 cwrobject *co;
2292 Py_ssize_t n;
2293 Py_ssize_t r;
2294 PyObject *pool = NULL;
2295 PyObject *iterable = NULL;
2296 Py_ssize_t *indices = NULL;
2297 Py_ssize_t i;
2298 static char *kwargs[] = {"iterable", "r", NULL};
2299
2300 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2301 &iterable, &r))
2302 return NULL;
2303
2304 pool = PySequence_Tuple(iterable);
2305 if (pool == NULL)
2306 goto error;
2307 n = PyTuple_GET_SIZE(pool);
2308 if (r < 0) {
2309 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2310 goto error;
2311 }
2312
2313 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2314 if (indices == NULL) {
2315 PyErr_NoMemory();
2316 goto error;
2317 }
2318
2319 for (i=0 ; i<r ; i++)
2320 indices[i] = 0;
2321
2322 /* create cwrobject structure */
2323 co = (cwrobject *)type->tp_alloc(type, 0);
2324 if (co == NULL)
2325 goto error;
2326
2327 co->pool = pool;
2328 co->indices = indices;
2329 co->result = NULL;
2330 co->r = r;
2331 co->stopped = !n && r;
2332
2333 return (PyObject *)co;
2334
2335error:
2336 if (indices != NULL)
2337 PyMem_Free(indices);
2338 Py_XDECREF(pool);
2339 return NULL;
2340}
2341
2342static void
2343cwr_dealloc(cwrobject *co)
2344{
2345 PyObject_GC_UnTrack(co);
2346 Py_XDECREF(co->pool);
2347 Py_XDECREF(co->result);
2348 if (co->indices != NULL)
2349 PyMem_Free(co->indices);
2350 Py_TYPE(co)->tp_free(co);
2351}
2352
2353static int
2354cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2355{
2356 Py_VISIT(co->pool);
2357 Py_VISIT(co->result);
2358 return 0;
2359}
2360
2361static PyObject *
2362cwr_next(cwrobject *co)
2363{
2364 PyObject *elem;
2365 PyObject *oldelem;
2366 PyObject *pool = co->pool;
2367 Py_ssize_t *indices = co->indices;
2368 PyObject *result = co->result;
2369 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2370 Py_ssize_t r = co->r;
2371 Py_ssize_t i, j, index;
2372
2373 if (co->stopped)
2374 return NULL;
2375
2376 if (result == NULL) {
2377 /* On the first pass, initialize result tuple using the indices */
2378 result = PyTuple_New(r);
2379 if (result == NULL)
2380 goto empty;
2381 co->result = result;
2382 for (i=0; i<r ; i++) {
2383 index = indices[i];
2384 elem = PyTuple_GET_ITEM(pool, index);
2385 Py_INCREF(elem);
2386 PyTuple_SET_ITEM(result, i, elem);
2387 }
2388 } else {
2389 /* Copy the previous result tuple or re-use it if available */
2390 if (Py_REFCNT(result) > 1) {
2391 PyObject *old_result = result;
2392 result = PyTuple_New(r);
2393 if (result == NULL)
2394 goto empty;
2395 co->result = result;
2396 for (i=0; i<r ; i++) {
2397 elem = PyTuple_GET_ITEM(old_result, i);
2398 Py_INCREF(elem);
2399 PyTuple_SET_ITEM(result, i, elem);
2400 }
2401 Py_DECREF(old_result);
2402 }
2403 /* Now, we've got the only copy so we can update it in-place CPython's
2404 empty tuple is a singleton and cached in PyTuple's freelist. */
2405 assert(r == 0 || Py_REFCNT(result) == 1);
2406
2407 /* Scan indices right-to-left until finding one that is not
2408 * at its maximum (n-1). */
2409 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2410 ;
2411
2412 /* If i is negative, then the indices are all at
2413 their maximum value and we're done. */
2414 if (i < 0)
2415 goto empty;
2416
2417 /* Increment the current index which we know is not at its
2418 maximum. Then set all to the right to the same value. */
2419 indices[i]++;
2420 for (j=i+1 ; j<r ; j++)
2421 indices[j] = indices[j-1];
2422
2423 /* Update the result tuple for the new indices
2424 starting with i, the leftmost index that changed */
2425 for ( ; i<r ; i++) {
2426 index = indices[i];
2427 elem = PyTuple_GET_ITEM(pool, index);
2428 Py_INCREF(elem);
2429 oldelem = PyTuple_GET_ITEM(result, i);
2430 PyTuple_SET_ITEM(result, i, elem);
2431 Py_DECREF(oldelem);
2432 }
2433 }
2434
2435 Py_INCREF(result);
2436 return result;
2437
2438empty:
2439 co->stopped = 1;
2440 return NULL;
2441}
2442
2443PyDoc_STRVAR(cwr_doc,
2444"combinations_with_replacement(iterable[, r]) --> combinations_with_replacement object\n\
2445\n\
2446Return successive r-length combinations of elements in the iterable\n\
2447allowing individual elements to have successive repeats.\n\
2448combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2449
2450static PyTypeObject cwr_type = {
2451 PyVarObject_HEAD_INIT(NULL, 0)
2452 "itertools.combinations_with_replacement", /* tp_name */
2453 sizeof(cwrobject), /* tp_basicsize */
2454 0, /* tp_itemsize */
2455 /* methods */
2456 (destructor)cwr_dealloc, /* tp_dealloc */
2457 0, /* tp_print */
2458 0, /* tp_getattr */
2459 0, /* tp_setattr */
2460 0, /* tp_compare */
2461 0, /* tp_repr */
2462 0, /* tp_as_number */
2463 0, /* tp_as_sequence */
2464 0, /* tp_as_mapping */
2465 0, /* tp_hash */
2466 0, /* tp_call */
2467 0, /* tp_str */
2468 PyObject_GenericGetAttr, /* tp_getattro */
2469 0, /* tp_setattro */
2470 0, /* tp_as_buffer */
2471 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2472 Py_TPFLAGS_BASETYPE, /* tp_flags */
2473 cwr_doc, /* tp_doc */
2474 (traverseproc)cwr_traverse, /* tp_traverse */
2475 0, /* tp_clear */
2476 0, /* tp_richcompare */
2477 0, /* tp_weaklistoffset */
2478 PyObject_SelfIter, /* tp_iter */
2479 (iternextfunc)cwr_next, /* tp_iternext */
2480 0, /* tp_methods */
2481 0, /* tp_members */
2482 0, /* tp_getset */
2483 0, /* tp_base */
2484 0, /* tp_dict */
2485 0, /* tp_descr_get */
2486 0, /* tp_descr_set */
2487 0, /* tp_dictoffset */
2488 0, /* tp_init */
2489 0, /* tp_alloc */
2490 cwr_new, /* tp_new */
2491 PyObject_GC_Del, /* tp_free */
2492};
2493
2494
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002495/* permutations object ************************************************************
2496
2497def permutations(iterable, r=None):
2498 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2499 pool = tuple(iterable)
2500 n = len(pool)
2501 r = n if r is None else r
2502 indices = range(n)
2503 cycles = range(n-r+1, n+1)[::-1]
2504 yield tuple(pool[i] for i in indices[:r])
2505 while n:
2506 for i in reversed(range(r)):
2507 cycles[i] -= 1
2508 if cycles[i] == 0:
2509 indices[i:] = indices[i+1:] + indices[i:i+1]
2510 cycles[i] = n - i
2511 else:
2512 j = cycles[i]
2513 indices[i], indices[-j] = indices[-j], indices[i]
2514 yield tuple(pool[i] for i in indices[:r])
2515 break
2516 else:
2517 return
2518*/
2519
2520typedef struct {
2521 PyObject_HEAD
2522 PyObject *pool; /* input converted to a tuple */
2523 Py_ssize_t *indices; /* one index per element in the pool */
2524 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2525 PyObject *result; /* most recently returned result tuple */
2526 Py_ssize_t r; /* size of result tuple */
2527 int stopped; /* set to 1 when the permutations iterator is exhausted */
2528} permutationsobject;
2529
2530static PyTypeObject permutations_type;
2531
2532static PyObject *
2533permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2534{
2535 permutationsobject *po;
2536 Py_ssize_t n;
2537 Py_ssize_t r;
2538 PyObject *robj = Py_None;
2539 PyObject *pool = NULL;
2540 PyObject *iterable = NULL;
2541 Py_ssize_t *indices = NULL;
2542 Py_ssize_t *cycles = NULL;
2543 Py_ssize_t i;
2544 static char *kwargs[] = {"iterable", "r", NULL};
2545
2546 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2547 &iterable, &robj))
2548 return NULL;
2549
2550 pool = PySequence_Tuple(iterable);
2551 if (pool == NULL)
2552 goto error;
2553 n = PyTuple_GET_SIZE(pool);
2554
2555 r = n;
2556 if (robj != Py_None) {
2557 r = PyInt_AsSsize_t(robj);
2558 if (r == -1 && PyErr_Occurred())
2559 goto error;
2560 }
2561 if (r < 0) {
2562 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2563 goto error;
2564 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002565
2566 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2567 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2568 if (indices == NULL || cycles == NULL) {
2569 PyErr_NoMemory();
2570 goto error;
2571 }
2572
2573 for (i=0 ; i<n ; i++)
2574 indices[i] = i;
2575 for (i=0 ; i<r ; i++)
2576 cycles[i] = n - i;
2577
2578 /* create permutationsobject structure */
2579 po = (permutationsobject *)type->tp_alloc(type, 0);
2580 if (po == NULL)
2581 goto error;
2582
2583 po->pool = pool;
2584 po->indices = indices;
2585 po->cycles = cycles;
2586 po->result = NULL;
2587 po->r = r;
Raymond Hettinger5b913e32009-01-08 06:39:04 +00002588 po->stopped = r > n ? 1 : 0;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002589
2590 return (PyObject *)po;
2591
2592error:
2593 if (indices != NULL)
2594 PyMem_Free(indices);
2595 if (cycles != NULL)
2596 PyMem_Free(cycles);
2597 Py_XDECREF(pool);
2598 return NULL;
2599}
2600
2601static void
2602permutations_dealloc(permutationsobject *po)
2603{
2604 PyObject_GC_UnTrack(po);
2605 Py_XDECREF(po->pool);
2606 Py_XDECREF(po->result);
2607 PyMem_Free(po->indices);
2608 PyMem_Free(po->cycles);
2609 Py_TYPE(po)->tp_free(po);
2610}
2611
2612static int
2613permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2614{
Raymond Hettinger6e3e4152008-03-05 21:04:32 +00002615 Py_VISIT(po->pool);
2616 Py_VISIT(po->result);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002617 return 0;
2618}
2619
2620static PyObject *
2621permutations_next(permutationsobject *po)
2622{
2623 PyObject *elem;
2624 PyObject *oldelem;
2625 PyObject *pool = po->pool;
2626 Py_ssize_t *indices = po->indices;
2627 Py_ssize_t *cycles = po->cycles;
2628 PyObject *result = po->result;
2629 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2630 Py_ssize_t r = po->r;
2631 Py_ssize_t i, j, k, index;
2632
2633 if (po->stopped)
2634 return NULL;
2635
2636 if (result == NULL) {
2637 /* On the first pass, initialize result tuple using the indices */
2638 result = PyTuple_New(r);
2639 if (result == NULL)
2640 goto empty;
2641 po->result = result;
2642 for (i=0; i<r ; i++) {
2643 index = indices[i];
2644 elem = PyTuple_GET_ITEM(pool, index);
2645 Py_INCREF(elem);
2646 PyTuple_SET_ITEM(result, i, elem);
2647 }
2648 } else {
2649 if (n == 0)
2650 goto empty;
2651
2652 /* Copy the previous result tuple or re-use it if available */
2653 if (Py_REFCNT(result) > 1) {
2654 PyObject *old_result = result;
2655 result = PyTuple_New(r);
2656 if (result == NULL)
2657 goto empty;
2658 po->result = result;
2659 for (i=0; i<r ; i++) {
2660 elem = PyTuple_GET_ITEM(old_result, i);
2661 Py_INCREF(elem);
2662 PyTuple_SET_ITEM(result, i, elem);
2663 }
2664 Py_DECREF(old_result);
2665 }
2666 /* Now, we've got the only copy so we can update it in-place */
2667 assert(r == 0 || Py_REFCNT(result) == 1);
2668
2669 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2670 for (i=r-1 ; i>=0 ; i--) {
2671 cycles[i] -= 1;
2672 if (cycles[i] == 0) {
2673 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2674 index = indices[i];
2675 for (j=i ; j<n-1 ; j++)
2676 indices[j] = indices[j+1];
2677 indices[n-1] = index;
2678 cycles[i] = n - i;
2679 } else {
2680 j = cycles[i];
2681 index = indices[i];
2682 indices[i] = indices[n-j];
2683 indices[n-j] = index;
2684
2685 for (k=i; k<r ; k++) {
2686 /* start with i, the leftmost element that changed */
2687 /* yield tuple(pool[k] for k in indices[:r]) */
2688 index = indices[k];
2689 elem = PyTuple_GET_ITEM(pool, index);
2690 Py_INCREF(elem);
2691 oldelem = PyTuple_GET_ITEM(result, k);
2692 PyTuple_SET_ITEM(result, k, elem);
2693 Py_DECREF(oldelem);
2694 }
2695 break;
2696 }
2697 }
2698 /* If i is negative, then the cycles have all
2699 rolled-over and we're done. */
2700 if (i < 0)
2701 goto empty;
2702 }
2703 Py_INCREF(result);
2704 return result;
2705
2706empty:
2707 po->stopped = 1;
2708 return NULL;
2709}
2710
2711PyDoc_STRVAR(permutations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +00002712"permutations(iterable[, r]) --> permutations object\n\
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002713\n\
2714Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl8550d692008-06-10 12:46:39 +00002715permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002716
2717static PyTypeObject permutations_type = {
2718 PyVarObject_HEAD_INIT(NULL, 0)
2719 "itertools.permutations", /* tp_name */
2720 sizeof(permutationsobject), /* tp_basicsize */
2721 0, /* tp_itemsize */
2722 /* methods */
2723 (destructor)permutations_dealloc, /* tp_dealloc */
2724 0, /* tp_print */
2725 0, /* tp_getattr */
2726 0, /* tp_setattr */
2727 0, /* tp_compare */
2728 0, /* tp_repr */
2729 0, /* tp_as_number */
2730 0, /* tp_as_sequence */
2731 0, /* tp_as_mapping */
2732 0, /* tp_hash */
2733 0, /* tp_call */
2734 0, /* tp_str */
2735 PyObject_GenericGetAttr, /* tp_getattro */
2736 0, /* tp_setattro */
2737 0, /* tp_as_buffer */
2738 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2739 Py_TPFLAGS_BASETYPE, /* tp_flags */
2740 permutations_doc, /* tp_doc */
2741 (traverseproc)permutations_traverse, /* tp_traverse */
2742 0, /* tp_clear */
2743 0, /* tp_richcompare */
2744 0, /* tp_weaklistoffset */
2745 PyObject_SelfIter, /* tp_iter */
2746 (iternextfunc)permutations_next, /* tp_iternext */
2747 0, /* tp_methods */
2748 0, /* tp_members */
2749 0, /* tp_getset */
2750 0, /* tp_base */
2751 0, /* tp_dict */
2752 0, /* tp_descr_get */
2753 0, /* tp_descr_set */
2754 0, /* tp_dictoffset */
2755 0, /* tp_init */
2756 0, /* tp_alloc */
2757 permutations_new, /* tp_new */
2758 PyObject_GC_Del, /* tp_free */
2759};
2760
2761
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002762/* compress object ************************************************************/
2763
2764/* Equivalent to:
2765
2766 def compress(data, selectors):
2767 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2768 return (d for d, s in izip(data, selectors) if s)
2769*/
2770
2771typedef struct {
2772 PyObject_HEAD
2773 PyObject *data;
2774 PyObject *selectors;
2775} compressobject;
2776
2777static PyTypeObject compress_type;
2778
2779static PyObject *
2780compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2781{
2782 PyObject *seq1, *seq2;
2783 PyObject *data=NULL, *selectors=NULL;
2784 compressobject *lz;
Raymond Hettinger2e2909f2009-02-19 02:15:14 +00002785 static char *kwargs[] = {"data", "selectors", NULL};
2786
2787 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002788 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,
Raymond Hettinger2e2909f2009-02-19 02:15:14 +00002865"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002866\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 Hettinger31c769c2009-02-12 05:39:46 +00003206 PyObject *long_cnt;
3207 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003208} countobject;
3209
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003210/* Counting logic and invariants:
3211
3212C_add_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
3213
3214 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3215 Advances with: cnt += 1
3216 When count hits Y_SSIZE_T_MAX, switch to Py_add_mode.
3217
3218Py_add_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
3219
3220 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3221 All counting is done with python objects (no overflows or underflows).
3222 Advances with: long_cnt += long_step
3223 Step may be zero -- effectively a slow version of repeat(cnt).
3224 Either long_cnt or long_step may be a float.
3225*/
3226
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003227static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003228
3229static PyObject *
3230count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3231{
3232 countobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00003233 Py_ssize_t cnt = 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003234 PyObject *long_cnt = NULL;
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003235 PyObject *long_step = NULL;
Raymond Hettingera4038032009-02-14 00:25:51 +00003236 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003237
Raymond Hettingera4038032009-02-14 00:25:51 +00003238 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3239 kwlist, &long_cnt, &long_step))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003240 return NULL;
3241
Benjamin Peterson062a7c32009-02-16 21:07:52 +00003242 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3243 (long_step != NULL && !PyNumber_Check(long_step))) {
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003244 PyErr_SetString(PyExc_TypeError, "a number is required");
3245 return NULL;
3246 }
3247
3248 if (long_step == NULL) {
3249 /* If not specified, step defaults to 1 */
3250 long_step = PyInt_FromLong(1);
3251 if (long_step == NULL)
3252 return NULL;
3253 } else
3254 Py_INCREF(long_step);
3255 assert(long_step != NULL);
3256
3257 if (long_cnt != NULL) {
3258 cnt = PyInt_AsSsize_t(long_cnt);
3259 if ((cnt == -1 && PyErr_Occurred()) ||
3260 !PyIndex_Check(long_cnt) ||
3261 !PyInt_Check(long_step) ||
3262 PyInt_AS_LONG(long_step) != 1) {
3263 /* Switch to Py_add_mode */
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003264 PyErr_Clear();
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003265 Py_INCREF(long_cnt);
3266 cnt = PY_SSIZE_T_MAX;
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003267 } else
3268 long_cnt = NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003269 }
Benjamin Peterson062a7c32009-02-16 21:07:52 +00003270 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL) ||
3271 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL));
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003272
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003273 /* create countobject structure */
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003274 lz = (countobject *)type->tp_alloc(type, 0);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003275 if (lz == NULL) {
3276 Py_XDECREF(long_cnt);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003277 return NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003278 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003279 lz->cnt = cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003280 lz->long_cnt = long_cnt;
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003281 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003282
3283 return (PyObject *)lz;
3284}
3285
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003286static void
3287count_dealloc(countobject *lz)
3288{
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003289 PyObject_GC_UnTrack(lz);
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003290 Py_XDECREF(lz->long_cnt);
3291 Py_XDECREF(lz->long_step);
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003292 Py_TYPE(lz)->tp_free(lz);
3293}
3294
Benjamin Peterson062a7c32009-02-16 21:07:52 +00003295static int
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003296count_traverse(countobject *lz, visitproc visit, void *arg)
3297{
3298 Py_VISIT(lz->long_cnt);
3299 Py_VISIT(lz->long_step);
3300 return 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003301}
3302
3303static PyObject *
3304count_nextlong(countobject *lz)
3305{
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003306 PyObject *long_cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003307 PyObject *stepped_up;
3308
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003309 long_cnt = lz->long_cnt;
3310 if (long_cnt == NULL) {
3311 /* Switch to Py_add_mode */
3312 long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
3313 if (long_cnt == NULL)
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003314 return NULL;
3315 }
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003316 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
3317
3318 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003319 if (stepped_up == NULL)
3320 return NULL;
3321 lz->long_cnt = stepped_up;
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003322 return long_cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003323}
3324
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003325static PyObject *
3326count_next(countobject *lz)
3327{
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003328 if (lz->cnt == PY_SSIZE_T_MAX)
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003329 return count_nextlong(lz);
Jack Diederich36234e82006-09-21 17:50:26 +00003330 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003331}
3332
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003333static PyObject *
3334count_repr(countobject *lz)
3335{
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003336 PyObject *cnt_repr, *step_repr = NULL;
3337 PyObject *result = NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003338
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003339 if (lz->cnt != PY_SSIZE_T_MAX)
Gregory P. Smithdd96db62008-06-09 04:58:54 +00003340 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003341
3342 cnt_repr = PyObject_Repr(lz->long_cnt);
3343 if (cnt_repr == NULL)
3344 return NULL;
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003345
3346 if (PyInt_Check(lz->long_step) && PyInt_AS_LONG(lz->long_step) == 1) {
3347 /* Don't display step when it is an integer equal to 1 */
3348 result = PyString_FromFormat("count(%s)",
3349 PyString_AS_STRING(cnt_repr));
3350 } else {
3351 step_repr = PyObject_Repr(lz->long_step);
3352 if (step_repr != NULL)
3353 result = PyString_FromFormat("count(%s, %s)",
3354 PyString_AS_STRING(cnt_repr),
3355 PyString_AS_STRING(step_repr));
3356 }
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003357 Py_DECREF(cnt_repr);
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003358 Py_XDECREF(step_repr);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003359 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003360}
3361
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003362PyDoc_STRVAR(count_doc,
Raymond Hettingera4038032009-02-14 00:25:51 +00003363 "count([start[, step]]) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003364\n\
3365Return a count object whose .next() method returns consecutive\n\
Raymond Hettingera4038032009-02-14 00:25:51 +00003366integers starting from zero or, if specified, from start.\n\
Raymond Hettinger66373ab2009-02-12 10:16:19 +00003367If step is specified, counts by that interval. Equivalent to:\n\n\
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003368 def count(firstval=0, step=1):\n\
3369 x = firstval\n\
Raymond Hettinger66373ab2009-02-12 10:16:19 +00003370 while 1:\n\
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003371 yield x\n\
Raymond Hettinger66373ab2009-02-12 10:16:19 +00003372 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003373
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003374static PyTypeObject count_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003375 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003376 "itertools.count", /* tp_name */
3377 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003378 0, /* tp_itemsize */
3379 /* methods */
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003380 (destructor)count_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003381 0, /* tp_print */
3382 0, /* tp_getattr */
3383 0, /* tp_setattr */
3384 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003385 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003386 0, /* tp_as_number */
3387 0, /* tp_as_sequence */
3388 0, /* tp_as_mapping */
3389 0, /* tp_hash */
3390 0, /* tp_call */
3391 0, /* tp_str */
3392 PyObject_GenericGetAttr, /* tp_getattro */
3393 0, /* tp_setattro */
3394 0, /* tp_as_buffer */
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003395 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3396 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003397 count_doc, /* tp_doc */
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003398 (traverseproc)count_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003399 0, /* tp_clear */
3400 0, /* tp_richcompare */
3401 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003402 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003403 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003404 0, /* tp_methods */
3405 0, /* tp_members */
3406 0, /* tp_getset */
3407 0, /* tp_base */
3408 0, /* tp_dict */
3409 0, /* tp_descr_get */
3410 0, /* tp_descr_set */
3411 0, /* tp_dictoffset */
3412 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003413 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003414 count_new, /* tp_new */
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003415 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003416};
3417
3418
3419/* izip object ************************************************************/
3420
3421#include "Python.h"
3422
3423typedef struct {
3424 PyObject_HEAD
Martin v. Löwisad0a4622006-02-16 14:30:23 +00003425 Py_ssize_t tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003426 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00003427 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003428} izipobject;
3429
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003430static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003431
3432static PyObject *
3433izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3434{
3435 izipobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00003436 Py_ssize_t i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003437 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00003438 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00003439 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003440
Georg Brandlb84c1372007-01-21 10:28:43 +00003441 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00003442 return NULL;
3443
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003444 /* args must be a tuple */
3445 assert(PyTuple_Check(args));
3446
3447 /* obtain iterators */
3448 ittuple = PyTuple_New(tuplesize);
Neal Norwitz2f3136b2006-05-27 05:18:57 +00003449 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003450 return NULL;
3451 for (i=0; i < tuplesize; ++i) {
3452 PyObject *item = PyTuple_GET_ITEM(args, i);
3453 PyObject *it = PyObject_GetIter(item);
3454 if (it == NULL) {
3455 if (PyErr_ExceptionMatches(PyExc_TypeError))
3456 PyErr_Format(PyExc_TypeError,
Neal Norwitz2f3136b2006-05-27 05:18:57 +00003457 "izip argument #%zd must support iteration",
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003458 i+1);
3459 Py_DECREF(ittuple);
3460 return NULL;
3461 }
3462 PyTuple_SET_ITEM(ittuple, i, it);
3463 }
3464
Raymond Hettinger2012f172003-02-07 05:32:58 +00003465 /* create a result holder */
3466 result = PyTuple_New(tuplesize);
3467 if (result == NULL) {
3468 Py_DECREF(ittuple);
3469 return NULL;
3470 }
3471 for (i=0 ; i < tuplesize ; i++) {
3472 Py_INCREF(Py_None);
3473 PyTuple_SET_ITEM(result, i, Py_None);
3474 }
3475
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003476 /* create izipobject structure */
3477 lz = (izipobject *)type->tp_alloc(type, 0);
3478 if (lz == NULL) {
3479 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003480 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003481 return NULL;
3482 }
3483 lz->ittuple = ittuple;
3484 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00003485 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003486
3487 return (PyObject *)lz;
3488}
3489
3490static void
3491izip_dealloc(izipobject *lz)
3492{
3493 PyObject_GC_UnTrack(lz);
3494 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003495 Py_XDECREF(lz->result);
Christian Heimese93237d2007-12-19 02:37:44 +00003496 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003497}
3498
3499static int
3500izip_traverse(izipobject *lz, visitproc visit, void *arg)
3501{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00003502 Py_VISIT(lz->ittuple);
3503 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003504 return 0;
3505}
3506
3507static PyObject *
3508izip_next(izipobject *lz)
3509{
Jack Diederich6c433a92006-05-26 11:15:17 +00003510 Py_ssize_t i;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00003511 Py_ssize_t tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00003512 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003513 PyObject *it;
3514 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00003515 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003516
Raymond Hettingerb5a42082003-08-08 05:10:41 +00003517 if (tuplesize == 0)
3518 return NULL;
Christian Heimese93237d2007-12-19 02:37:44 +00003519 if (Py_REFCNT(result) == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00003520 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003521 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003522 it = PyTuple_GET_ITEM(lz->ittuple, i);
Christian Heimese93237d2007-12-19 02:37:44 +00003523 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00003524 if (item == NULL) {
3525 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003526 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00003527 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00003528 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003529 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00003530 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003531 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003532 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00003533 result = PyTuple_New(tuplesize);
3534 if (result == NULL)
3535 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003536 for (i=0 ; i < tuplesize ; i++) {
3537 it = PyTuple_GET_ITEM(lz->ittuple, i);
Christian Heimese93237d2007-12-19 02:37:44 +00003538 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003539 if (item == NULL) {
3540 Py_DECREF(result);
3541 return NULL;
3542 }
3543 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003544 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003545 }
3546 return result;
3547}
3548
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003549PyDoc_STRVAR(izip_doc,
3550"izip(iter1 [,iter2 [...]]) --> izip object\n\
3551\n\
3552Return a izip object whose .next() method returns a tuple where\n\
3553the i-th element comes from the i-th iterable argument. The .next()\n\
3554method continues until the shortest iterable in the argument sequence\n\
3555is exhausted and then it raises StopIteration. Works like the zip()\n\
3556function but consumes less memory by returning an iterator instead of\n\
3557a list.");
3558
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003559static PyTypeObject izip_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003560 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003561 "itertools.izip", /* tp_name */
3562 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003563 0, /* tp_itemsize */
3564 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003565 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003566 0, /* tp_print */
3567 0, /* tp_getattr */
3568 0, /* tp_setattr */
3569 0, /* tp_compare */
3570 0, /* tp_repr */
3571 0, /* tp_as_number */
3572 0, /* tp_as_sequence */
3573 0, /* tp_as_mapping */
3574 0, /* tp_hash */
3575 0, /* tp_call */
3576 0, /* tp_str */
3577 PyObject_GenericGetAttr, /* tp_getattro */
3578 0, /* tp_setattro */
3579 0, /* tp_as_buffer */
3580 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3581 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003582 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003583 (traverseproc)izip_traverse, /* tp_traverse */
3584 0, /* tp_clear */
3585 0, /* tp_richcompare */
3586 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003587 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003588 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003589 0, /* tp_methods */
3590 0, /* tp_members */
3591 0, /* tp_getset */
3592 0, /* tp_base */
3593 0, /* tp_dict */
3594 0, /* tp_descr_get */
3595 0, /* tp_descr_set */
3596 0, /* tp_dictoffset */
3597 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003598 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003599 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003600 PyObject_GC_Del, /* tp_free */
3601};
3602
3603
3604/* repeat object ************************************************************/
3605
3606typedef struct {
3607 PyObject_HEAD
3608 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00003609 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003610} repeatobject;
3611
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003612static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003613
3614static PyObject *
3615repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3616{
3617 repeatobject *ro;
3618 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00003619 Py_ssize_t cnt = -1;
Raymond Hettinger182edae2009-02-19 02:38:25 +00003620 static char *kwargs[] = {"object", "times", NULL};
3621
3622 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3623 &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003624 return NULL;
3625
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003626 if (PyTuple_Size(args) == 2 && cnt < 0)
3627 cnt = 0;
3628
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003629 ro = (repeatobject *)type->tp_alloc(type, 0);
3630 if (ro == NULL)
3631 return NULL;
3632 Py_INCREF(element);
3633 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003634 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003635 return (PyObject *)ro;
3636}
3637
3638static void
3639repeat_dealloc(repeatobject *ro)
3640{
3641 PyObject_GC_UnTrack(ro);
3642 Py_XDECREF(ro->element);
Christian Heimese93237d2007-12-19 02:37:44 +00003643 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003644}
3645
3646static int
3647repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3648{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00003649 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003650 return 0;
3651}
3652
3653static PyObject *
3654repeat_next(repeatobject *ro)
3655{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003656 if (ro->cnt == 0)
3657 return NULL;
3658 if (ro->cnt > 0)
3659 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003660 Py_INCREF(ro->element);
3661 return ro->element;
3662}
3663
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003664static PyObject *
3665repeat_repr(repeatobject *ro)
3666{
3667 PyObject *result, *objrepr;
3668
3669 objrepr = PyObject_Repr(ro->element);
3670 if (objrepr == NULL)
3671 return NULL;
3672
3673 if (ro->cnt == -1)
Gregory P. Smithdd96db62008-06-09 04:58:54 +00003674 result = PyString_FromFormat("repeat(%s)",
3675 PyString_AS_STRING(objrepr));
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003676 else
Gregory P. Smithdd96db62008-06-09 04:58:54 +00003677 result = PyString_FromFormat("repeat(%s, %zd)",
3678 PyString_AS_STRING(objrepr), ro->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003679 Py_DECREF(objrepr);
3680 return result;
3681}
3682
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003683static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003684repeat_len(repeatobject *ro)
3685{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003686 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003687 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003688 return NULL;
3689 }
Jack Diederich6c433a92006-05-26 11:15:17 +00003690 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003691}
3692
Armin Rigof5b3e362006-02-11 21:32:43 +00003693PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003694
3695static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00003696 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003697 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003698};
3699
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003700PyDoc_STRVAR(repeat_doc,
Raymond Hettinger182edae2009-02-19 02:38:25 +00003701"repeat(object [,times]) -> create an iterator which returns the object\n\
3702for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003703endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003704
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003705static PyTypeObject repeat_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003706 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003707 "itertools.repeat", /* tp_name */
3708 sizeof(repeatobject), /* tp_basicsize */
3709 0, /* tp_itemsize */
3710 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003711 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003712 0, /* tp_print */
3713 0, /* tp_getattr */
3714 0, /* tp_setattr */
3715 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003716 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003717 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003718 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003719 0, /* tp_as_mapping */
3720 0, /* tp_hash */
3721 0, /* tp_call */
3722 0, /* tp_str */
3723 PyObject_GenericGetAttr, /* tp_getattro */
3724 0, /* tp_setattro */
3725 0, /* tp_as_buffer */
3726 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3727 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003728 repeat_doc, /* tp_doc */
3729 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003730 0, /* tp_clear */
3731 0, /* tp_richcompare */
3732 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003733 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003734 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003735 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003736 0, /* tp_members */
3737 0, /* tp_getset */
3738 0, /* tp_base */
3739 0, /* tp_dict */
3740 0, /* tp_descr_get */
3741 0, /* tp_descr_set */
3742 0, /* tp_dictoffset */
3743 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003744 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003745 repeat_new, /* tp_new */
3746 PyObject_GC_Del, /* tp_free */
3747};
3748
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003749/* iziplongest object ************************************************************/
3750
3751#include "Python.h"
3752
3753typedef struct {
3754 PyObject_HEAD
3755 Py_ssize_t tuplesize;
3756 Py_ssize_t numactive;
3757 PyObject *ittuple; /* tuple of iterators */
3758 PyObject *result;
3759 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003760} iziplongestobject;
3761
3762static PyTypeObject iziplongest_type;
3763
3764static PyObject *
3765izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3766{
3767 iziplongestobject *lz;
3768 Py_ssize_t i;
3769 PyObject *ittuple; /* tuple of iterators */
3770 PyObject *result;
3771 PyObject *fillvalue = Py_None;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003772 Py_ssize_t tuplesize = PySequence_Length(args);
3773
3774 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3775 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3776 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3777 PyErr_SetString(PyExc_TypeError,
3778 "izip_longest() got an unexpected keyword argument");
3779 return NULL;
3780 }
3781 }
3782
3783 /* args must be a tuple */
3784 assert(PyTuple_Check(args));
3785
3786 /* obtain iterators */
3787 ittuple = PyTuple_New(tuplesize);
3788 if (ittuple == NULL)
3789 return NULL;
3790 for (i=0; i < tuplesize; ++i) {
3791 PyObject *item = PyTuple_GET_ITEM(args, i);
3792 PyObject *it = PyObject_GetIter(item);
3793 if (it == NULL) {
3794 if (PyErr_ExceptionMatches(PyExc_TypeError))
3795 PyErr_Format(PyExc_TypeError,
3796 "izip_longest argument #%zd must support iteration",
3797 i+1);
3798 Py_DECREF(ittuple);
3799 return NULL;
3800 }
3801 PyTuple_SET_ITEM(ittuple, i, it);
3802 }
3803
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003804 /* create a result holder */
3805 result = PyTuple_New(tuplesize);
3806 if (result == NULL) {
3807 Py_DECREF(ittuple);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003808 return NULL;
3809 }
3810 for (i=0 ; i < tuplesize ; i++) {
3811 Py_INCREF(Py_None);
3812 PyTuple_SET_ITEM(result, i, Py_None);
3813 }
3814
3815 /* create iziplongestobject structure */
3816 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3817 if (lz == NULL) {
3818 Py_DECREF(ittuple);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003819 Py_DECREF(result);
3820 return NULL;
3821 }
3822 lz->ittuple = ittuple;
3823 lz->tuplesize = tuplesize;
3824 lz->numactive = tuplesize;
3825 lz->result = result;
3826 Py_INCREF(fillvalue);
3827 lz->fillvalue = fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003828 return (PyObject *)lz;
3829}
3830
3831static void
3832izip_longest_dealloc(iziplongestobject *lz)
3833{
3834 PyObject_GC_UnTrack(lz);
3835 Py_XDECREF(lz->ittuple);
3836 Py_XDECREF(lz->result);
3837 Py_XDECREF(lz->fillvalue);
Christian Heimese93237d2007-12-19 02:37:44 +00003838 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003839}
3840
3841static int
3842izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3843{
3844 Py_VISIT(lz->ittuple);
3845 Py_VISIT(lz->result);
3846 Py_VISIT(lz->fillvalue);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003847 return 0;
3848}
3849
3850static PyObject *
3851izip_longest_next(iziplongestobject *lz)
3852{
3853 Py_ssize_t i;
3854 Py_ssize_t tuplesize = lz->tuplesize;
3855 PyObject *result = lz->result;
3856 PyObject *it;
3857 PyObject *item;
3858 PyObject *olditem;
3859
3860 if (tuplesize == 0)
3861 return NULL;
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003862 if (lz->numactive == 0)
3863 return NULL;
Christian Heimese93237d2007-12-19 02:37:44 +00003864 if (Py_REFCNT(result) == 1) {
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003865 Py_INCREF(result);
3866 for (i=0 ; i < tuplesize ; i++) {
3867 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003868 if (it == NULL) {
3869 Py_INCREF(lz->fillvalue);
3870 item = lz->fillvalue;
3871 } else {
Christian Heimese93237d2007-12-19 02:37:44 +00003872 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003873 if (item == NULL) {
3874 lz->numactive -= 1;
3875 if (lz->numactive == 0) {
3876 Py_DECREF(result);
3877 return NULL;
3878 } else {
3879 Py_INCREF(lz->fillvalue);
3880 item = lz->fillvalue;
3881 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3882 Py_DECREF(it);
3883 }
3884 }
3885 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003886 olditem = PyTuple_GET_ITEM(result, i);
3887 PyTuple_SET_ITEM(result, i, item);
3888 Py_DECREF(olditem);
3889 }
3890 } else {
3891 result = PyTuple_New(tuplesize);
3892 if (result == NULL)
3893 return NULL;
3894 for (i=0 ; i < tuplesize ; i++) {
3895 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003896 if (it == NULL) {
3897 Py_INCREF(lz->fillvalue);
3898 item = lz->fillvalue;
3899 } else {
Christian Heimese93237d2007-12-19 02:37:44 +00003900 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003901 if (item == NULL) {
3902 lz->numactive -= 1;
3903 if (lz->numactive == 0) {
3904 Py_DECREF(result);
3905 return NULL;
3906 } else {
3907 Py_INCREF(lz->fillvalue);
3908 item = lz->fillvalue;
3909 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3910 Py_DECREF(it);
3911 }
3912 }
3913 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003914 PyTuple_SET_ITEM(result, i, item);
3915 }
3916 }
3917 return result;
3918}
3919
3920PyDoc_STRVAR(izip_longest_doc,
3921"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3922\n\
3923Return an izip_longest object whose .next() method returns a tuple where\n\
3924the i-th element comes from the i-th iterable argument. The .next()\n\
3925method continues until the longest iterable in the argument sequence\n\
3926is exhausted and then it raises StopIteration. When the shorter iterables\n\
3927are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3928defaults to None or can be specified by a keyword argument.\n\
3929");
3930
3931static PyTypeObject iziplongest_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003932 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003933 "itertools.izip_longest", /* tp_name */
3934 sizeof(iziplongestobject), /* tp_basicsize */
3935 0, /* tp_itemsize */
3936 /* methods */
3937 (destructor)izip_longest_dealloc, /* tp_dealloc */
3938 0, /* tp_print */
3939 0, /* tp_getattr */
3940 0, /* tp_setattr */
3941 0, /* tp_compare */
3942 0, /* tp_repr */
3943 0, /* tp_as_number */
3944 0, /* tp_as_sequence */
3945 0, /* tp_as_mapping */
3946 0, /* tp_hash */
3947 0, /* tp_call */
3948 0, /* tp_str */
3949 PyObject_GenericGetAttr, /* tp_getattro */
3950 0, /* tp_setattro */
3951 0, /* tp_as_buffer */
3952 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3953 Py_TPFLAGS_BASETYPE, /* tp_flags */
3954 izip_longest_doc, /* tp_doc */
3955 (traverseproc)izip_longest_traverse, /* tp_traverse */
3956 0, /* tp_clear */
3957 0, /* tp_richcompare */
3958 0, /* tp_weaklistoffset */
3959 PyObject_SelfIter, /* tp_iter */
3960 (iternextfunc)izip_longest_next, /* tp_iternext */
3961 0, /* tp_methods */
3962 0, /* tp_members */
3963 0, /* tp_getset */
3964 0, /* tp_base */
3965 0, /* tp_dict */
3966 0, /* tp_descr_get */
3967 0, /* tp_descr_set */
3968 0, /* tp_dictoffset */
3969 0, /* tp_init */
3970 0, /* tp_alloc */
3971 izip_longest_new, /* tp_new */
3972 PyObject_GC_Del, /* tp_free */
3973};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003974
3975/* module level code ********************************************************/
3976
3977PyDoc_STRVAR(module_doc,
3978"Functional tools for creating and using iterators.\n\
3979\n\
3980Infinite iterators:\n\
3981count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003982cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003983repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003984\n\
3985Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00003986chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettingere76816b2009-01-29 03:43:44 +00003987compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00003988dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
3989groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003990ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
3991ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003992islice(seq, [start,] stop [, step]) --> elements from\n\
3993 seq[start:stop:step]\n\
3994imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
3995starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003996tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003997takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00003998izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
3999izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4000\n\
4001Combinatoric generators:\n\
4002product(p, q, ... [repeat=1]) --> cartesian product\n\
4003permutations(p[, r])\n\
4004combinations(p[, r])\n\
4005combinations_with_replacement(p[, r])\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004006");
4007
4008
Raymond Hettingerad983e72003-11-12 14:32:26 +00004009static PyMethodDef module_methods[] = {
4010 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4011 {NULL, NULL} /* sentinel */
4012};
4013
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004014PyMODINIT_FUNC
4015inititertools(void)
4016{
Raymond Hettinger60eca932003-02-09 06:40:58 +00004017 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004018 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00004019 char *name;
4020 PyTypeObject *typelist[] = {
Raymond Hettinger93e804d2008-02-26 23:40:50 +00004021 &combinations_type,
Raymond Hettingerd081abc2009-01-27 02:58:49 +00004022 &cwr_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004023 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00004024 &dropwhile_type,
4025 &takewhile_type,
4026 &islice_type,
4027 &starmap_type,
4028 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004029 &chain_type,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00004030 &compress_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00004031 &ifilter_type,
4032 &ifilterfalse_type,
4033 &count_type,
4034 &izip_type,
Raymond Hettinger50986cc2008-02-22 03:16:42 +00004035 &iziplongest_type,
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00004036 &permutations_type,
David Wolever2724ab92008-03-19 02:35:45 +00004037 &product_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00004038 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00004039 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00004040 NULL
4041 };
4042
Christian Heimese93237d2007-12-19 02:37:44 +00004043 Py_TYPE(&teedataobject_type) = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00004044 m = Py_InitModule3("itertools", module_methods, module_doc);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00004045 if (m == NULL)
4046 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004047
Raymond Hettinger60eca932003-02-09 06:40:58 +00004048 for (i=0 ; typelist[i] != NULL ; i++) {
4049 if (PyType_Ready(typelist[i]) < 0)
4050 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00004051 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004052 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00004053 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00004054 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00004055 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004056
4057 if (PyType_Ready(&teedataobject_type) < 0)
4058 return;
4059 if (PyType_Ready(&tee_type) < 0)
4060 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00004061 if (PyType_Ready(&_grouper_type) < 0)
4062 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004063}