blob: 58c30f930d9450ca3e20f16e0635091970930395 [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 Heimes90aa7642007-12-19 02:45:37 +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öwis9f2e3462007-07-21 17:22:18 +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 */
Mark Dickinsone94c6792009-02-02 20:36:42 +0000151 0, /* tp_reserved */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000152 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
Christian Heimesdd15f6c2008-03-16 00:07:10 +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
Christian Heimesdd15f6c2008-03-16 00:07:10 +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{
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000216 PyObject_GC_UnTrack(igo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000217 Py_DECREF(igo->parent);
218 Py_DECREF(igo->tgtkey);
Christian Heimesdd15f6c2008-03-16 00:07:10 +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öwis9f2e3462007-07-21 17:22:18 +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 */
Mark Dickinsone94c6792009-02-02 20:36:42 +0000282 0, /* tp_reserved */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000283 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 */
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000293 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000294 0, /* tp_doc */
Christian Heimesdd15f6c2008-03-16 00:07:10 +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 */
Christian Heimesdd15f6c2008-03-16 00:07:10 +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 Wouters49fd7fa2006-04-21 10:40:58 +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 Wouters49fd7fa2006-04-21 10:40:58 +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);
Thomas Wouters0e3f5912006-08-11 14:57:12 +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 Wouters49fd7fa2006-04-21 10:40:58 +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 Wouters49fd7fa2006-04-21 10:40:58 +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öwis9f2e3462007-07-21 17:22:18 +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 */
Mark Dickinsone94c6792009-02-02 20:36:42 +0000435 0, /* tp_reserved */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000436 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 Wouters49fd7fa2006-04-21 10:40:58 +0000446 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000447 teedataobject_doc, /* tp_doc */
Thomas Wouters49fd7fa2006-04-21 10:40:58 +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);
Thomas Wouters0e3f5912006-08-11 14:57:12 +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 Wouters49fd7fa2006-04-21 10:40:58 +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 Wouters49fd7fa2006-04-21 10:40:58 +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 Wouters49fd7fa2006-04-21 10:40:58 +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 Wouters49fd7fa2006-04-21 10:40:58 +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);
Thomas Wouters0e3f5912006-08-11 14:57:12 +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 Wouters49fd7fa2006-04-21 10:40:58 +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 Wouters49fd7fa2006-04-21 10:40:58 +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 Wouters49fd7fa2006-04-21 10:40:58 +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öwis9f2e3462007-07-21 17:22:18 +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 */
Mark Dickinsone94c6792009-02-02 20:36:42 +0000591 0, /* tp_reserved */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000592 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 Wouters49fd7fa2006-04-21 10:40:58 +0000602 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000603 teeobject_doc, /* tp_doc */
Thomas Wouters49fd7fa2006-04-21 10:40:58 +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 Wouters49fd7fa2006-04-21 10:40:58 +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{
Thomas Wouters89f507f2006-12-13 04:49:30 +0000627 Py_ssize_t i, n=2;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000628 PyObject *it, *iterable, *copyable, *result;
629
Thomas Wouters89f507f2006-12-13 04:49:30 +0000630 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000631 return NULL;
Thomas Wouters89f507f2006-12-13 04:49:30 +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
Thomas Woutersb2137042007-02-01 18:02:27 +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 Heimes90aa7642007-12-19 02:45:37 +0000727 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000728}
729
730static int
731cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
732{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000733 Py_VISIT(lz->it);
734 Py_VISIT(lz->saved);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000735 return 0;
736}
737
738static PyObject *
739cycle_next(cycleobject *lz)
740{
741 PyObject *item;
742 PyObject *it;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000743 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000744
745 while (1) {
746 item = PyIter_Next(lz->it);
747 if (item != NULL) {
748 if (!lz->firstpass)
749 PyList_Append(lz->saved, item);
750 return item;
751 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000752 if (PyErr_Occurred()) {
753 if (PyErr_ExceptionMatches(PyExc_StopIteration))
754 PyErr_Clear();
755 else
756 return NULL;
757 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000758 if (PyList_Size(lz->saved) == 0)
759 return NULL;
760 it = PyObject_GetIter(lz->saved);
761 if (it == NULL)
762 return NULL;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000763 tmp = lz->it;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000764 lz->it = it;
765 lz->firstpass = 1;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000766 Py_DECREF(tmp);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000767 }
768}
769
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000770PyDoc_STRVAR(cycle_doc,
771"cycle(iterable) --> cycle object\n\
772\n\
773Return elements from the iterable until it is exhausted.\n\
774Then repeat the sequence indefinitely.");
775
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000776static PyTypeObject cycle_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +0000777 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000778 "itertools.cycle", /* tp_name */
779 sizeof(cycleobject), /* tp_basicsize */
780 0, /* tp_itemsize */
781 /* methods */
782 (destructor)cycle_dealloc, /* tp_dealloc */
783 0, /* tp_print */
784 0, /* tp_getattr */
785 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +0000786 0, /* tp_reserved */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000787 0, /* tp_repr */
788 0, /* tp_as_number */
789 0, /* tp_as_sequence */
790 0, /* tp_as_mapping */
791 0, /* tp_hash */
792 0, /* tp_call */
793 0, /* tp_str */
794 PyObject_GenericGetAttr, /* tp_getattro */
795 0, /* tp_setattro */
796 0, /* tp_as_buffer */
797 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
798 Py_TPFLAGS_BASETYPE, /* tp_flags */
799 cycle_doc, /* tp_doc */
800 (traverseproc)cycle_traverse, /* tp_traverse */
801 0, /* tp_clear */
802 0, /* tp_richcompare */
803 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000804 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000805 (iternextfunc)cycle_next, /* tp_iternext */
806 0, /* tp_methods */
807 0, /* tp_members */
808 0, /* tp_getset */
809 0, /* tp_base */
810 0, /* tp_dict */
811 0, /* tp_descr_get */
812 0, /* tp_descr_set */
813 0, /* tp_dictoffset */
814 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000815 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000816 cycle_new, /* tp_new */
817 PyObject_GC_Del, /* tp_free */
818};
819
820
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000821/* dropwhile object **********************************************************/
822
823typedef struct {
824 PyObject_HEAD
825 PyObject *func;
826 PyObject *it;
827 long start;
828} dropwhileobject;
829
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000830static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000831
832static PyObject *
833dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
834{
835 PyObject *func, *seq;
836 PyObject *it;
837 dropwhileobject *lz;
838
Thomas Woutersb2137042007-02-01 18:02:27 +0000839 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000840 return NULL;
841
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000842 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
843 return NULL;
844
845 /* Get iterator. */
846 it = PyObject_GetIter(seq);
847 if (it == NULL)
848 return NULL;
849
850 /* create dropwhileobject structure */
851 lz = (dropwhileobject *)type->tp_alloc(type, 0);
852 if (lz == NULL) {
853 Py_DECREF(it);
854 return NULL;
855 }
856 Py_INCREF(func);
857 lz->func = func;
858 lz->it = it;
859 lz->start = 0;
860
861 return (PyObject *)lz;
862}
863
864static void
865dropwhile_dealloc(dropwhileobject *lz)
866{
867 PyObject_GC_UnTrack(lz);
868 Py_XDECREF(lz->func);
869 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +0000870 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000871}
872
873static int
874dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
875{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000876 Py_VISIT(lz->it);
877 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000878 return 0;
879}
880
881static PyObject *
882dropwhile_next(dropwhileobject *lz)
883{
884 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000885 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000886 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000887 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000888
Christian Heimes90aa7642007-12-19 02:45:37 +0000889 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000890 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000891 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000892 if (item == NULL)
893 return NULL;
894 if (lz->start == 1)
895 return item;
896
897 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
898 if (good == NULL) {
899 Py_DECREF(item);
900 return NULL;
901 }
902 ok = PyObject_IsTrue(good);
903 Py_DECREF(good);
904 if (!ok) {
905 lz->start = 1;
906 return item;
907 }
908 Py_DECREF(item);
909 }
910}
911
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000912PyDoc_STRVAR(dropwhile_doc,
913"dropwhile(predicate, iterable) --> dropwhile object\n\
914\n\
915Drop items from the iterable while predicate(item) is true.\n\
916Afterwards, return every element until the iterable is exhausted.");
917
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000918static PyTypeObject dropwhile_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +0000919 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000920 "itertools.dropwhile", /* tp_name */
921 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000922 0, /* tp_itemsize */
923 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000924 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000925 0, /* tp_print */
926 0, /* tp_getattr */
927 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +0000928 0, /* tp_reserved */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000929 0, /* tp_repr */
930 0, /* tp_as_number */
931 0, /* tp_as_sequence */
932 0, /* tp_as_mapping */
933 0, /* tp_hash */
934 0, /* tp_call */
935 0, /* tp_str */
936 PyObject_GenericGetAttr, /* tp_getattro */
937 0, /* tp_setattro */
938 0, /* tp_as_buffer */
939 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
940 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000941 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000942 (traverseproc)dropwhile_traverse, /* tp_traverse */
943 0, /* tp_clear */
944 0, /* tp_richcompare */
945 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000946 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000947 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000948 0, /* tp_methods */
949 0, /* tp_members */
950 0, /* tp_getset */
951 0, /* tp_base */
952 0, /* tp_dict */
953 0, /* tp_descr_get */
954 0, /* tp_descr_set */
955 0, /* tp_dictoffset */
956 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000957 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000958 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000959 PyObject_GC_Del, /* tp_free */
960};
961
962
963/* takewhile object **********************************************************/
964
965typedef struct {
966 PyObject_HEAD
967 PyObject *func;
968 PyObject *it;
969 long stop;
970} takewhileobject;
971
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000972static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000973
974static PyObject *
975takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
976{
977 PyObject *func, *seq;
978 PyObject *it;
979 takewhileobject *lz;
980
Thomas Woutersb2137042007-02-01 18:02:27 +0000981 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000982 return NULL;
983
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000984 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
985 return NULL;
986
987 /* Get iterator. */
988 it = PyObject_GetIter(seq);
989 if (it == NULL)
990 return NULL;
991
992 /* create takewhileobject structure */
993 lz = (takewhileobject *)type->tp_alloc(type, 0);
994 if (lz == NULL) {
995 Py_DECREF(it);
996 return NULL;
997 }
998 Py_INCREF(func);
999 lz->func = func;
1000 lz->it = it;
1001 lz->stop = 0;
1002
1003 return (PyObject *)lz;
1004}
1005
1006static void
1007takewhile_dealloc(takewhileobject *lz)
1008{
1009 PyObject_GC_UnTrack(lz);
1010 Py_XDECREF(lz->func);
1011 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00001012 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001013}
1014
1015static int
1016takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1017{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001018 Py_VISIT(lz->it);
1019 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001020 return 0;
1021}
1022
1023static PyObject *
1024takewhile_next(takewhileobject *lz)
1025{
1026 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001027 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001028 long ok;
1029
1030 if (lz->stop == 1)
1031 return NULL;
1032
Christian Heimes90aa7642007-12-19 02:45:37 +00001033 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001034 if (item == NULL)
1035 return NULL;
1036
1037 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1038 if (good == NULL) {
1039 Py_DECREF(item);
1040 return NULL;
1041 }
1042 ok = PyObject_IsTrue(good);
1043 Py_DECREF(good);
1044 if (ok)
1045 return item;
1046 Py_DECREF(item);
1047 lz->stop = 1;
1048 return NULL;
1049}
1050
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001051PyDoc_STRVAR(takewhile_doc,
1052"takewhile(predicate, iterable) --> takewhile object\n\
1053\n\
1054Return successive entries from an iterable as long as the \n\
1055predicate evaluates to true for each entry.");
1056
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001057static PyTypeObject takewhile_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001058 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001059 "itertools.takewhile", /* tp_name */
1060 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001061 0, /* tp_itemsize */
1062 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001063 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001064 0, /* tp_print */
1065 0, /* tp_getattr */
1066 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00001067 0, /* tp_reserved */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001068 0, /* tp_repr */
1069 0, /* tp_as_number */
1070 0, /* tp_as_sequence */
1071 0, /* tp_as_mapping */
1072 0, /* tp_hash */
1073 0, /* tp_call */
1074 0, /* tp_str */
1075 PyObject_GenericGetAttr, /* tp_getattro */
1076 0, /* tp_setattro */
1077 0, /* tp_as_buffer */
1078 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1079 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001080 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001081 (traverseproc)takewhile_traverse, /* tp_traverse */
1082 0, /* tp_clear */
1083 0, /* tp_richcompare */
1084 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001085 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001086 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001087 0, /* tp_methods */
1088 0, /* tp_members */
1089 0, /* tp_getset */
1090 0, /* tp_base */
1091 0, /* tp_dict */
1092 0, /* tp_descr_get */
1093 0, /* tp_descr_set */
1094 0, /* tp_dictoffset */
1095 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001096 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001097 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001098 PyObject_GC_Del, /* tp_free */
1099};
1100
1101
1102/* islice object ************************************************************/
1103
1104typedef struct {
1105 PyObject_HEAD
1106 PyObject *it;
Thomas Wouters477c8d52006-05-27 19:21:47 +00001107 Py_ssize_t next;
1108 Py_ssize_t stop;
1109 Py_ssize_t step;
1110 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001111} isliceobject;
1112
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001113static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001114
1115static PyObject *
1116islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1117{
1118 PyObject *seq;
Thomas Wouters477c8d52006-05-27 19:21:47 +00001119 Py_ssize_t start=0, stop=-1, step=1;
Raymond Hettingerb2594052004-12-05 09:25:51 +00001120 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001121 Py_ssize_t numargs;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001122 isliceobject *lz;
1123
Thomas Woutersb2137042007-02-01 18:02:27 +00001124 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001125 return NULL;
1126
Raymond Hettingerb2594052004-12-05 09:25:51 +00001127 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001128 return NULL;
1129
Raymond Hettingerb2594052004-12-05 09:25:51 +00001130 numargs = PyTuple_Size(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001131 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001132 if (a1 != Py_None) {
Christian Heimes217cfd12007-12-02 14:31:20 +00001133 stop = PyLong_AsSsize_t(a1);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001134 if (stop == -1) {
1135 if (PyErr_Occurred())
1136 PyErr_Clear();
1137 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001138 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001139 return NULL;
1140 }
1141 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001142 } else {
Raymond Hettingerb2594052004-12-05 09:25:51 +00001143 if (a1 != Py_None)
Christian Heimes217cfd12007-12-02 14:31:20 +00001144 start = PyLong_AsSsize_t(a1);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001145 if (start == -1 && PyErr_Occurred())
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001146 PyErr_Clear();
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001147 if (a2 != Py_None) {
Christian Heimes217cfd12007-12-02 14:31:20 +00001148 stop = PyLong_AsSsize_t(a2);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001149 if (stop == -1) {
1150 if (PyErr_Occurred())
1151 PyErr_Clear();
1152 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001153 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001154 return NULL;
1155 }
1156 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001157 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001158 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001159 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001160 "Indices for islice() must be non-negative integers or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001161 return NULL;
1162 }
1163
Raymond Hettingerb2594052004-12-05 09:25:51 +00001164 if (a3 != NULL) {
1165 if (a3 != Py_None)
Christian Heimes217cfd12007-12-02 14:31:20 +00001166 step = PyLong_AsSsize_t(a3);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001167 if (step == -1 && PyErr_Occurred())
1168 PyErr_Clear();
1169 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001170 if (step<1) {
1171 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001172 "Step for islice() must be a positive integer or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001173 return NULL;
1174 }
1175
1176 /* Get iterator. */
1177 it = PyObject_GetIter(seq);
1178 if (it == NULL)
1179 return NULL;
1180
1181 /* create isliceobject structure */
1182 lz = (isliceobject *)type->tp_alloc(type, 0);
1183 if (lz == NULL) {
1184 Py_DECREF(it);
1185 return NULL;
1186 }
1187 lz->it = it;
1188 lz->next = start;
1189 lz->stop = stop;
1190 lz->step = step;
1191 lz->cnt = 0L;
1192
1193 return (PyObject *)lz;
1194}
1195
1196static void
1197islice_dealloc(isliceobject *lz)
1198{
1199 PyObject_GC_UnTrack(lz);
1200 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00001201 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001202}
1203
1204static int
1205islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1206{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001207 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001208 return 0;
1209}
1210
1211static PyObject *
1212islice_next(isliceobject *lz)
1213{
1214 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001215 PyObject *it = lz->it;
Thomas Wouters477c8d52006-05-27 19:21:47 +00001216 Py_ssize_t oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001217 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001218
Christian Heimes90aa7642007-12-19 02:45:37 +00001219 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001220 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001221 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001222 if (item == NULL)
1223 return NULL;
1224 Py_DECREF(item);
1225 lz->cnt++;
1226 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001227 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001228 return NULL;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001229 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001230 if (item == NULL)
1231 return NULL;
1232 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001233 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001234 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001235 if (lz->next < oldnext) /* Check for overflow */
1236 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001237 return item;
1238}
1239
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001240PyDoc_STRVAR(islice_doc,
1241"islice(iterable, [start,] stop [, step]) --> islice object\n\
1242\n\
1243Return an iterator whose next() method returns selected values from an\n\
1244iterable. If start is specified, will skip all preceding elements;\n\
1245otherwise, start defaults to zero. Step defaults to one. If\n\
1246specified as another value, step determines how many values are \n\
1247skipped between successive calls. Works like a slice() on a list\n\
1248but returns an iterator.");
1249
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001250static PyTypeObject islice_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001251 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001252 "itertools.islice", /* tp_name */
1253 sizeof(isliceobject), /* tp_basicsize */
1254 0, /* tp_itemsize */
1255 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001256 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001257 0, /* tp_print */
1258 0, /* tp_getattr */
1259 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00001260 0, /* tp_reserved */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001261 0, /* tp_repr */
1262 0, /* tp_as_number */
1263 0, /* tp_as_sequence */
1264 0, /* tp_as_mapping */
1265 0, /* tp_hash */
1266 0, /* tp_call */
1267 0, /* tp_str */
1268 PyObject_GenericGetAttr, /* tp_getattro */
1269 0, /* tp_setattro */
1270 0, /* tp_as_buffer */
1271 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1272 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001273 islice_doc, /* tp_doc */
1274 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001275 0, /* tp_clear */
1276 0, /* tp_richcompare */
1277 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001278 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001279 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001280 0, /* tp_methods */
1281 0, /* tp_members */
1282 0, /* tp_getset */
1283 0, /* tp_base */
1284 0, /* tp_dict */
1285 0, /* tp_descr_get */
1286 0, /* tp_descr_set */
1287 0, /* tp_dictoffset */
1288 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001289 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001290 islice_new, /* tp_new */
1291 PyObject_GC_Del, /* tp_free */
1292};
1293
1294
1295/* starmap object ************************************************************/
1296
1297typedef struct {
1298 PyObject_HEAD
1299 PyObject *func;
1300 PyObject *it;
1301} starmapobject;
1302
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001303static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001304
1305static PyObject *
1306starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1307{
1308 PyObject *func, *seq;
1309 PyObject *it;
1310 starmapobject *lz;
1311
Thomas Woutersb2137042007-02-01 18:02:27 +00001312 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001313 return NULL;
1314
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001315 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1316 return NULL;
1317
1318 /* Get iterator. */
1319 it = PyObject_GetIter(seq);
1320 if (it == NULL)
1321 return NULL;
1322
1323 /* create starmapobject structure */
1324 lz = (starmapobject *)type->tp_alloc(type, 0);
1325 if (lz == NULL) {
1326 Py_DECREF(it);
1327 return NULL;
1328 }
1329 Py_INCREF(func);
1330 lz->func = func;
1331 lz->it = it;
1332
1333 return (PyObject *)lz;
1334}
1335
1336static void
1337starmap_dealloc(starmapobject *lz)
1338{
1339 PyObject_GC_UnTrack(lz);
1340 Py_XDECREF(lz->func);
1341 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00001342 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001343}
1344
1345static int
1346starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1347{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001348 Py_VISIT(lz->it);
1349 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001350 return 0;
1351}
1352
1353static PyObject *
1354starmap_next(starmapobject *lz)
1355{
1356 PyObject *args;
1357 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001358 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001359
Christian Heimes90aa7642007-12-19 02:45:37 +00001360 args = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001361 if (args == NULL)
1362 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001363 if (!PyTuple_CheckExact(args)) {
Christian Heimes679db4a2008-01-18 09:56:22 +00001364 PyObject *newargs = PySequence_Tuple(args);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001365 Py_DECREF(args);
Christian Heimes679db4a2008-01-18 09:56:22 +00001366 if (newargs == NULL)
1367 return NULL;
1368 args = newargs;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001369 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001370 result = PyObject_Call(lz->func, args, NULL);
1371 Py_DECREF(args);
1372 return result;
1373}
1374
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001375PyDoc_STRVAR(starmap_doc,
1376"starmap(function, sequence) --> starmap object\n\
1377\n\
1378Return an iterator whose values are returned from the function evaluated\n\
1379with a argument tuple taken from the given sequence.");
1380
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001381static PyTypeObject starmap_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001382 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001383 "itertools.starmap", /* tp_name */
1384 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001385 0, /* tp_itemsize */
1386 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001387 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001388 0, /* tp_print */
1389 0, /* tp_getattr */
1390 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00001391 0, /* tp_reserved */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001392 0, /* tp_repr */
1393 0, /* tp_as_number */
1394 0, /* tp_as_sequence */
1395 0, /* tp_as_mapping */
1396 0, /* tp_hash */
1397 0, /* tp_call */
1398 0, /* tp_str */
1399 PyObject_GenericGetAttr, /* tp_getattro */
1400 0, /* tp_setattro */
1401 0, /* tp_as_buffer */
1402 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1403 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001404 starmap_doc, /* tp_doc */
1405 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001406 0, /* tp_clear */
1407 0, /* tp_richcompare */
1408 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001409 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001410 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001411 0, /* tp_methods */
1412 0, /* tp_members */
1413 0, /* tp_getset */
1414 0, /* tp_base */
1415 0, /* tp_dict */
1416 0, /* tp_descr_get */
1417 0, /* tp_descr_set */
1418 0, /* tp_dictoffset */
1419 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001420 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001421 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001422 PyObject_GC_Del, /* tp_free */
1423};
1424
1425
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001426/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001427
1428typedef struct {
1429 PyObject_HEAD
Christian Heimesf16baeb2008-02-29 14:57:44 +00001430 PyObject *source; /* Iterator over input iterables */
1431 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001432} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001433
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001434static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001435
Christian Heimesf16baeb2008-02-29 14:57:44 +00001436static PyObject *
1437chain_new_internal(PyTypeObject *type, PyObject *source)
1438{
1439 chainobject *lz;
1440
1441 lz = (chainobject *)type->tp_alloc(type, 0);
1442 if (lz == NULL) {
1443 Py_DECREF(source);
1444 return NULL;
1445 }
1446
1447 lz->source = source;
1448 lz->active = NULL;
1449 return (PyObject *)lz;
1450}
1451
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001452static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001453chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001454{
Christian Heimesf16baeb2008-02-29 14:57:44 +00001455 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001456
Thomas Woutersb2137042007-02-01 18:02:27 +00001457 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001458 return NULL;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001459
1460 source = PyObject_GetIter(args);
1461 if (source == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001462 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001463
Christian Heimesf16baeb2008-02-29 14:57:44 +00001464 return chain_new_internal(type, source);
1465}
1466
1467static PyObject *
1468chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1469{
1470 PyObject *source;
1471
1472 source = PyObject_GetIter(arg);
1473 if (source == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001474 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001475
Christian Heimesf16baeb2008-02-29 14:57:44 +00001476 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001477}
1478
1479static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001480chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001481{
1482 PyObject_GC_UnTrack(lz);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001483 Py_XDECREF(lz->active);
1484 Py_XDECREF(lz->source);
Christian Heimes90aa7642007-12-19 02:45:37 +00001485 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001486}
1487
Raymond Hettinger2012f172003-02-07 05:32:58 +00001488static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001489chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001490{
Christian Heimesf16baeb2008-02-29 14:57:44 +00001491 Py_VISIT(lz->source);
1492 Py_VISIT(lz->active);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001493 return 0;
1494}
1495
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001496static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001497chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001498{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001499 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001500
Christian Heimesf16baeb2008-02-29 14:57:44 +00001501 if (lz->source == NULL)
1502 return NULL; /* already stopped */
1503
1504 if (lz->active == NULL) {
1505 PyObject *iterable = PyIter_Next(lz->source);
1506 if (iterable == NULL) {
1507 Py_CLEAR(lz->source);
1508 return NULL; /* no more input sources */
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001509 }
Christian Heimesf16baeb2008-02-29 14:57:44 +00001510 lz->active = PyObject_GetIter(iterable);
Christian Heimes78644762008-03-04 23:39:23 +00001511 Py_DECREF(iterable);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001512 if (lz->active == NULL) {
Christian Heimesf16baeb2008-02-29 14:57:44 +00001513 Py_CLEAR(lz->source);
1514 return NULL; /* input not iterable */
1515 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001516 }
Christian Heimesf16baeb2008-02-29 14:57:44 +00001517 item = PyIter_Next(lz->active);
1518 if (item != NULL)
1519 return item;
1520 if (PyErr_Occurred()) {
1521 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1522 PyErr_Clear();
1523 else
1524 return NULL; /* input raised an exception */
1525 }
1526 Py_CLEAR(lz->active);
1527 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001528}
1529
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001530PyDoc_STRVAR(chain_doc,
1531"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001532\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001533Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001534first iterable until it is exhausted, then elements from the next\n\
1535iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001536
Christian Heimesf16baeb2008-02-29 14:57:44 +00001537PyDoc_STRVAR(chain_from_iterable_doc,
1538"chain.from_iterable(iterable) --> chain object\n\
1539\n\
1540Alternate chain() contructor taking a single iterable argument\n\
1541that evaluates lazily.");
1542
1543static PyMethodDef chain_methods[] = {
1544 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1545 chain_from_iterable_doc},
1546 {NULL, NULL} /* sentinel */
1547};
1548
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001549static PyTypeObject chain_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001550 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001551 "itertools.chain", /* tp_name */
1552 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001553 0, /* tp_itemsize */
1554 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001555 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001556 0, /* tp_print */
1557 0, /* tp_getattr */
1558 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00001559 0, /* tp_reserved */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001560 0, /* tp_repr */
1561 0, /* tp_as_number */
1562 0, /* tp_as_sequence */
1563 0, /* tp_as_mapping */
1564 0, /* tp_hash */
1565 0, /* tp_call */
1566 0, /* tp_str */
1567 PyObject_GenericGetAttr, /* tp_getattro */
1568 0, /* tp_setattro */
1569 0, /* tp_as_buffer */
1570 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1571 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001572 chain_doc, /* tp_doc */
1573 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001574 0, /* tp_clear */
1575 0, /* tp_richcompare */
1576 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001577 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001578 (iternextfunc)chain_next, /* tp_iternext */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001579 chain_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001580 0, /* tp_members */
1581 0, /* tp_getset */
1582 0, /* tp_base */
1583 0, /* tp_dict */
1584 0, /* tp_descr_get */
1585 0, /* tp_descr_set */
1586 0, /* tp_dictoffset */
1587 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001588 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001589 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001590 PyObject_GC_Del, /* tp_free */
1591};
1592
1593
Christian Heimesc3f30c42008-02-22 16:37:40 +00001594/* product object ************************************************************/
1595
1596typedef struct {
1597 PyObject_HEAD
1598 PyObject *pools; /* tuple of pool tuples */
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001599 Py_ssize_t *indices; /* one index per pool */
1600 PyObject *result; /* most recently returned result tuple */
1601 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001602} productobject;
1603
1604static PyTypeObject product_type;
1605
1606static PyObject *
1607product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1608{
1609 productobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001610 Py_ssize_t nargs, npools, repeat=1;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001611 PyObject *pools = NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001612 Py_ssize_t *indices = NULL;
1613 Py_ssize_t i;
1614
Christian Heimesf16baeb2008-02-29 14:57:44 +00001615 if (kwds != NULL) {
1616 char *kwlist[] = {"repeat", 0};
1617 PyObject *tmpargs = PyTuple_New(0);
1618 if (tmpargs == NULL)
1619 return NULL;
1620 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1621 Py_DECREF(tmpargs);
1622 return NULL;
1623 }
1624 Py_DECREF(tmpargs);
1625 if (repeat < 0) {
1626 PyErr_SetString(PyExc_ValueError,
1627 "repeat argument cannot be negative");
1628 return NULL;
1629 }
1630 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001631
1632 assert(PyTuple_Check(args));
Christian Heimesf16baeb2008-02-29 14:57:44 +00001633 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1634 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001635
Christian Heimesc3f30c42008-02-22 16:37:40 +00001636 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
Christian Heimesb558a2e2008-03-02 22:46:37 +00001637 if (indices == NULL) {
Christian Heimesc3f30c42008-02-22 16:37:40 +00001638 PyErr_NoMemory();
1639 goto error;
1640 }
1641
1642 pools = PyTuple_New(npools);
1643 if (pools == NULL)
1644 goto error;
1645
Christian Heimesf16baeb2008-02-29 14:57:44 +00001646 for (i=0; i < nargs ; ++i) {
Christian Heimesc3f30c42008-02-22 16:37:40 +00001647 PyObject *item = PyTuple_GET_ITEM(args, i);
1648 PyObject *pool = PySequence_Tuple(item);
1649 if (pool == NULL)
1650 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001651 PyTuple_SET_ITEM(pools, i, pool);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001652 indices[i] = 0;
1653 }
Christian Heimesf16baeb2008-02-29 14:57:44 +00001654 for ( ; i < npools; ++i) {
1655 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1656 Py_INCREF(pool);
1657 PyTuple_SET_ITEM(pools, i, pool);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001658 indices[i] = 0;
1659 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001660
1661 /* create productobject structure */
1662 lz = (productobject *)type->tp_alloc(type, 0);
Christian Heimes380f7f22008-02-28 11:19:05 +00001663 if (lz == NULL)
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001664 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001665
1666 lz->pools = pools;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001667 lz->indices = indices;
1668 lz->result = NULL;
1669 lz->stopped = 0;
1670
1671 return (PyObject *)lz;
1672
1673error:
Christian Heimesc3f30c42008-02-22 16:37:40 +00001674 if (indices != NULL)
1675 PyMem_Free(indices);
1676 Py_XDECREF(pools);
1677 return NULL;
1678}
1679
1680static void
1681product_dealloc(productobject *lz)
1682{
1683 PyObject_GC_UnTrack(lz);
1684 Py_XDECREF(lz->pools);
1685 Py_XDECREF(lz->result);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001686 if (lz->indices != NULL)
1687 PyMem_Free(lz->indices);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001688 Py_TYPE(lz)->tp_free(lz);
1689}
1690
1691static int
1692product_traverse(productobject *lz, visitproc visit, void *arg)
1693{
1694 Py_VISIT(lz->pools);
1695 Py_VISIT(lz->result);
1696 return 0;
1697}
1698
1699static PyObject *
1700product_next(productobject *lz)
1701{
1702 PyObject *pool;
1703 PyObject *elem;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001704 PyObject *oldelem;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001705 PyObject *pools = lz->pools;
1706 PyObject *result = lz->result;
1707 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1708 Py_ssize_t i;
1709
1710 if (lz->stopped)
1711 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001712
Christian Heimesc3f30c42008-02-22 16:37:40 +00001713 if (result == NULL) {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001714 /* On the first pass, return an initial tuple filled with the
Christian Heimes78644762008-03-04 23:39:23 +00001715 first element from each pool. */
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001716 result = PyTuple_New(npools);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001717 if (result == NULL)
1718 goto empty;
1719 lz->result = result;
1720 for (i=0; i < npools; i++) {
1721 pool = PyTuple_GET_ITEM(pools, i);
1722 if (PyTuple_GET_SIZE(pool) == 0)
1723 goto empty;
1724 elem = PyTuple_GET_ITEM(pool, 0);
1725 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001726 PyTuple_SET_ITEM(result, i, elem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001727 }
1728 } else {
1729 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001730
1731 /* Copy the previous result tuple or re-use it if available */
1732 if (Py_REFCNT(result) > 1) {
1733 PyObject *old_result = result;
1734 result = PyTuple_New(npools);
1735 if (result == NULL)
1736 goto empty;
1737 lz->result = result;
1738 for (i=0; i < npools; i++) {
1739 elem = PyTuple_GET_ITEM(old_result, i);
1740 Py_INCREF(elem);
1741 PyTuple_SET_ITEM(result, i, elem);
1742 }
1743 Py_DECREF(old_result);
1744 }
1745 /* Now, we've got the only copy so we can update it in-place */
Christian Heimesb558a2e2008-03-02 22:46:37 +00001746 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001747
1748 /* Update the pool indices right-to-left. Only advance to the
1749 next pool when the previous one rolls-over */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001750 for (i=npools-1 ; i >= 0 ; i--) {
1751 pool = PyTuple_GET_ITEM(pools, i);
1752 indices[i]++;
Christian Heimesb558a2e2008-03-02 22:46:37 +00001753 if (indices[i] == PyTuple_GET_SIZE(pool)) {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001754 /* Roll-over and advance to next pool */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001755 indices[i] = 0;
1756 elem = PyTuple_GET_ITEM(pool, 0);
1757 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001758 oldelem = PyTuple_GET_ITEM(result, i);
1759 PyTuple_SET_ITEM(result, i, elem);
1760 Py_DECREF(oldelem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001761 } else {
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001762 /* No rollover. Just increment and stop here. */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001763 elem = PyTuple_GET_ITEM(pool, indices[i]);
1764 Py_INCREF(elem);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001765 oldelem = PyTuple_GET_ITEM(result, i);
1766 PyTuple_SET_ITEM(result, i, elem);
1767 Py_DECREF(oldelem);
Christian Heimesc3f30c42008-02-22 16:37:40 +00001768 break;
1769 }
1770 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001771
1772 /* If i is negative, then the indices have all rolled-over
1773 and we're done. */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001774 if (i < 0)
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001775 goto empty;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001776 }
1777
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001778 Py_INCREF(result);
1779 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001780
1781empty:
1782 lz->stopped = 1;
1783 return NULL;
1784}
1785
1786PyDoc_STRVAR(product_doc,
1787"product(*iterables) --> product object\n\
1788\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001789Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00001790For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1791The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1792cycle in a manner similar to an odometer (with the rightmost element changing\n\
1793on every iteration).\n\n\
1794product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1795product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1796
1797static PyTypeObject product_type = {
1798 PyVarObject_HEAD_INIT(NULL, 0)
1799 "itertools.product", /* tp_name */
1800 sizeof(productobject), /* tp_basicsize */
1801 0, /* tp_itemsize */
1802 /* methods */
1803 (destructor)product_dealloc, /* tp_dealloc */
1804 0, /* tp_print */
1805 0, /* tp_getattr */
1806 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00001807 0, /* tp_reserved */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001808 0, /* tp_repr */
1809 0, /* tp_as_number */
1810 0, /* tp_as_sequence */
1811 0, /* tp_as_mapping */
1812 0, /* tp_hash */
1813 0, /* tp_call */
1814 0, /* tp_str */
1815 PyObject_GenericGetAttr, /* tp_getattro */
1816 0, /* tp_setattro */
1817 0, /* tp_as_buffer */
1818 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1819 Py_TPFLAGS_BASETYPE, /* tp_flags */
1820 product_doc, /* tp_doc */
1821 (traverseproc)product_traverse, /* tp_traverse */
1822 0, /* tp_clear */
1823 0, /* tp_richcompare */
1824 0, /* tp_weaklistoffset */
1825 PyObject_SelfIter, /* tp_iter */
1826 (iternextfunc)product_next, /* tp_iternext */
1827 0, /* tp_methods */
1828 0, /* tp_members */
1829 0, /* tp_getset */
1830 0, /* tp_base */
1831 0, /* tp_dict */
1832 0, /* tp_descr_get */
1833 0, /* tp_descr_set */
1834 0, /* tp_dictoffset */
1835 0, /* tp_init */
1836 0, /* tp_alloc */
1837 product_new, /* tp_new */
1838 PyObject_GC_Del, /* tp_free */
1839};
1840
1841
Christian Heimes380f7f22008-02-28 11:19:05 +00001842/* combinations object ************************************************************/
1843
1844typedef struct {
1845 PyObject_HEAD
1846 PyObject *pool; /* input converted to a tuple */
1847 Py_ssize_t *indices; /* one index per result element */
1848 PyObject *result; /* most recently returned result tuple */
1849 Py_ssize_t r; /* size of result tuple */
1850 int stopped; /* set to 1 when the combinations iterator is exhausted */
1851} combinationsobject;
1852
1853static PyTypeObject combinations_type;
1854
1855static PyObject *
1856combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1857{
1858 combinationsobject *co;
1859 Py_ssize_t n;
1860 Py_ssize_t r;
1861 PyObject *pool = NULL;
1862 PyObject *iterable = NULL;
1863 Py_ssize_t *indices = NULL;
1864 Py_ssize_t i;
1865 static char *kwargs[] = {"iterable", "r", NULL};
1866
1867 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
1868 &iterable, &r))
1869 return NULL;
1870
1871 pool = PySequence_Tuple(iterable);
1872 if (pool == NULL)
1873 goto error;
1874 n = PyTuple_GET_SIZE(pool);
1875 if (r < 0) {
1876 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
1877 goto error;
1878 }
Christian Heimes380f7f22008-02-28 11:19:05 +00001879
1880 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
1881 if (indices == NULL) {
1882 PyErr_NoMemory();
1883 goto error;
1884 }
1885
1886 for (i=0 ; i<r ; i++)
1887 indices[i] = i;
1888
1889 /* create combinationsobject structure */
1890 co = (combinationsobject *)type->tp_alloc(type, 0);
1891 if (co == NULL)
1892 goto error;
1893
1894 co->pool = pool;
1895 co->indices = indices;
1896 co->result = NULL;
1897 co->r = r;
Raymond Hettinger5bad41e2009-01-08 21:01:54 +00001898 co->stopped = r > n ? 1 : 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00001899
1900 return (PyObject *)co;
1901
1902error:
1903 if (indices != NULL)
1904 PyMem_Free(indices);
1905 Py_XDECREF(pool);
1906 return NULL;
1907}
1908
1909static void
1910combinations_dealloc(combinationsobject *co)
1911{
1912 PyObject_GC_UnTrack(co);
1913 Py_XDECREF(co->pool);
1914 Py_XDECREF(co->result);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001915 if (co->indices != NULL)
1916 PyMem_Free(co->indices);
Christian Heimes380f7f22008-02-28 11:19:05 +00001917 Py_TYPE(co)->tp_free(co);
1918}
1919
1920static int
1921combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
1922{
1923 Py_VISIT(co->pool);
1924 Py_VISIT(co->result);
1925 return 0;
1926}
1927
1928static PyObject *
1929combinations_next(combinationsobject *co)
1930{
1931 PyObject *elem;
1932 PyObject *oldelem;
1933 PyObject *pool = co->pool;
1934 Py_ssize_t *indices = co->indices;
1935 PyObject *result = co->result;
1936 Py_ssize_t n = PyTuple_GET_SIZE(pool);
1937 Py_ssize_t r = co->r;
1938 Py_ssize_t i, j, index;
1939
1940 if (co->stopped)
1941 return NULL;
1942
1943 if (result == NULL) {
1944 /* On the first pass, initialize result tuple using the indices */
1945 result = PyTuple_New(r);
1946 if (result == NULL)
1947 goto empty;
1948 co->result = result;
1949 for (i=0; i<r ; i++) {
1950 index = indices[i];
1951 elem = PyTuple_GET_ITEM(pool, index);
1952 Py_INCREF(elem);
1953 PyTuple_SET_ITEM(result, i, elem);
1954 }
1955 } else {
1956 /* Copy the previous result tuple or re-use it if available */
1957 if (Py_REFCNT(result) > 1) {
1958 PyObject *old_result = result;
1959 result = PyTuple_New(r);
1960 if (result == NULL)
1961 goto empty;
1962 co->result = result;
1963 for (i=0; i<r ; i++) {
1964 elem = PyTuple_GET_ITEM(old_result, i);
1965 Py_INCREF(elem);
1966 PyTuple_SET_ITEM(result, i, elem);
1967 }
1968 Py_DECREF(old_result);
1969 }
1970 /* Now, we've got the only copy so we can update it in-place
1971 * CPython's empty tuple is a singleton and cached in
1972 * PyTuple's freelist.
1973 */
1974 assert(r == 0 || Py_REFCNT(result) == 1);
1975
1976 /* Scan indices right-to-left until finding one that is not
1977 at its maximum (i + n - r). */
1978 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
1979 ;
1980
1981 /* If i is negative, then the indices are all at
1982 their maximum value and we're done. */
1983 if (i < 0)
1984 goto empty;
1985
1986 /* Increment the current index which we know is not at its
1987 maximum. Then move back to the right setting each index
1988 to its lowest possible value (one higher than the index
1989 to its left -- this maintains the sort order invariant). */
1990 indices[i]++;
1991 for (j=i+1 ; j<r ; j++)
1992 indices[j] = indices[j-1] + 1;
1993
1994 /* Update the result tuple for the new indices
1995 starting with i, the leftmost index that changed */
1996 for ( ; i<r ; i++) {
1997 index = indices[i];
1998 elem = PyTuple_GET_ITEM(pool, index);
1999 Py_INCREF(elem);
2000 oldelem = PyTuple_GET_ITEM(result, i);
2001 PyTuple_SET_ITEM(result, i, elem);
2002 Py_DECREF(oldelem);
2003 }
2004 }
2005
2006 Py_INCREF(result);
2007 return result;
2008
2009empty:
2010 co->stopped = 1;
2011 return NULL;
2012}
2013
2014PyDoc_STRVAR(combinations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00002015"combinations(iterable[, r]) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002016\n\
2017Return successive r-length combinations of elements in the iterable.\n\n\
2018combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2019
2020static PyTypeObject combinations_type = {
2021 PyVarObject_HEAD_INIT(NULL, 0)
2022 "itertools.combinations", /* tp_name */
2023 sizeof(combinationsobject), /* tp_basicsize */
2024 0, /* tp_itemsize */
2025 /* methods */
2026 (destructor)combinations_dealloc, /* tp_dealloc */
2027 0, /* tp_print */
2028 0, /* tp_getattr */
2029 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00002030 0, /* tp_reserved */
Christian Heimes380f7f22008-02-28 11:19:05 +00002031 0, /* tp_repr */
2032 0, /* tp_as_number */
2033 0, /* tp_as_sequence */
2034 0, /* tp_as_mapping */
2035 0, /* tp_hash */
2036 0, /* tp_call */
2037 0, /* tp_str */
2038 PyObject_GenericGetAttr, /* tp_getattro */
2039 0, /* tp_setattro */
2040 0, /* tp_as_buffer */
2041 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2042 Py_TPFLAGS_BASETYPE, /* tp_flags */
2043 combinations_doc, /* tp_doc */
2044 (traverseproc)combinations_traverse, /* tp_traverse */
2045 0, /* tp_clear */
2046 0, /* tp_richcompare */
2047 0, /* tp_weaklistoffset */
2048 PyObject_SelfIter, /* tp_iter */
2049 (iternextfunc)combinations_next, /* tp_iternext */
2050 0, /* tp_methods */
2051 0, /* tp_members */
2052 0, /* tp_getset */
2053 0, /* tp_base */
2054 0, /* tp_dict */
2055 0, /* tp_descr_get */
2056 0, /* tp_descr_set */
2057 0, /* tp_dictoffset */
2058 0, /* tp_init */
2059 0, /* tp_alloc */
2060 combinations_new, /* tp_new */
2061 PyObject_GC_Del, /* tp_free */
2062};
2063
2064
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002065/* combinations with replacement object *******************************************/
2066
2067/* Equivalent to:
2068
2069 def combinations_with_replacement(iterable, r):
2070 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2071 # number items returned: (n+r-1)! / r! / (n-1)!
2072 pool = tuple(iterable)
2073 n = len(pool)
2074 indices = [0] * r
2075 yield tuple(pool[i] for i in indices)
2076 while 1:
2077 for i in reversed(range(r)):
2078 if indices[i] != n - 1:
2079 break
2080 else:
2081 return
2082 indices[i:] = [indices[i] + 1] * (r - i)
2083 yield tuple(pool[i] for i in indices)
2084
2085 def combinations_with_replacement2(iterable, r):
2086 'Alternate version that filters from product()'
2087 pool = tuple(iterable)
2088 n = len(pool)
2089 for indices in product(range(n), repeat=r):
2090 if sorted(indices) == list(indices):
2091 yield tuple(pool[i] for i in indices)
2092*/
2093typedef struct {
2094 PyObject_HEAD
2095 PyObject *pool; /* input converted to a tuple */
2096 Py_ssize_t *indices; /* one index per result element */
2097 PyObject *result; /* most recently returned result tuple */
2098 Py_ssize_t r; /* size of result tuple */
2099 int stopped; /* set to 1 when the cwr iterator is exhausted */
2100} cwrobject;
2101
2102static PyTypeObject cwr_type;
2103
2104static PyObject *
2105cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2106{
2107 cwrobject *co;
2108 Py_ssize_t n;
2109 Py_ssize_t r;
2110 PyObject *pool = NULL;
2111 PyObject *iterable = NULL;
2112 Py_ssize_t *indices = NULL;
2113 Py_ssize_t i;
2114 static char *kwargs[] = {"iterable", "r", NULL};
2115
2116 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2117 &iterable, &r))
2118 return NULL;
2119
2120 pool = PySequence_Tuple(iterable);
2121 if (pool == NULL)
2122 goto error;
2123 n = PyTuple_GET_SIZE(pool);
2124 if (r < 0) {
2125 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2126 goto error;
2127 }
2128
2129 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2130 if (indices == NULL) {
2131 PyErr_NoMemory();
2132 goto error;
2133 }
2134
2135 for (i=0 ; i<r ; i++)
2136 indices[i] = 0;
2137
2138 /* create cwrobject structure */
2139 co = (cwrobject *)type->tp_alloc(type, 0);
2140 if (co == NULL)
2141 goto error;
2142
2143 co->pool = pool;
2144 co->indices = indices;
2145 co->result = NULL;
2146 co->r = r;
2147 co->stopped = !n && r;
2148
2149 return (PyObject *)co;
2150
2151error:
2152 if (indices != NULL)
2153 PyMem_Free(indices);
2154 Py_XDECREF(pool);
2155 return NULL;
2156}
2157
2158static void
2159cwr_dealloc(cwrobject *co)
2160{
2161 PyObject_GC_UnTrack(co);
2162 Py_XDECREF(co->pool);
2163 Py_XDECREF(co->result);
2164 if (co->indices != NULL)
2165 PyMem_Free(co->indices);
2166 Py_TYPE(co)->tp_free(co);
2167}
2168
2169static int
2170cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2171{
2172 Py_VISIT(co->pool);
2173 Py_VISIT(co->result);
2174 return 0;
2175}
2176
2177static PyObject *
2178cwr_next(cwrobject *co)
2179{
2180 PyObject *elem;
2181 PyObject *oldelem;
2182 PyObject *pool = co->pool;
2183 Py_ssize_t *indices = co->indices;
2184 PyObject *result = co->result;
2185 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2186 Py_ssize_t r = co->r;
2187 Py_ssize_t i, j, index;
2188
2189 if (co->stopped)
2190 return NULL;
2191
2192 if (result == NULL) {
2193 /* On the first pass, initialize result tuple using the indices */
2194 result = PyTuple_New(r);
2195 if (result == NULL)
2196 goto empty;
2197 co->result = result;
2198 for (i=0; i<r ; i++) {
2199 index = indices[i];
2200 elem = PyTuple_GET_ITEM(pool, index);
2201 Py_INCREF(elem);
2202 PyTuple_SET_ITEM(result, i, elem);
2203 }
2204 } else {
2205 /* Copy the previous result tuple or re-use it if available */
2206 if (Py_REFCNT(result) > 1) {
2207 PyObject *old_result = result;
2208 result = PyTuple_New(r);
2209 if (result == NULL)
2210 goto empty;
2211 co->result = result;
2212 for (i=0; i<r ; i++) {
2213 elem = PyTuple_GET_ITEM(old_result, i);
2214 Py_INCREF(elem);
2215 PyTuple_SET_ITEM(result, i, elem);
2216 }
2217 Py_DECREF(old_result);
2218 }
2219 /* Now, we've got the only copy so we can update it in-place CPython's
2220 empty tuple is a singleton and cached in PyTuple's freelist. */
2221 assert(r == 0 || Py_REFCNT(result) == 1);
2222
2223 /* Scan indices right-to-left until finding one that is not
2224 * at its maximum (n-1). */
2225 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2226 ;
2227
2228 /* If i is negative, then the indices are all at
2229 their maximum value and we're done. */
2230 if (i < 0)
2231 goto empty;
2232
2233 /* Increment the current index which we know is not at its
2234 maximum. Then set all to the right to the same value. */
2235 indices[i]++;
2236 for (j=i+1 ; j<r ; j++)
2237 indices[j] = indices[j-1];
2238
2239 /* Update the result tuple for the new indices
2240 starting with i, the leftmost index that changed */
2241 for ( ; i<r ; i++) {
2242 index = indices[i];
2243 elem = PyTuple_GET_ITEM(pool, index);
2244 Py_INCREF(elem);
2245 oldelem = PyTuple_GET_ITEM(result, i);
2246 PyTuple_SET_ITEM(result, i, elem);
2247 Py_DECREF(oldelem);
2248 }
2249 }
2250
2251 Py_INCREF(result);
2252 return result;
2253
2254empty:
2255 co->stopped = 1;
2256 return NULL;
2257}
2258
2259PyDoc_STRVAR(cwr_doc,
2260"combinations_with_replacement(iterable[, r]) --> combinations_with_replacement object\n\
2261\n\
2262Return successive r-length combinations of elements in the iterable\n\
2263allowing individual elements to have successive repeats.\n\
2264combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2265
2266static PyTypeObject cwr_type = {
2267 PyVarObject_HEAD_INIT(NULL, 0)
2268 "itertools.combinations_with_replacement", /* tp_name */
2269 sizeof(cwrobject), /* tp_basicsize */
2270 0, /* tp_itemsize */
2271 /* methods */
2272 (destructor)cwr_dealloc, /* tp_dealloc */
2273 0, /* tp_print */
2274 0, /* tp_getattr */
2275 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00002276 0, /* tp_reserved */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002277 0, /* tp_repr */
2278 0, /* tp_as_number */
2279 0, /* tp_as_sequence */
2280 0, /* tp_as_mapping */
2281 0, /* tp_hash */
2282 0, /* tp_call */
2283 0, /* tp_str */
2284 PyObject_GenericGetAttr, /* tp_getattro */
2285 0, /* tp_setattro */
2286 0, /* tp_as_buffer */
2287 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2288 Py_TPFLAGS_BASETYPE, /* tp_flags */
2289 cwr_doc, /* tp_doc */
2290 (traverseproc)cwr_traverse, /* tp_traverse */
2291 0, /* tp_clear */
2292 0, /* tp_richcompare */
2293 0, /* tp_weaklistoffset */
2294 PyObject_SelfIter, /* tp_iter */
2295 (iternextfunc)cwr_next, /* tp_iternext */
2296 0, /* tp_methods */
2297 0, /* tp_members */
2298 0, /* tp_getset */
2299 0, /* tp_base */
2300 0, /* tp_dict */
2301 0, /* tp_descr_get */
2302 0, /* tp_descr_set */
2303 0, /* tp_dictoffset */
2304 0, /* tp_init */
2305 0, /* tp_alloc */
2306 cwr_new, /* tp_new */
2307 PyObject_GC_Del, /* tp_free */
2308};
2309
2310
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002311/* permutations object ************************************************************
2312
2313def permutations(iterable, r=None):
2314 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2315 pool = tuple(iterable)
2316 n = len(pool)
2317 r = n if r is None else r
2318 indices = range(n)
2319 cycles = range(n-r+1, n+1)[::-1]
2320 yield tuple(pool[i] for i in indices[:r])
2321 while n:
2322 for i in reversed(range(r)):
2323 cycles[i] -= 1
2324 if cycles[i] == 0:
2325 indices[i:] = indices[i+1:] + indices[i:i+1]
2326 cycles[i] = n - i
2327 else:
2328 j = cycles[i]
2329 indices[i], indices[-j] = indices[-j], indices[i]
2330 yield tuple(pool[i] for i in indices[:r])
2331 break
2332 else:
2333 return
2334*/
2335
2336typedef struct {
2337 PyObject_HEAD
2338 PyObject *pool; /* input converted to a tuple */
2339 Py_ssize_t *indices; /* one index per element in the pool */
2340 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2341 PyObject *result; /* most recently returned result tuple */
2342 Py_ssize_t r; /* size of result tuple */
2343 int stopped; /* set to 1 when the permutations iterator is exhausted */
2344} permutationsobject;
2345
2346static PyTypeObject permutations_type;
2347
2348static PyObject *
2349permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2350{
2351 permutationsobject *po;
2352 Py_ssize_t n;
2353 Py_ssize_t r;
2354 PyObject *robj = Py_None;
2355 PyObject *pool = NULL;
2356 PyObject *iterable = NULL;
2357 Py_ssize_t *indices = NULL;
2358 Py_ssize_t *cycles = NULL;
2359 Py_ssize_t i;
2360 static char *kwargs[] = {"iterable", "r", NULL};
2361
2362 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2363 &iterable, &robj))
2364 return NULL;
2365
2366 pool = PySequence_Tuple(iterable);
2367 if (pool == NULL)
2368 goto error;
2369 n = PyTuple_GET_SIZE(pool);
2370
2371 r = n;
2372 if (robj != Py_None) {
2373 if (!PyLong_Check(robj)) {
2374 PyErr_SetString(PyExc_TypeError, "Expected int as r");
Neal Norwitzd2bee322008-04-01 07:37:58 +00002375 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002376 }
2377 r = PyLong_AsSsize_t(robj);
2378 if (r == -1 && PyErr_Occurred())
2379 goto error;
2380 }
2381 if (r < 0) {
2382 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2383 goto error;
2384 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002385
2386 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2387 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2388 if (indices == NULL || cycles == NULL) {
2389 PyErr_NoMemory();
2390 goto error;
2391 }
2392
2393 for (i=0 ; i<n ; i++)
2394 indices[i] = i;
2395 for (i=0 ; i<r ; i++)
2396 cycles[i] = n - i;
2397
2398 /* create permutationsobject structure */
2399 po = (permutationsobject *)type->tp_alloc(type, 0);
2400 if (po == NULL)
2401 goto error;
2402
2403 po->pool = pool;
2404 po->indices = indices;
2405 po->cycles = cycles;
2406 po->result = NULL;
2407 po->r = r;
Raymond Hettinger5bad41e2009-01-08 21:01:54 +00002408 po->stopped = r > n ? 1 : 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002409
2410 return (PyObject *)po;
2411
2412error:
2413 if (indices != NULL)
2414 PyMem_Free(indices);
2415 if (cycles != NULL)
2416 PyMem_Free(cycles);
2417 Py_XDECREF(pool);
2418 return NULL;
2419}
2420
2421static void
2422permutations_dealloc(permutationsobject *po)
2423{
2424 PyObject_GC_UnTrack(po);
2425 Py_XDECREF(po->pool);
2426 Py_XDECREF(po->result);
2427 PyMem_Free(po->indices);
2428 PyMem_Free(po->cycles);
2429 Py_TYPE(po)->tp_free(po);
2430}
2431
2432static int
2433permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2434{
2435 Py_VISIT(po->pool);
2436 Py_VISIT(po->result);
2437 return 0;
2438}
2439
2440static PyObject *
2441permutations_next(permutationsobject *po)
2442{
2443 PyObject *elem;
2444 PyObject *oldelem;
2445 PyObject *pool = po->pool;
2446 Py_ssize_t *indices = po->indices;
2447 Py_ssize_t *cycles = po->cycles;
2448 PyObject *result = po->result;
2449 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2450 Py_ssize_t r = po->r;
2451 Py_ssize_t i, j, k, index;
2452
2453 if (po->stopped)
2454 return NULL;
2455
2456 if (result == NULL) {
2457 /* On the first pass, initialize result tuple using the indices */
2458 result = PyTuple_New(r);
2459 if (result == NULL)
2460 goto empty;
2461 po->result = result;
2462 for (i=0; i<r ; i++) {
2463 index = indices[i];
2464 elem = PyTuple_GET_ITEM(pool, index);
2465 Py_INCREF(elem);
2466 PyTuple_SET_ITEM(result, i, elem);
2467 }
2468 } else {
2469 if (n == 0)
2470 goto empty;
2471
2472 /* Copy the previous result tuple or re-use it if available */
2473 if (Py_REFCNT(result) > 1) {
2474 PyObject *old_result = result;
2475 result = PyTuple_New(r);
2476 if (result == NULL)
2477 goto empty;
2478 po->result = result;
2479 for (i=0; i<r ; i++) {
2480 elem = PyTuple_GET_ITEM(old_result, i);
2481 Py_INCREF(elem);
2482 PyTuple_SET_ITEM(result, i, elem);
2483 }
2484 Py_DECREF(old_result);
2485 }
2486 /* Now, we've got the only copy so we can update it in-place */
2487 assert(r == 0 || Py_REFCNT(result) == 1);
2488
2489 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2490 for (i=r-1 ; i>=0 ; i--) {
2491 cycles[i] -= 1;
2492 if (cycles[i] == 0) {
2493 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2494 index = indices[i];
2495 for (j=i ; j<n-1 ; j++)
2496 indices[j] = indices[j+1];
2497 indices[n-1] = index;
2498 cycles[i] = n - i;
2499 } else {
2500 j = cycles[i];
2501 index = indices[i];
2502 indices[i] = indices[n-j];
2503 indices[n-j] = index;
2504
2505 for (k=i; k<r ; k++) {
2506 /* start with i, the leftmost element that changed */
2507 /* yield tuple(pool[k] for k in indices[:r]) */
2508 index = indices[k];
2509 elem = PyTuple_GET_ITEM(pool, index);
2510 Py_INCREF(elem);
2511 oldelem = PyTuple_GET_ITEM(result, k);
2512 PyTuple_SET_ITEM(result, k, elem);
2513 Py_DECREF(oldelem);
2514 }
2515 break;
2516 }
2517 }
2518 /* If i is negative, then the cycles have all
2519 rolled-over and we're done. */
2520 if (i < 0)
2521 goto empty;
2522 }
2523 Py_INCREF(result);
2524 return result;
2525
2526empty:
2527 po->stopped = 1;
2528 return NULL;
2529}
2530
2531PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00002532"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002533\n\
2534Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00002535permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002536
2537static PyTypeObject permutations_type = {
2538 PyVarObject_HEAD_INIT(NULL, 0)
2539 "itertools.permutations", /* tp_name */
2540 sizeof(permutationsobject), /* tp_basicsize */
2541 0, /* tp_itemsize */
2542 /* methods */
2543 (destructor)permutations_dealloc, /* tp_dealloc */
2544 0, /* tp_print */
2545 0, /* tp_getattr */
2546 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00002547 0, /* tp_reserved */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002548 0, /* tp_repr */
2549 0, /* tp_as_number */
2550 0, /* tp_as_sequence */
2551 0, /* tp_as_mapping */
2552 0, /* tp_hash */
2553 0, /* tp_call */
2554 0, /* tp_str */
2555 PyObject_GenericGetAttr, /* tp_getattro */
2556 0, /* tp_setattro */
2557 0, /* tp_as_buffer */
2558 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2559 Py_TPFLAGS_BASETYPE, /* tp_flags */
2560 permutations_doc, /* tp_doc */
2561 (traverseproc)permutations_traverse, /* tp_traverse */
2562 0, /* tp_clear */
2563 0, /* tp_richcompare */
2564 0, /* tp_weaklistoffset */
2565 PyObject_SelfIter, /* tp_iter */
2566 (iternextfunc)permutations_next, /* tp_iternext */
2567 0, /* tp_methods */
2568 0, /* tp_members */
2569 0, /* tp_getset */
2570 0, /* tp_base */
2571 0, /* tp_dict */
2572 0, /* tp_descr_get */
2573 0, /* tp_descr_set */
2574 0, /* tp_dictoffset */
2575 0, /* tp_init */
2576 0, /* tp_alloc */
2577 permutations_new, /* tp_new */
2578 PyObject_GC_Del, /* tp_free */
2579};
2580
2581
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002582/* compress object ************************************************************/
2583
2584/* Equivalent to:
2585
2586 def compress(data, selectors):
2587 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2588 return (d for d, s in zip(data, selectors) if s)
2589*/
2590
2591typedef struct {
2592 PyObject_HEAD
2593 PyObject *data;
2594 PyObject *selectors;
2595} compressobject;
2596
2597static PyTypeObject compress_type;
2598
2599static PyObject *
2600compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2601{
2602 PyObject *seq1, *seq2;
2603 PyObject *data=NULL, *selectors=NULL;
2604 compressobject *lz;
2605
2606 if (type == &compress_type && !_PyArg_NoKeywords("compress()", kwds))
2607 return NULL;
2608
2609 if (!PyArg_UnpackTuple(args, "compress", 2, 2, &seq1, &seq2))
2610 return NULL;
2611
2612 data = PyObject_GetIter(seq1);
2613 if (data == NULL)
2614 goto fail;
2615 selectors = PyObject_GetIter(seq2);
2616 if (selectors == NULL)
2617 goto fail;
2618
2619 /* create compressobject structure */
2620 lz = (compressobject *)type->tp_alloc(type, 0);
2621 if (lz == NULL)
2622 goto fail;
2623 lz->data = data;
2624 lz->selectors = selectors;
2625 return (PyObject *)lz;
2626
2627fail:
2628 Py_XDECREF(data);
2629 Py_XDECREF(selectors);
2630 return NULL;
2631}
2632
2633static void
2634compress_dealloc(compressobject *lz)
2635{
2636 PyObject_GC_UnTrack(lz);
2637 Py_XDECREF(lz->data);
2638 Py_XDECREF(lz->selectors);
2639 Py_TYPE(lz)->tp_free(lz);
2640}
2641
2642static int
2643compress_traverse(compressobject *lz, visitproc visit, void *arg)
2644{
2645 Py_VISIT(lz->data);
2646 Py_VISIT(lz->selectors);
2647 return 0;
2648}
2649
2650static PyObject *
2651compress_next(compressobject *lz)
2652{
2653 PyObject *data = lz->data, *selectors = lz->selectors;
2654 PyObject *datum, *selector;
2655 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2656 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2657 int ok;
2658
2659 while (1) {
2660 /* Steps: get datum, get selector, evaluate selector.
2661 Order is important (to match the pure python version
2662 in terms of which input gets a chance to raise an
2663 exception first).
2664 */
2665
2666 datum = datanext(data);
2667 if (datum == NULL)
2668 return NULL;
2669
2670 selector = selectornext(selectors);
2671 if (selector == NULL) {
2672 Py_DECREF(datum);
2673 return NULL;
2674 }
2675
2676 ok = PyObject_IsTrue(selector);
2677 Py_DECREF(selector);
2678 if (ok == 1)
2679 return datum;
2680 Py_DECREF(datum);
2681 if (ok == -1)
2682 return NULL;
2683 }
2684}
2685
2686PyDoc_STRVAR(compress_doc,
2687"compress(data sequence, selector sequence) --> iterator over selected data\n\
2688\n\
2689Return data elements corresponding to true selector elements.\n\
2690Forms a shorter iterator from selected data elements using the\n\
2691selectors to choose the data elements.");
2692
2693static PyTypeObject compress_type = {
2694 PyVarObject_HEAD_INIT(NULL, 0)
2695 "itertools.compress", /* tp_name */
2696 sizeof(compressobject), /* tp_basicsize */
2697 0, /* tp_itemsize */
2698 /* methods */
2699 (destructor)compress_dealloc, /* tp_dealloc */
2700 0, /* tp_print */
2701 0, /* tp_getattr */
2702 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00002703 0, /* tp_reserved */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002704 0, /* tp_repr */
2705 0, /* tp_as_number */
2706 0, /* tp_as_sequence */
2707 0, /* tp_as_mapping */
2708 0, /* tp_hash */
2709 0, /* tp_call */
2710 0, /* tp_str */
2711 PyObject_GenericGetAttr, /* tp_getattro */
2712 0, /* tp_setattro */
2713 0, /* tp_as_buffer */
2714 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2715 Py_TPFLAGS_BASETYPE, /* tp_flags */
2716 compress_doc, /* tp_doc */
2717 (traverseproc)compress_traverse, /* tp_traverse */
2718 0, /* tp_clear */
2719 0, /* tp_richcompare */
2720 0, /* tp_weaklistoffset */
2721 PyObject_SelfIter, /* tp_iter */
2722 (iternextfunc)compress_next, /* tp_iternext */
2723 0, /* tp_methods */
2724 0, /* tp_members */
2725 0, /* tp_getset */
2726 0, /* tp_base */
2727 0, /* tp_dict */
2728 0, /* tp_descr_get */
2729 0, /* tp_descr_set */
2730 0, /* tp_dictoffset */
2731 0, /* tp_init */
2732 0, /* tp_alloc */
2733 compress_new, /* tp_new */
2734 PyObject_GC_Del, /* tp_free */
2735};
2736
2737
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002738/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00002739
2740typedef struct {
2741 PyObject_HEAD
2742 PyObject *func;
2743 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002744} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002745
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002746static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002747
2748static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002749filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002750{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00002751 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002752 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002753 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002754
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002755 if (type == &filterfalse_type &&
2756 !_PyArg_NoKeywords("filterfalse()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002757 return NULL;
2758
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002759 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
Raymond Hettinger60eca932003-02-09 06:40:58 +00002760 return NULL;
2761
2762 /* Get iterator. */
2763 it = PyObject_GetIter(seq);
2764 if (it == NULL)
2765 return NULL;
2766
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002767 /* create filterfalseobject structure */
2768 lz = (filterfalseobject *)type->tp_alloc(type, 0);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002769 if (lz == NULL) {
2770 Py_DECREF(it);
2771 return NULL;
2772 }
2773 Py_INCREF(func);
2774 lz->func = func;
2775 lz->it = it;
2776
2777 return (PyObject *)lz;
2778}
2779
2780static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002781filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002782{
2783 PyObject_GC_UnTrack(lz);
2784 Py_XDECREF(lz->func);
2785 Py_XDECREF(lz->it);
Christian Heimes90aa7642007-12-19 02:45:37 +00002786 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002787}
2788
2789static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002790filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002791{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002792 Py_VISIT(lz->it);
2793 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002794 return 0;
2795}
2796
2797static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002798filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002799{
2800 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002801 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002802 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002803 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002804
Christian Heimes90aa7642007-12-19 02:45:37 +00002805 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002806 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002807 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002808 if (item == NULL)
2809 return NULL;
2810
Christian Heimes836baa52008-02-26 08:18:30 +00002811 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
Raymond Hettinger60eca932003-02-09 06:40:58 +00002812 ok = PyObject_IsTrue(item);
2813 } else {
2814 PyObject *good;
2815 good = PyObject_CallFunctionObjArgs(lz->func,
2816 item, NULL);
2817 if (good == NULL) {
2818 Py_DECREF(item);
2819 return NULL;
2820 }
2821 ok = PyObject_IsTrue(good);
2822 Py_DECREF(good);
2823 }
2824 if (!ok)
2825 return item;
2826 Py_DECREF(item);
2827 }
2828}
2829
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002830PyDoc_STRVAR(filterfalse_doc,
2831"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002832\n\
2833Return those items of sequence for which function(item) is false.\n\
2834If function is None, return the items that are false.");
2835
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002836static PyTypeObject filterfalse_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002837 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002838 "itertools.filterfalse", /* tp_name */
2839 sizeof(filterfalseobject), /* tp_basicsize */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002840 0, /* tp_itemsize */
2841 /* methods */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002842 (destructor)filterfalse_dealloc, /* tp_dealloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002843 0, /* tp_print */
2844 0, /* tp_getattr */
2845 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00002846 0, /* tp_reserved */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002847 0, /* tp_repr */
2848 0, /* tp_as_number */
2849 0, /* tp_as_sequence */
2850 0, /* tp_as_mapping */
2851 0, /* tp_hash */
2852 0, /* tp_call */
2853 0, /* tp_str */
2854 PyObject_GenericGetAttr, /* tp_getattro */
2855 0, /* tp_setattro */
2856 0, /* tp_as_buffer */
2857 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2858 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002859 filterfalse_doc, /* tp_doc */
2860 (traverseproc)filterfalse_traverse, /* tp_traverse */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002861 0, /* tp_clear */
2862 0, /* tp_richcompare */
2863 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002864 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002865 (iternextfunc)filterfalse_next, /* tp_iternext */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002866 0, /* tp_methods */
2867 0, /* tp_members */
2868 0, /* tp_getset */
2869 0, /* tp_base */
2870 0, /* tp_dict */
2871 0, /* tp_descr_get */
2872 0, /* tp_descr_set */
2873 0, /* tp_dictoffset */
2874 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002875 0, /* tp_alloc */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002876 filterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002877 PyObject_GC_Del, /* tp_free */
2878};
2879
2880
2881/* count object ************************************************************/
2882
2883typedef struct {
2884 PyObject_HEAD
Thomas Wouters477c8d52006-05-27 19:21:47 +00002885 Py_ssize_t cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002886 PyObject *long_cnt; /* Arbitrarily large count when cnt >= PY_SSIZE_T_MAX */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002887} countobject;
2888
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002889static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002890
2891static PyObject *
2892count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2893{
2894 countobject *lz;
Thomas Wouters477c8d52006-05-27 19:21:47 +00002895 Py_ssize_t cnt = 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002896 PyObject *cnt_arg = NULL;
2897 PyObject *long_cnt = NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002898
Thomas Woutersb2137042007-02-01 18:02:27 +00002899 if (type == &count_type && !_PyArg_NoKeywords("count()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002900 return NULL;
2901
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002902 if (!PyArg_UnpackTuple(args, "count", 0, 1, &cnt_arg))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002903 return NULL;
2904
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002905 if (cnt_arg != NULL) {
Christian Heimes217cfd12007-12-02 14:31:20 +00002906 cnt = PyLong_AsSsize_t(cnt_arg);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002907 if (cnt == -1 && PyErr_Occurred()) {
2908 PyErr_Clear();
2909 if (!PyLong_Check(cnt_arg)) {
2910 PyErr_SetString(PyExc_TypeError, "an integer is required");
2911 return NULL;
2912 }
2913 long_cnt = cnt_arg;
2914 Py_INCREF(long_cnt);
2915 cnt = PY_SSIZE_T_MAX;
2916 }
2917 }
2918
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002919 /* create countobject structure */
2920 lz = (countobject *)PyObject_New(countobject, &count_type);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002921 if (lz == NULL) {
2922 Py_XDECREF(long_cnt);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002923 return NULL;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002924 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002925 lz->cnt = cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002926 lz->long_cnt = long_cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002927
2928 return (PyObject *)lz;
2929}
2930
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002931static void
2932count_dealloc(countobject *lz)
2933{
2934 Py_XDECREF(lz->long_cnt);
2935 PyObject_Del(lz);
2936}
2937
2938static PyObject *
2939count_nextlong(countobject *lz)
2940{
2941 static PyObject *one = NULL;
2942 PyObject *cnt;
2943 PyObject *stepped_up;
2944
2945 if (lz->long_cnt == NULL) {
Christian Heimes217cfd12007-12-02 14:31:20 +00002946 lz->long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002947 if (lz->long_cnt == NULL)
2948 return NULL;
2949 }
2950 if (one == NULL) {
Christian Heimes217cfd12007-12-02 14:31:20 +00002951 one = PyLong_FromLong(1);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002952 if (one == NULL)
2953 return NULL;
2954 }
2955 cnt = lz->long_cnt;
2956 assert(cnt != NULL);
2957 stepped_up = PyNumber_Add(cnt, one);
2958 if (stepped_up == NULL)
2959 return NULL;
2960 lz->long_cnt = stepped_up;
2961 return cnt;
2962}
2963
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002964static PyObject *
2965count_next(countobject *lz)
2966{
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002967 if (lz->cnt == PY_SSIZE_T_MAX)
2968 return count_nextlong(lz);
Christian Heimes217cfd12007-12-02 14:31:20 +00002969 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002970}
2971
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002972static PyObject *
2973count_repr(countobject *lz)
2974{
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002975 if (lz->cnt != PY_SSIZE_T_MAX)
2976 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
2977
2978 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002979}
2980
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002981PyDoc_STRVAR(count_doc,
2982"count([firstval]) --> count object\n\
2983\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002984Return a count object whose .__next__() method returns consecutive\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002985integers starting from zero or, if specified, from firstval.");
2986
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002987static PyTypeObject count_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00002988 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002989 "itertools.count", /* tp_name */
2990 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002991 0, /* tp_itemsize */
2992 /* methods */
Guido van Rossum8ce8a782007-11-01 19:42:39 +00002993 (destructor)count_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002994 0, /* tp_print */
2995 0, /* tp_getattr */
2996 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00002997 0, /* tp_reserved */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002998 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002999 0, /* tp_as_number */
3000 0, /* tp_as_sequence */
3001 0, /* tp_as_mapping */
3002 0, /* tp_hash */
3003 0, /* tp_call */
3004 0, /* tp_str */
3005 PyObject_GenericGetAttr, /* tp_getattro */
3006 0, /* tp_setattro */
3007 0, /* tp_as_buffer */
3008 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003009 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003010 0, /* tp_traverse */
3011 0, /* tp_clear */
3012 0, /* tp_richcompare */
3013 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003014 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003015 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003016 0, /* tp_methods */
3017 0, /* tp_members */
3018 0, /* tp_getset */
3019 0, /* tp_base */
3020 0, /* tp_dict */
3021 0, /* tp_descr_get */
3022 0, /* tp_descr_set */
3023 0, /* tp_dictoffset */
3024 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003025 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003026 count_new, /* tp_new */
3027};
3028
3029
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003030/* repeat object ************************************************************/
3031
3032typedef struct {
3033 PyObject_HEAD
3034 PyObject *element;
Thomas Wouters477c8d52006-05-27 19:21:47 +00003035 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003036} repeatobject;
3037
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003038static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003039
3040static PyObject *
3041repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3042{
3043 repeatobject *ro;
3044 PyObject *element;
Thomas Wouters477c8d52006-05-27 19:21:47 +00003045 Py_ssize_t cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003046
Thomas Woutersb2137042007-02-01 18:02:27 +00003047 if (type == &repeat_type && !_PyArg_NoKeywords("repeat()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00003048 return NULL;
3049
Thomas Wouters477c8d52006-05-27 19:21:47 +00003050 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003051 return NULL;
3052
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003053 if (PyTuple_Size(args) == 2 && cnt < 0)
3054 cnt = 0;
3055
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003056 ro = (repeatobject *)type->tp_alloc(type, 0);
3057 if (ro == NULL)
3058 return NULL;
3059 Py_INCREF(element);
3060 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003061 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003062 return (PyObject *)ro;
3063}
3064
3065static void
3066repeat_dealloc(repeatobject *ro)
3067{
3068 PyObject_GC_UnTrack(ro);
3069 Py_XDECREF(ro->element);
Christian Heimes90aa7642007-12-19 02:45:37 +00003070 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003071}
3072
3073static int
3074repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3075{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00003076 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003077 return 0;
3078}
3079
3080static PyObject *
3081repeat_next(repeatobject *ro)
3082{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003083 if (ro->cnt == 0)
3084 return NULL;
3085 if (ro->cnt > 0)
3086 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003087 Py_INCREF(ro->element);
3088 return ro->element;
3089}
3090
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003091static PyObject *
3092repeat_repr(repeatobject *ro)
3093{
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003094 if (ro->cnt == -1)
Walter Dörwald7569dfe2007-05-19 21:49:49 +00003095 return PyUnicode_FromFormat("repeat(%R)", ro->element);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003096 else
Walter Dörwald7569dfe2007-05-19 21:49:49 +00003097 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003098}
3099
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003100static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003101repeat_len(repeatobject *ro)
3102{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003103 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003104 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003105 return NULL;
3106 }
Christian Heimes217cfd12007-12-02 14:31:20 +00003107 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003108}
3109
Armin Rigof5b3e362006-02-11 21:32:43 +00003110PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003111
3112static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00003113 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003114 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003115};
3116
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003117PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003118"repeat(element [,times]) -> create an iterator which returns the element\n\
3119for the specified number of times. If not specified, returns the element\n\
3120endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003121
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003122static PyTypeObject repeat_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00003123 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003124 "itertools.repeat", /* tp_name */
3125 sizeof(repeatobject), /* tp_basicsize */
3126 0, /* tp_itemsize */
3127 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003128 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003129 0, /* tp_print */
3130 0, /* tp_getattr */
3131 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00003132 0, /* tp_reserved */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003133 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003134 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003135 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003136 0, /* tp_as_mapping */
3137 0, /* tp_hash */
3138 0, /* tp_call */
3139 0, /* tp_str */
3140 PyObject_GenericGetAttr, /* tp_getattro */
3141 0, /* tp_setattro */
3142 0, /* tp_as_buffer */
3143 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3144 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003145 repeat_doc, /* tp_doc */
3146 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003147 0, /* tp_clear */
3148 0, /* tp_richcompare */
3149 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003150 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003151 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003152 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003153 0, /* tp_members */
3154 0, /* tp_getset */
3155 0, /* tp_base */
3156 0, /* tp_dict */
3157 0, /* tp_descr_get */
3158 0, /* tp_descr_set */
3159 0, /* tp_dictoffset */
3160 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003161 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003162 repeat_new, /* tp_new */
3163 PyObject_GC_Del, /* tp_free */
3164};
3165
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003166/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00003167
3168#include "Python.h"
3169
3170typedef struct {
3171 PyObject_HEAD
3172 Py_ssize_t tuplesize;
3173 Py_ssize_t numactive;
3174 PyObject *ittuple; /* tuple of iterators */
3175 PyObject *result;
3176 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003177} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003178
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003179static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003180
3181static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003182zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003183{
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003184 ziplongestobject *lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00003185 Py_ssize_t i;
3186 PyObject *ittuple; /* tuple of iterators */
3187 PyObject *result;
3188 PyObject *fillvalue = Py_None;
3189 Py_ssize_t tuplesize = PySequence_Length(args);
3190
3191 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3192 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3193 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3194 PyErr_SetString(PyExc_TypeError,
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003195 "zip_longest() got an unexpected keyword argument");
Thomas Wouterscf297e42007-02-23 15:07:44 +00003196 return NULL;
3197 }
3198 }
3199
3200 /* args must be a tuple */
3201 assert(PyTuple_Check(args));
3202
3203 /* obtain iterators */
3204 ittuple = PyTuple_New(tuplesize);
3205 if (ittuple == NULL)
3206 return NULL;
3207 for (i=0; i < tuplesize; ++i) {
3208 PyObject *item = PyTuple_GET_ITEM(args, i);
3209 PyObject *it = PyObject_GetIter(item);
3210 if (it == NULL) {
3211 if (PyErr_ExceptionMatches(PyExc_TypeError))
3212 PyErr_Format(PyExc_TypeError,
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003213 "zip_longest argument #%zd must support iteration",
Thomas Wouterscf297e42007-02-23 15:07:44 +00003214 i+1);
3215 Py_DECREF(ittuple);
3216 return NULL;
3217 }
3218 PyTuple_SET_ITEM(ittuple, i, it);
3219 }
3220
3221 /* create a result holder */
3222 result = PyTuple_New(tuplesize);
3223 if (result == NULL) {
3224 Py_DECREF(ittuple);
3225 return NULL;
3226 }
3227 for (i=0 ; i < tuplesize ; i++) {
3228 Py_INCREF(Py_None);
3229 PyTuple_SET_ITEM(result, i, Py_None);
3230 }
3231
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003232 /* create ziplongestobject structure */
3233 lz = (ziplongestobject *)type->tp_alloc(type, 0);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003234 if (lz == NULL) {
3235 Py_DECREF(ittuple);
3236 Py_DECREF(result);
3237 return NULL;
3238 }
3239 lz->ittuple = ittuple;
3240 lz->tuplesize = tuplesize;
3241 lz->numactive = tuplesize;
3242 lz->result = result;
3243 Py_INCREF(fillvalue);
3244 lz->fillvalue = fillvalue;
3245 return (PyObject *)lz;
3246}
3247
3248static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003249zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003250{
3251 PyObject_GC_UnTrack(lz);
3252 Py_XDECREF(lz->ittuple);
3253 Py_XDECREF(lz->result);
3254 Py_XDECREF(lz->fillvalue);
Christian Heimes90aa7642007-12-19 02:45:37 +00003255 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003256}
3257
3258static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003259zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003260{
3261 Py_VISIT(lz->ittuple);
3262 Py_VISIT(lz->result);
3263 Py_VISIT(lz->fillvalue);
3264 return 0;
3265}
3266
3267static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003268zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00003269{
3270 Py_ssize_t i;
3271 Py_ssize_t tuplesize = lz->tuplesize;
3272 PyObject *result = lz->result;
3273 PyObject *it;
3274 PyObject *item;
3275 PyObject *olditem;
3276
3277 if (tuplesize == 0)
3278 return NULL;
3279 if (lz->numactive == 0)
3280 return NULL;
Christian Heimes90aa7642007-12-19 02:45:37 +00003281 if (Py_REFCNT(result) == 1) {
Thomas Wouterscf297e42007-02-23 15:07:44 +00003282 Py_INCREF(result);
3283 for (i=0 ; i < tuplesize ; i++) {
3284 it = PyTuple_GET_ITEM(lz->ittuple, i);
3285 if (it == NULL) {
3286 Py_INCREF(lz->fillvalue);
3287 item = lz->fillvalue;
3288 } else {
Christian Heimes90aa7642007-12-19 02:45:37 +00003289 item = (*Py_TYPE(it)->tp_iternext)(it);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003290 if (item == NULL) {
3291 lz->numactive -= 1;
3292 if (lz->numactive == 0) {
3293 Py_DECREF(result);
3294 return NULL;
3295 } else {
3296 Py_INCREF(lz->fillvalue);
3297 item = lz->fillvalue;
3298 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3299 Py_DECREF(it);
3300 }
3301 }
3302 }
3303 olditem = PyTuple_GET_ITEM(result, i);
3304 PyTuple_SET_ITEM(result, i, item);
3305 Py_DECREF(olditem);
3306 }
3307 } else {
3308 result = PyTuple_New(tuplesize);
3309 if (result == NULL)
3310 return NULL;
3311 for (i=0 ; i < tuplesize ; i++) {
3312 it = PyTuple_GET_ITEM(lz->ittuple, i);
3313 if (it == NULL) {
3314 Py_INCREF(lz->fillvalue);
3315 item = lz->fillvalue;
3316 } else {
Christian Heimes90aa7642007-12-19 02:45:37 +00003317 item = (*Py_TYPE(it)->tp_iternext)(it);
Thomas Wouterscf297e42007-02-23 15:07:44 +00003318 if (item == NULL) {
3319 lz->numactive -= 1;
3320 if (lz->numactive == 0) {
3321 Py_DECREF(result);
3322 return NULL;
3323 } else {
3324 Py_INCREF(lz->fillvalue);
3325 item = lz->fillvalue;
3326 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3327 Py_DECREF(it);
3328 }
3329 }
3330 }
3331 PyTuple_SET_ITEM(result, i, item);
3332 }
3333 }
3334 return result;
3335}
3336
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003337PyDoc_STRVAR(zip_longest_doc,
3338"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003339\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003340Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00003341the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00003342method continues until the longest iterable in the argument sequence\n\
3343is exhausted and then it raises StopIteration. When the shorter iterables\n\
3344are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3345defaults to None or can be specified by a keyword argument.\n\
3346");
3347
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003348static PyTypeObject ziplongest_type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00003349 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003350 "itertools.zip_longest", /* tp_name */
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003351 sizeof(ziplongestobject), /* tp_basicsize */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003352 0, /* tp_itemsize */
3353 /* methods */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003354 (destructor)zip_longest_dealloc, /* tp_dealloc */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003355 0, /* tp_print */
3356 0, /* tp_getattr */
3357 0, /* tp_setattr */
Mark Dickinsone94c6792009-02-02 20:36:42 +00003358 0, /* tp_reserved */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003359 0, /* tp_repr */
3360 0, /* tp_as_number */
3361 0, /* tp_as_sequence */
3362 0, /* tp_as_mapping */
3363 0, /* tp_hash */
3364 0, /* tp_call */
3365 0, /* tp_str */
3366 PyObject_GenericGetAttr, /* tp_getattro */
3367 0, /* tp_setattro */
3368 0, /* tp_as_buffer */
3369 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3370 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003371 zip_longest_doc, /* tp_doc */
3372 (traverseproc)zip_longest_traverse, /* tp_traverse */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003373 0, /* tp_clear */
3374 0, /* tp_richcompare */
3375 0, /* tp_weaklistoffset */
3376 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003377 (iternextfunc)zip_longest_next, /* tp_iternext */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003378 0, /* tp_methods */
3379 0, /* tp_members */
3380 0, /* tp_getset */
3381 0, /* tp_base */
3382 0, /* tp_dict */
3383 0, /* tp_descr_get */
3384 0, /* tp_descr_set */
3385 0, /* tp_dictoffset */
3386 0, /* tp_init */
3387 0, /* tp_alloc */
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003388 zip_longest_new, /* tp_new */
Thomas Wouterscf297e42007-02-23 15:07:44 +00003389 PyObject_GC_Del, /* tp_free */
3390};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003391
3392/* module level code ********************************************************/
3393
3394PyDoc_STRVAR(module_doc,
3395"Functional tools for creating and using iterators.\n\
3396\n\
3397Infinite iterators:\n\
3398count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003399cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003400repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003401\n\
3402Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003403chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
3404compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
3405dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
3406groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003407filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003408islice(seq, [start,] stop [, step]) --> elements from\n\
3409 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003410starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003411tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003412takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00003413zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00003414\n\
3415Combinatoric generators:\n\
3416product(p, q, ... [repeat=1]) --> cartesian product\n\
3417permutations(p[, r])\n\
3418combinations(p[, r])\n\
3419combinations_with_replacement(p[, r])\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003420");
3421
3422
Raymond Hettingerad983e72003-11-12 14:32:26 +00003423static PyMethodDef module_methods[] = {
3424 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3425 {NULL, NULL} /* sentinel */
3426};
3427
Martin v. Löwis1a214512008-06-11 05:26:20 +00003428
3429static struct PyModuleDef itertoolsmodule = {
3430 PyModuleDef_HEAD_INIT,
3431 "itertools",
3432 module_doc,
3433 -1,
3434 module_methods,
3435 NULL,
3436 NULL,
3437 NULL,
3438 NULL
3439};
3440
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003441PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00003442PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003443{
Raymond Hettinger60eca932003-02-09 06:40:58 +00003444 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003445 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003446 char *name;
3447 PyTypeObject *typelist[] = {
Christian Heimes380f7f22008-02-28 11:19:05 +00003448 &combinations_type,
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003449 &cwr_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003450 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003451 &dropwhile_type,
3452 &takewhile_type,
3453 &islice_type,
3454 &starmap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003455 &chain_type,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003456 &compress_type,
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003457 &filterfalse_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003458 &count_type,
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00003459 &ziplongest_type,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003460 &permutations_type,
Christian Heimes380f7f22008-02-28 11:19:05 +00003461 &product_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003462 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003463 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003464 NULL
3465 };
3466
Christian Heimes90aa7642007-12-19 02:45:37 +00003467 Py_TYPE(&teedataobject_type) = &PyType_Type;
Martin v. Löwis1a214512008-06-11 05:26:20 +00003468 m = PyModule_Create(&itertoolsmodule);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00003469 if (m == NULL)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003470 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003471
Raymond Hettinger60eca932003-02-09 06:40:58 +00003472 for (i=0 ; typelist[i] != NULL ; i++) {
3473 if (PyType_Ready(typelist[i]) < 0)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003474 return NULL;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003475 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003476 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003477 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003478 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003479 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003480
3481 if (PyType_Ready(&teedataobject_type) < 0)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003482 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +00003483 if (PyType_Ready(&tee_type) < 0)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003484 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003485 if (PyType_Ready(&_grouper_type) < 0)
Martin v. Löwis1a214512008-06-11 05:26:20 +00003486 return NULL;
3487 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003488}