blob: 1c91a1959dfe05fedc77abe590c8db831d97251d [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);
62 gbo->ob_type->tp_free(gbo);
63}
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 = {
142 PyObject_HEAD_INIT(NULL)
143 0, /* ob_size */
144 "itertools.groupby", /* tp_name */
145 sizeof(groupbyobject), /* tp_basicsize */
146 0, /* tp_itemsize */
147 /* methods */
148 (destructor)groupby_dealloc, /* tp_dealloc */
149 0, /* tp_print */
150 0, /* tp_getattr */
151 0, /* tp_setattr */
152 0, /* tp_compare */
153 0, /* tp_repr */
154 0, /* tp_as_number */
155 0, /* tp_as_sequence */
156 0, /* tp_as_mapping */
157 0, /* tp_hash */
158 0, /* tp_call */
159 0, /* tp_str */
160 PyObject_GenericGetAttr, /* tp_getattro */
161 0, /* tp_setattro */
162 0, /* tp_as_buffer */
163 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
164 Py_TPFLAGS_BASETYPE, /* tp_flags */
165 groupby_doc, /* tp_doc */
166 (traverseproc)groupby_traverse, /* tp_traverse */
167 0, /* tp_clear */
168 0, /* tp_richcompare */
169 0, /* tp_weaklistoffset */
170 PyObject_SelfIter, /* tp_iter */
171 (iternextfunc)groupby_next, /* tp_iternext */
172 0, /* tp_methods */
173 0, /* tp_members */
174 0, /* tp_getset */
175 0, /* tp_base */
176 0, /* tp_dict */
177 0, /* tp_descr_get */
178 0, /* tp_descr_set */
179 0, /* tp_dictoffset */
180 0, /* tp_init */
181 0, /* tp_alloc */
182 groupby_new, /* tp_new */
183 PyObject_GC_Del, /* tp_free */
184};
185
186
187/* _grouper object (internal) ************************************************/
188
189typedef struct {
190 PyObject_HEAD
191 PyObject *parent;
192 PyObject *tgtkey;
193} _grouperobject;
194
195static PyTypeObject _grouper_type;
196
197static PyObject *
198_grouper_create(groupbyobject *parent, PyObject *tgtkey)
199{
200 _grouperobject *igo;
201
202 igo = PyObject_New(_grouperobject, &_grouper_type);
203 if (igo == NULL)
204 return NULL;
205 igo->parent = (PyObject *)parent;
206 Py_INCREF(parent);
207 igo->tgtkey = tgtkey;
208 Py_INCREF(tgtkey);
209
210 return (PyObject *)igo;
211}
212
213static void
214_grouper_dealloc(_grouperobject *igo)
215{
216 Py_DECREF(igo->parent);
217 Py_DECREF(igo->tgtkey);
218 PyObject_Del(igo);
219}
220
221static PyObject *
222_grouper_next(_grouperobject *igo)
223{
224 groupbyobject *gbo = (groupbyobject *)igo->parent;
225 PyObject *newvalue, *newkey, *r;
226 int rcmp;
227
228 if (gbo->currvalue == NULL) {
229 newvalue = PyIter_Next(gbo->it);
230 if (newvalue == NULL)
231 return NULL;
232
233 if (gbo->keyfunc == Py_None) {
234 newkey = newvalue;
235 Py_INCREF(newvalue);
236 } else {
237 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
238 newvalue, NULL);
239 if (newkey == NULL) {
240 Py_DECREF(newvalue);
241 return NULL;
242 }
243 }
244
245 assert(gbo->currkey == NULL);
246 gbo->currkey = newkey;
247 gbo->currvalue = newvalue;
248 }
249
250 assert(gbo->currkey != NULL);
251 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
252 if (rcmp <= 0)
253 /* got any error or current group is end */
254 return NULL;
255
256 r = gbo->currvalue;
257 gbo->currvalue = NULL;
Raymond Hettinger75ccea32004-09-01 07:02:44 +0000258 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000259
260 return r;
261}
262
263static PyTypeObject _grouper_type = {
264 PyObject_HEAD_INIT(NULL)
265 0, /* ob_size */
266 "itertools._grouper", /* tp_name */
267 sizeof(_grouperobject), /* tp_basicsize */
268 0, /* tp_itemsize */
269 /* methods */
270 (destructor)_grouper_dealloc, /* tp_dealloc */
271 0, /* tp_print */
272 0, /* tp_getattr */
273 0, /* tp_setattr */
274 0, /* tp_compare */
275 0, /* tp_repr */
276 0, /* tp_as_number */
277 0, /* tp_as_sequence */
278 0, /* tp_as_mapping */
279 0, /* tp_hash */
280 0, /* tp_call */
281 0, /* tp_str */
282 PyObject_GenericGetAttr, /* tp_getattro */
283 0, /* tp_setattro */
284 0, /* tp_as_buffer */
285 Py_TPFLAGS_DEFAULT, /* tp_flags */
286 0, /* tp_doc */
287 0, /* tp_traverse */
288 0, /* tp_clear */
289 0, /* tp_richcompare */
290 0, /* tp_weaklistoffset */
291 PyObject_SelfIter, /* tp_iter */
292 (iternextfunc)_grouper_next, /* tp_iternext */
293 0, /* tp_methods */
294 0, /* tp_members */
295 0, /* tp_getset */
296 0, /* tp_base */
297 0, /* tp_dict */
298 0, /* tp_descr_get */
299 0, /* tp_descr_set */
300 0, /* tp_dictoffset */
301 0, /* tp_init */
302 0, /* tp_alloc */
303 0, /* tp_new */
304 PyObject_Del, /* tp_free */
305};
306
307
308
Raymond Hettingerad983e72003-11-12 14:32:26 +0000309/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000310
Raymond Hettingerad983e72003-11-12 14:32:26 +0000311/* The teedataobject pre-allocates space for LINKCELLS number of objects.
312 To help the object fit neatly inside cache lines (space for 16 to 32
313 pointers), the value should be a multiple of 16 minus space for
314 the other structure members including PyHEAD overhead. The larger the
315 value, the less memory overhead per object and the less time spent
316 allocating/deallocating new links. The smaller the number, the less
317 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000318*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000319#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000320
321typedef struct {
322 PyObject_HEAD
323 PyObject *it;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000324 int numread;
325 PyObject *nextlink;
326 PyObject *(values[LINKCELLS]);
327} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000328
329typedef struct {
330 PyObject_HEAD
Raymond Hettingerad983e72003-11-12 14:32:26 +0000331 teedataobject *dataobj;
332 int index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000333 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000334} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000335
Raymond Hettingerad983e72003-11-12 14:32:26 +0000336static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000337
338static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000339teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000340{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000341 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000342
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000343 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000344 if (tdo == NULL)
Raymond Hettinger45143692003-10-25 06:37:47 +0000345 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000346
347 tdo->numread = 0;
348 tdo->nextlink = NULL;
349 Py_INCREF(it);
350 tdo->it = it;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000351 PyObject_GC_Track(tdo);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000352 return (PyObject *)tdo;
353}
354
355static PyObject *
356teedataobject_jumplink(teedataobject *tdo)
357{
358 if (tdo->nextlink == NULL)
359 tdo->nextlink = teedataobject_new(tdo->it);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000360 Py_XINCREF(tdo->nextlink);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000361 return tdo->nextlink;
362}
363
364static PyObject *
365teedataobject_getitem(teedataobject *tdo, int i)
366{
367 PyObject *value;
368
369 assert(i < LINKCELLS);
370 if (i < tdo->numread)
371 value = tdo->values[i];
372 else {
373 /* this is the lead iterator, so fetch more data */
374 assert(i == tdo->numread);
375 value = PyIter_Next(tdo->it);
376 if (value == NULL)
377 return NULL;
378 tdo->numread++;
379 tdo->values[i] = value;
Raymond Hettinger45143692003-10-25 06:37:47 +0000380 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000381 Py_INCREF(value);
382 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000383}
384
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000385static int
386teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
387{
388 int i;
389 Py_VISIT(tdo->it);
390 for (i = 0; i < tdo->numread; i++)
391 Py_VISIT(tdo->values[i]);
392 Py_VISIT(tdo->nextlink);
393 return 0;
394}
395
396static int
397teedataobject_clear(teedataobject *tdo)
398{
399 int i;
400 Py_CLEAR(tdo->it);
401 for (i=0 ; i<tdo->numread ; i++)
402 Py_CLEAR(tdo->values[i]);
403 Py_CLEAR(tdo->nextlink);
404 return 0;
405}
406
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000407static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000408teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000409{
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000410 PyObject_GC_UnTrack(tdo);
411 teedataobject_clear(tdo);
412 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000413}
414
Raymond Hettingerad983e72003-11-12 14:32:26 +0000415PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000416
Raymond Hettingerad983e72003-11-12 14:32:26 +0000417static PyTypeObject teedataobject_type = {
Skip Montanarof3938fd2004-02-10 20:27:40 +0000418 PyObject_HEAD_INIT(0) /* Must fill in type value later */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000419 0, /* ob_size */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000420 "itertools.tee_dataobject", /* tp_name */
421 sizeof(teedataobject), /* tp_basicsize */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000422 0, /* tp_itemsize */
423 /* methods */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000424 (destructor)teedataobject_dealloc, /* tp_dealloc */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000425 0, /* tp_print */
426 0, /* tp_getattr */
427 0, /* tp_setattr */
428 0, /* tp_compare */
429 0, /* tp_repr */
430 0, /* tp_as_number */
431 0, /* tp_as_sequence */
432 0, /* tp_as_mapping */
433 0, /* tp_hash */
434 0, /* tp_call */
435 0, /* tp_str */
436 PyObject_GenericGetAttr, /* tp_getattro */
437 0, /* tp_setattro */
438 0, /* tp_as_buffer */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000439 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000440 teedataobject_doc, /* tp_doc */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000441 (traverseproc)teedataobject_traverse, /* tp_traverse */
442 (inquiry)teedataobject_clear, /* tp_clear */
443 0, /* tp_richcompare */
444 0, /* tp_weaklistoffset */
445 0, /* tp_iter */
446 0, /* tp_iternext */
447 0, /* tp_methods */
448 0, /* tp_members */
449 0, /* tp_getset */
450 0, /* tp_base */
451 0, /* tp_dict */
452 0, /* tp_descr_get */
453 0, /* tp_descr_set */
454 0, /* tp_dictoffset */
455 0, /* tp_init */
456 0, /* tp_alloc */
457 0, /* tp_new */
458 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000459};
460
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000461
462static PyTypeObject tee_type;
463
464static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000465tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000466{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000467 PyObject *value, *link;
468
469 if (to->index >= LINKCELLS) {
470 link = teedataobject_jumplink(to->dataobj);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000471 Py_DECREF(to->dataobj);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000472 to->dataobj = (teedataobject *)link;
473 to->index = 0;
474 }
475 value = teedataobject_getitem(to->dataobj, to->index);
476 if (value == NULL)
477 return NULL;
478 to->index++;
479 return value;
480}
481
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000482static int
483tee_traverse(teeobject *to, visitproc visit, void *arg)
484{
485 Py_VISIT((PyObject *)to->dataobj);
486 return 0;
487}
488
Raymond Hettingerad983e72003-11-12 14:32:26 +0000489static PyObject *
490tee_copy(teeobject *to)
491{
492 teeobject *newto;
493
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000494 newto = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000495 if (newto == NULL)
496 return NULL;
497 Py_INCREF(to->dataobj);
498 newto->dataobj = to->dataobj;
499 newto->index = to->index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000500 newto->weakreflist = NULL;
Thomas Woutersb3deb942006-04-15 22:33:13 +0000501 PyObject_GC_Track(newto);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000502 return (PyObject *)newto;
503}
504
505PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
506
507static PyObject *
508tee_fromiterable(PyObject *iterable)
509{
510 teeobject *to;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000511 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000512
513 it = PyObject_GetIter(iterable);
514 if (it == NULL)
515 return NULL;
516 if (PyObject_TypeCheck(it, &tee_type)) {
517 to = (teeobject *)tee_copy((teeobject *)it);
518 goto done;
519 }
520
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000521 to = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000522 if (to == NULL)
523 goto done;
524 to->dataobj = (teedataobject *)teedataobject_new(it);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000525 if (!to->dataobj) {
526 PyObject_GC_Del(to);
527 to = NULL;
528 goto done;
529 }
530
Raymond Hettingerad983e72003-11-12 14:32:26 +0000531 to->index = 0;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000532 to->weakreflist = NULL;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000533 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000534done:
535 Py_XDECREF(it);
536 return (PyObject *)to;
537}
538
539static PyObject *
540tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
541{
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000542 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000543
544 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
545 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000546 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000547}
548
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000549static int
550tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000551{
Raymond Hettingera9f60922004-10-17 16:40:14 +0000552 if (to->weakreflist != NULL)
553 PyObject_ClearWeakRefs((PyObject *) to);
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000554 Py_CLEAR(to->dataobj);
555 return 0;
556}
557
558static void
559tee_dealloc(teeobject *to)
560{
561 PyObject_GC_UnTrack(to);
562 tee_clear(to);
563 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000564}
565
Raymond Hettingerad983e72003-11-12 14:32:26 +0000566PyDoc_STRVAR(teeobject_doc,
567"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000568
Raymond Hettingerad983e72003-11-12 14:32:26 +0000569static PyMethodDef tee_methods[] = {
570 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
571 {NULL, NULL} /* sentinel */
572};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000573
574static PyTypeObject tee_type = {
575 PyObject_HEAD_INIT(NULL)
576 0, /* ob_size */
577 "itertools.tee", /* tp_name */
578 sizeof(teeobject), /* tp_basicsize */
579 0, /* tp_itemsize */
580 /* methods */
581 (destructor)tee_dealloc, /* tp_dealloc */
582 0, /* tp_print */
583 0, /* tp_getattr */
584 0, /* tp_setattr */
585 0, /* tp_compare */
586 0, /* tp_repr */
587 0, /* tp_as_number */
588 0, /* tp_as_sequence */
589 0, /* tp_as_mapping */
590 0, /* tp_hash */
591 0, /* tp_call */
592 0, /* tp_str */
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000593 0, /* tp_getattro */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000594 0, /* tp_setattro */
595 0, /* tp_as_buffer */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000596 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000597 teeobject_doc, /* tp_doc */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000598 (traverseproc)tee_traverse, /* tp_traverse */
599 (inquiry)tee_clear, /* tp_clear */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000600 0, /* tp_richcompare */
Raymond Hettingera9f60922004-10-17 16:40:14 +0000601 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000602 PyObject_SelfIter, /* tp_iter */
603 (iternextfunc)tee_next, /* tp_iternext */
604 tee_methods, /* tp_methods */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000605 0, /* tp_members */
606 0, /* tp_getset */
607 0, /* tp_base */
608 0, /* tp_dict */
609 0, /* tp_descr_get */
610 0, /* tp_descr_set */
611 0, /* tp_dictoffset */
612 0, /* tp_init */
613 0, /* tp_alloc */
614 tee_new, /* tp_new */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000615 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000616};
617
Raymond Hettingerad983e72003-11-12 14:32:26 +0000618static PyObject *
619tee(PyObject *self, PyObject *args)
620{
Neal Norwitz69e88972006-09-02 02:58:13 +0000621 Py_ssize_t i, n=2;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000622 PyObject *it, *iterable, *copyable, *result;
623
Neal Norwitz69e88972006-09-02 02:58:13 +0000624 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000625 return NULL;
Neal Norwitz69e88972006-09-02 02:58:13 +0000626 if (n < 0) {
627 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
628 return NULL;
629 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000630 result = PyTuple_New(n);
631 if (result == NULL)
632 return NULL;
633 if (n == 0)
634 return result;
635 it = PyObject_GetIter(iterable);
636 if (it == NULL) {
637 Py_DECREF(result);
638 return NULL;
639 }
640 if (!PyObject_HasAttrString(it, "__copy__")) {
641 copyable = tee_fromiterable(it);
642 Py_DECREF(it);
643 if (copyable == NULL) {
644 Py_DECREF(result);
645 return NULL;
646 }
647 } else
648 copyable = it;
649 PyTuple_SET_ITEM(result, 0, copyable);
650 for (i=1 ; i<n ; i++) {
651 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
652 if (copyable == NULL) {
653 Py_DECREF(result);
654 return NULL;
655 }
656 PyTuple_SET_ITEM(result, i, copyable);
657 }
658 return result;
659}
660
661PyDoc_STRVAR(tee_doc,
662"tee(iterable, n=2) --> tuple of n independent iterators.");
663
664
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000665/* cycle object **********************************************************/
666
667typedef struct {
668 PyObject_HEAD
669 PyObject *it;
670 PyObject *saved;
671 int firstpass;
672} cycleobject;
673
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000674static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000675
676static PyObject *
677cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
678{
679 PyObject *it;
680 PyObject *iterable;
681 PyObject *saved;
682 cycleobject *lz;
683
Georg Brandlb84c1372007-01-21 10:28:43 +0000684 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000685 return NULL;
686
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000687 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
688 return NULL;
689
690 /* Get iterator. */
691 it = PyObject_GetIter(iterable);
692 if (it == NULL)
693 return NULL;
694
695 saved = PyList_New(0);
696 if (saved == NULL) {
697 Py_DECREF(it);
698 return NULL;
699 }
700
701 /* create cycleobject structure */
702 lz = (cycleobject *)type->tp_alloc(type, 0);
703 if (lz == NULL) {
704 Py_DECREF(it);
705 Py_DECREF(saved);
706 return NULL;
707 }
708 lz->it = it;
709 lz->saved = saved;
710 lz->firstpass = 0;
711
712 return (PyObject *)lz;
713}
714
715static void
716cycle_dealloc(cycleobject *lz)
717{
718 PyObject_GC_UnTrack(lz);
719 Py_XDECREF(lz->saved);
720 Py_XDECREF(lz->it);
721 lz->ob_type->tp_free(lz);
722}
723
724static int
725cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
726{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000727 Py_VISIT(lz->it);
728 Py_VISIT(lz->saved);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000729 return 0;
730}
731
732static PyObject *
733cycle_next(cycleobject *lz)
734{
735 PyObject *item;
736 PyObject *it;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000737 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000738
739 while (1) {
740 item = PyIter_Next(lz->it);
741 if (item != NULL) {
742 if (!lz->firstpass)
743 PyList_Append(lz->saved, item);
744 return item;
745 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000746 if (PyErr_Occurred()) {
747 if (PyErr_ExceptionMatches(PyExc_StopIteration))
748 PyErr_Clear();
749 else
750 return NULL;
751 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000752 if (PyList_Size(lz->saved) == 0)
753 return NULL;
754 it = PyObject_GetIter(lz->saved);
755 if (it == NULL)
756 return NULL;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000757 tmp = lz->it;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000758 lz->it = it;
759 lz->firstpass = 1;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000760 Py_DECREF(tmp);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000761 }
762}
763
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000764PyDoc_STRVAR(cycle_doc,
765"cycle(iterable) --> cycle object\n\
766\n\
767Return elements from the iterable until it is exhausted.\n\
768Then repeat the sequence indefinitely.");
769
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000770static PyTypeObject cycle_type = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000771 PyObject_HEAD_INIT(NULL)
772 0, /* ob_size */
773 "itertools.cycle", /* tp_name */
774 sizeof(cycleobject), /* tp_basicsize */
775 0, /* tp_itemsize */
776 /* methods */
777 (destructor)cycle_dealloc, /* tp_dealloc */
778 0, /* tp_print */
779 0, /* tp_getattr */
780 0, /* tp_setattr */
781 0, /* tp_compare */
782 0, /* tp_repr */
783 0, /* tp_as_number */
784 0, /* tp_as_sequence */
785 0, /* tp_as_mapping */
786 0, /* tp_hash */
787 0, /* tp_call */
788 0, /* tp_str */
789 PyObject_GenericGetAttr, /* tp_getattro */
790 0, /* tp_setattro */
791 0, /* tp_as_buffer */
792 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
793 Py_TPFLAGS_BASETYPE, /* tp_flags */
794 cycle_doc, /* tp_doc */
795 (traverseproc)cycle_traverse, /* tp_traverse */
796 0, /* tp_clear */
797 0, /* tp_richcompare */
798 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000799 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000800 (iternextfunc)cycle_next, /* tp_iternext */
801 0, /* tp_methods */
802 0, /* tp_members */
803 0, /* tp_getset */
804 0, /* tp_base */
805 0, /* tp_dict */
806 0, /* tp_descr_get */
807 0, /* tp_descr_set */
808 0, /* tp_dictoffset */
809 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000810 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000811 cycle_new, /* tp_new */
812 PyObject_GC_Del, /* tp_free */
813};
814
815
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000816/* dropwhile object **********************************************************/
817
818typedef struct {
819 PyObject_HEAD
820 PyObject *func;
821 PyObject *it;
822 long start;
823} dropwhileobject;
824
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000825static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000826
827static PyObject *
828dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
829{
830 PyObject *func, *seq;
831 PyObject *it;
832 dropwhileobject *lz;
833
Georg Brandlb84c1372007-01-21 10:28:43 +0000834 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000835 return NULL;
836
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000837 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
838 return NULL;
839
840 /* Get iterator. */
841 it = PyObject_GetIter(seq);
842 if (it == NULL)
843 return NULL;
844
845 /* create dropwhileobject structure */
846 lz = (dropwhileobject *)type->tp_alloc(type, 0);
847 if (lz == NULL) {
848 Py_DECREF(it);
849 return NULL;
850 }
851 Py_INCREF(func);
852 lz->func = func;
853 lz->it = it;
854 lz->start = 0;
855
856 return (PyObject *)lz;
857}
858
859static void
860dropwhile_dealloc(dropwhileobject *lz)
861{
862 PyObject_GC_UnTrack(lz);
863 Py_XDECREF(lz->func);
864 Py_XDECREF(lz->it);
865 lz->ob_type->tp_free(lz);
866}
867
868static int
869dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
870{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000871 Py_VISIT(lz->it);
872 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000873 return 0;
874}
875
876static PyObject *
877dropwhile_next(dropwhileobject *lz)
878{
879 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000880 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000881 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000882 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000883
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000884 assert(PyIter_Check(it));
885 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000886 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000887 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000888 if (item == NULL)
889 return NULL;
890 if (lz->start == 1)
891 return item;
892
893 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
894 if (good == NULL) {
895 Py_DECREF(item);
896 return NULL;
897 }
898 ok = PyObject_IsTrue(good);
899 Py_DECREF(good);
900 if (!ok) {
901 lz->start = 1;
902 return item;
903 }
904 Py_DECREF(item);
905 }
906}
907
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000908PyDoc_STRVAR(dropwhile_doc,
909"dropwhile(predicate, iterable) --> dropwhile object\n\
910\n\
911Drop items from the iterable while predicate(item) is true.\n\
912Afterwards, return every element until the iterable is exhausted.");
913
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000914static PyTypeObject dropwhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000915 PyObject_HEAD_INIT(NULL)
916 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000917 "itertools.dropwhile", /* tp_name */
918 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000919 0, /* tp_itemsize */
920 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000921 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000922 0, /* tp_print */
923 0, /* tp_getattr */
924 0, /* tp_setattr */
925 0, /* tp_compare */
926 0, /* tp_repr */
927 0, /* tp_as_number */
928 0, /* tp_as_sequence */
929 0, /* tp_as_mapping */
930 0, /* tp_hash */
931 0, /* tp_call */
932 0, /* tp_str */
933 PyObject_GenericGetAttr, /* tp_getattro */
934 0, /* tp_setattro */
935 0, /* tp_as_buffer */
936 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
937 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000938 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000939 (traverseproc)dropwhile_traverse, /* tp_traverse */
940 0, /* tp_clear */
941 0, /* tp_richcompare */
942 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000943 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000944 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000945 0, /* tp_methods */
946 0, /* tp_members */
947 0, /* tp_getset */
948 0, /* tp_base */
949 0, /* tp_dict */
950 0, /* tp_descr_get */
951 0, /* tp_descr_set */
952 0, /* tp_dictoffset */
953 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000954 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000955 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000956 PyObject_GC_Del, /* tp_free */
957};
958
959
960/* takewhile object **********************************************************/
961
962typedef struct {
963 PyObject_HEAD
964 PyObject *func;
965 PyObject *it;
966 long stop;
967} takewhileobject;
968
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000969static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000970
971static PyObject *
972takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
973{
974 PyObject *func, *seq;
975 PyObject *it;
976 takewhileobject *lz;
977
Georg Brandlb84c1372007-01-21 10:28:43 +0000978 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000979 return NULL;
980
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000981 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
982 return NULL;
983
984 /* Get iterator. */
985 it = PyObject_GetIter(seq);
986 if (it == NULL)
987 return NULL;
988
989 /* create takewhileobject structure */
990 lz = (takewhileobject *)type->tp_alloc(type, 0);
991 if (lz == NULL) {
992 Py_DECREF(it);
993 return NULL;
994 }
995 Py_INCREF(func);
996 lz->func = func;
997 lz->it = it;
998 lz->stop = 0;
999
1000 return (PyObject *)lz;
1001}
1002
1003static void
1004takewhile_dealloc(takewhileobject *lz)
1005{
1006 PyObject_GC_UnTrack(lz);
1007 Py_XDECREF(lz->func);
1008 Py_XDECREF(lz->it);
1009 lz->ob_type->tp_free(lz);
1010}
1011
1012static int
1013takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1014{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001015 Py_VISIT(lz->it);
1016 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001017 return 0;
1018}
1019
1020static PyObject *
1021takewhile_next(takewhileobject *lz)
1022{
1023 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001024 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001025 long ok;
1026
1027 if (lz->stop == 1)
1028 return NULL;
1029
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001030 assert(PyIter_Check(it));
1031 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001032 if (item == NULL)
1033 return NULL;
1034
1035 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1036 if (good == NULL) {
1037 Py_DECREF(item);
1038 return NULL;
1039 }
1040 ok = PyObject_IsTrue(good);
1041 Py_DECREF(good);
1042 if (ok)
1043 return item;
1044 Py_DECREF(item);
1045 lz->stop = 1;
1046 return NULL;
1047}
1048
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001049PyDoc_STRVAR(takewhile_doc,
1050"takewhile(predicate, iterable) --> takewhile object\n\
1051\n\
1052Return successive entries from an iterable as long as the \n\
1053predicate evaluates to true for each entry.");
1054
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001055static PyTypeObject takewhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001056 PyObject_HEAD_INIT(NULL)
1057 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001058 "itertools.takewhile", /* tp_name */
1059 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001060 0, /* tp_itemsize */
1061 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001062 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001063 0, /* tp_print */
1064 0, /* tp_getattr */
1065 0, /* tp_setattr */
1066 0, /* tp_compare */
1067 0, /* tp_repr */
1068 0, /* tp_as_number */
1069 0, /* tp_as_sequence */
1070 0, /* tp_as_mapping */
1071 0, /* tp_hash */
1072 0, /* tp_call */
1073 0, /* tp_str */
1074 PyObject_GenericGetAttr, /* tp_getattro */
1075 0, /* tp_setattro */
1076 0, /* tp_as_buffer */
1077 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1078 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001079 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001080 (traverseproc)takewhile_traverse, /* tp_traverse */
1081 0, /* tp_clear */
1082 0, /* tp_richcompare */
1083 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001084 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001085 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001086 0, /* tp_methods */
1087 0, /* tp_members */
1088 0, /* tp_getset */
1089 0, /* tp_base */
1090 0, /* tp_dict */
1091 0, /* tp_descr_get */
1092 0, /* tp_descr_set */
1093 0, /* tp_dictoffset */
1094 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001095 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001096 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001097 PyObject_GC_Del, /* tp_free */
1098};
1099
1100
1101/* islice object ************************************************************/
1102
1103typedef struct {
1104 PyObject_HEAD
1105 PyObject *it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001106 Py_ssize_t next;
1107 Py_ssize_t stop;
1108 Py_ssize_t step;
1109 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001110} isliceobject;
1111
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001112static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001113
1114static PyObject *
1115islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1116{
1117 PyObject *seq;
Jack Diederich6c433a92006-05-26 11:15:17 +00001118 Py_ssize_t start=0, stop=-1, step=1;
Raymond Hettingerb2594052004-12-05 09:25:51 +00001119 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001120 Py_ssize_t numargs;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001121 isliceobject *lz;
1122
Georg Brandlb84c1372007-01-21 10:28:43 +00001123 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001124 return NULL;
1125
Raymond Hettingerb2594052004-12-05 09:25:51 +00001126 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001127 return NULL;
1128
Raymond Hettingerb2594052004-12-05 09:25:51 +00001129 numargs = PyTuple_Size(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001130 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001131 if (a1 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001132 stop = PyInt_AsSsize_t(a1);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001133 if (stop == -1) {
1134 if (PyErr_Occurred())
1135 PyErr_Clear();
1136 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001137 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001138 return NULL;
1139 }
1140 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001141 } else {
Raymond Hettingerb2594052004-12-05 09:25:51 +00001142 if (a1 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001143 start = PyInt_AsSsize_t(a1);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001144 if (start == -1 && PyErr_Occurred())
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001145 PyErr_Clear();
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001146 if (a2 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001147 stop = PyInt_AsSsize_t(a2);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001148 if (stop == -1) {
1149 if (PyErr_Occurred())
1150 PyErr_Clear();
1151 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001152 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001153 return NULL;
1154 }
1155 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001156 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001157 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001158 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001159 "Indices for islice() must be non-negative integers or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001160 return NULL;
1161 }
1162
Raymond Hettingerb2594052004-12-05 09:25:51 +00001163 if (a3 != NULL) {
1164 if (a3 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001165 step = PyInt_AsSsize_t(a3);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001166 if (step == -1 && PyErr_Occurred())
1167 PyErr_Clear();
1168 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001169 if (step<1) {
1170 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001171 "Step for islice() must be a positive integer or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001172 return NULL;
1173 }
1174
1175 /* Get iterator. */
1176 it = PyObject_GetIter(seq);
1177 if (it == NULL)
1178 return NULL;
1179
1180 /* create isliceobject structure */
1181 lz = (isliceobject *)type->tp_alloc(type, 0);
1182 if (lz == NULL) {
1183 Py_DECREF(it);
1184 return NULL;
1185 }
1186 lz->it = it;
1187 lz->next = start;
1188 lz->stop = stop;
1189 lz->step = step;
1190 lz->cnt = 0L;
1191
1192 return (PyObject *)lz;
1193}
1194
1195static void
1196islice_dealloc(isliceobject *lz)
1197{
1198 PyObject_GC_UnTrack(lz);
1199 Py_XDECREF(lz->it);
1200 lz->ob_type->tp_free(lz);
1201}
1202
1203static int
1204islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1205{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001206 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001207 return 0;
1208}
1209
1210static PyObject *
1211islice_next(isliceobject *lz)
1212{
1213 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001214 PyObject *it = lz->it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001215 Py_ssize_t oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001216 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001217
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001218 assert(PyIter_Check(it));
1219 iternext = *it->ob_type->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 Hettingerd1a283b2003-03-01 01:48:24 +00001229 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001230 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001231 if (item == NULL)
1232 return NULL;
1233 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001234 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001235 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001236 if (lz->next < oldnext) /* Check for overflow */
1237 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001238 return item;
1239}
1240
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001241PyDoc_STRVAR(islice_doc,
1242"islice(iterable, [start,] stop [, step]) --> islice object\n\
1243\n\
1244Return an iterator whose next() method returns selected values from an\n\
1245iterable. If start is specified, will skip all preceding elements;\n\
1246otherwise, start defaults to zero. Step defaults to one. If\n\
1247specified as another value, step determines how many values are \n\
1248skipped between successive calls. Works like a slice() on a list\n\
1249but returns an iterator.");
1250
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001251static PyTypeObject islice_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001252 PyObject_HEAD_INIT(NULL)
1253 0, /* ob_size */
1254 "itertools.islice", /* tp_name */
1255 sizeof(isliceobject), /* tp_basicsize */
1256 0, /* tp_itemsize */
1257 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001258 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001259 0, /* tp_print */
1260 0, /* tp_getattr */
1261 0, /* tp_setattr */
1262 0, /* tp_compare */
1263 0, /* tp_repr */
1264 0, /* tp_as_number */
1265 0, /* tp_as_sequence */
1266 0, /* tp_as_mapping */
1267 0, /* tp_hash */
1268 0, /* tp_call */
1269 0, /* tp_str */
1270 PyObject_GenericGetAttr, /* tp_getattro */
1271 0, /* tp_setattro */
1272 0, /* tp_as_buffer */
1273 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1274 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001275 islice_doc, /* tp_doc */
1276 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001277 0, /* tp_clear */
1278 0, /* tp_richcompare */
1279 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001280 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001281 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001282 0, /* tp_methods */
1283 0, /* tp_members */
1284 0, /* tp_getset */
1285 0, /* tp_base */
1286 0, /* tp_dict */
1287 0, /* tp_descr_get */
1288 0, /* tp_descr_set */
1289 0, /* tp_dictoffset */
1290 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001291 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001292 islice_new, /* tp_new */
1293 PyObject_GC_Del, /* tp_free */
1294};
1295
1296
1297/* starmap object ************************************************************/
1298
1299typedef struct {
1300 PyObject_HEAD
1301 PyObject *func;
1302 PyObject *it;
1303} starmapobject;
1304
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001305static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001306
1307static PyObject *
1308starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1309{
1310 PyObject *func, *seq;
1311 PyObject *it;
1312 starmapobject *lz;
1313
Georg Brandlb84c1372007-01-21 10:28:43 +00001314 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001315 return NULL;
1316
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001317 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1318 return NULL;
1319
1320 /* Get iterator. */
1321 it = PyObject_GetIter(seq);
1322 if (it == NULL)
1323 return NULL;
1324
1325 /* create starmapobject structure */
1326 lz = (starmapobject *)type->tp_alloc(type, 0);
1327 if (lz == NULL) {
1328 Py_DECREF(it);
1329 return NULL;
1330 }
1331 Py_INCREF(func);
1332 lz->func = func;
1333 lz->it = it;
1334
1335 return (PyObject *)lz;
1336}
1337
1338static void
1339starmap_dealloc(starmapobject *lz)
1340{
1341 PyObject_GC_UnTrack(lz);
1342 Py_XDECREF(lz->func);
1343 Py_XDECREF(lz->it);
1344 lz->ob_type->tp_free(lz);
1345}
1346
1347static int
1348starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1349{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001350 Py_VISIT(lz->it);
1351 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001352 return 0;
1353}
1354
1355static PyObject *
1356starmap_next(starmapobject *lz)
1357{
1358 PyObject *args;
1359 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001360 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001361
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001362 assert(PyIter_Check(it));
1363 args = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001364 if (args == NULL)
1365 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001366 if (!PyTuple_CheckExact(args)) {
1367 Py_DECREF(args);
1368 PyErr_SetString(PyExc_TypeError,
1369 "iterator must return a tuple");
1370 return NULL;
1371 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001372 result = PyObject_Call(lz->func, args, NULL);
1373 Py_DECREF(args);
1374 return result;
1375}
1376
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001377PyDoc_STRVAR(starmap_doc,
1378"starmap(function, sequence) --> starmap object\n\
1379\n\
1380Return an iterator whose values are returned from the function evaluated\n\
1381with a argument tuple taken from the given sequence.");
1382
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001383static PyTypeObject starmap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001384 PyObject_HEAD_INIT(NULL)
1385 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001386 "itertools.starmap", /* tp_name */
1387 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001388 0, /* tp_itemsize */
1389 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001390 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001391 0, /* tp_print */
1392 0, /* tp_getattr */
1393 0, /* tp_setattr */
1394 0, /* tp_compare */
1395 0, /* tp_repr */
1396 0, /* tp_as_number */
1397 0, /* tp_as_sequence */
1398 0, /* tp_as_mapping */
1399 0, /* tp_hash */
1400 0, /* tp_call */
1401 0, /* tp_str */
1402 PyObject_GenericGetAttr, /* tp_getattro */
1403 0, /* tp_setattro */
1404 0, /* tp_as_buffer */
1405 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1406 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001407 starmap_doc, /* tp_doc */
1408 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001409 0, /* tp_clear */
1410 0, /* tp_richcompare */
1411 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001412 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001413 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001414 0, /* tp_methods */
1415 0, /* tp_members */
1416 0, /* tp_getset */
1417 0, /* tp_base */
1418 0, /* tp_dict */
1419 0, /* tp_descr_get */
1420 0, /* tp_descr_set */
1421 0, /* tp_dictoffset */
1422 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001423 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001424 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001425 PyObject_GC_Del, /* tp_free */
1426};
1427
1428
1429/* imap object ************************************************************/
1430
1431typedef struct {
1432 PyObject_HEAD
1433 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001434 PyObject *func;
1435} imapobject;
1436
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001437static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001438
1439static PyObject *
1440imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1441{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001442 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001443 imapobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001444 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001445
Georg Brandlb84c1372007-01-21 10:28:43 +00001446 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001447 return NULL;
1448
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001449 numargs = PyTuple_Size(args);
1450 if (numargs < 2) {
1451 PyErr_SetString(PyExc_TypeError,
1452 "imap() must have at least two arguments.");
1453 return NULL;
1454 }
1455
1456 iters = PyTuple_New(numargs-1);
1457 if (iters == NULL)
1458 return NULL;
1459
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001460 for (i=1 ; i<numargs ; i++) {
1461 /* Get iterator. */
1462 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1463 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001464 Py_DECREF(iters);
1465 return NULL;
1466 }
1467 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001468 }
1469
1470 /* create imapobject structure */
1471 lz = (imapobject *)type->tp_alloc(type, 0);
1472 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001473 Py_DECREF(iters);
1474 return NULL;
1475 }
1476 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001477 func = PyTuple_GET_ITEM(args, 0);
1478 Py_INCREF(func);
1479 lz->func = func;
1480
1481 return (PyObject *)lz;
1482}
1483
1484static void
1485imap_dealloc(imapobject *lz)
1486{
1487 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001488 Py_XDECREF(lz->iters);
1489 Py_XDECREF(lz->func);
1490 lz->ob_type->tp_free(lz);
1491}
1492
1493static int
1494imap_traverse(imapobject *lz, visitproc visit, void *arg)
1495{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001496 Py_VISIT(lz->iters);
1497 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001498 return 0;
1499}
1500
Raymond Hettinger2012f172003-02-07 05:32:58 +00001501/*
1502imap() is an iterator version of __builtins__.map() except that it does
1503not have the None fill-in feature. That was intentionally left out for
1504the following reasons:
1505
1506 1) Itertools are designed to be easily combined and chained together.
1507 Having all tools stop with the shortest input is a unifying principle
1508 that makes it easier to combine finite iterators (supplying data) with
1509 infinite iterators like count() and repeat() (for supplying sequential
1510 or constant arguments to a function).
1511
1512 2) In typical use cases for combining itertools, having one finite data
1513 supplier run out before another is likely to be an error condition which
1514 should not pass silently by automatically supplying None.
1515
1516 3) The use cases for automatic None fill-in are rare -- not many functions
1517 do something useful when a parameter suddenly switches type and becomes
1518 None.
1519
1520 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001521 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001522
1523 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1524*/
1525
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001526static PyObject *
1527imap_next(imapobject *lz)
1528{
1529 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001530 PyObject *argtuple;
1531 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001532 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001533
1534 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001535 argtuple = PyTuple_New(numargs);
1536 if (argtuple == NULL)
1537 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001538
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001539 for (i=0 ; i<numargs ; i++) {
1540 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1541 if (val == NULL) {
1542 Py_DECREF(argtuple);
1543 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001544 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001545 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001546 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001547 if (lz->func == Py_None)
1548 return argtuple;
1549 result = PyObject_Call(lz->func, argtuple, NULL);
1550 Py_DECREF(argtuple);
1551 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001552}
1553
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001554PyDoc_STRVAR(imap_doc,
1555"imap(func, *iterables) --> imap object\n\
1556\n\
1557Make an iterator that computes the function using arguments from\n\
1558each of the iterables. Like map() except that it returns\n\
1559an iterator instead of a list and that it stops when the shortest\n\
1560iterable is exhausted instead of filling in None for shorter\n\
1561iterables.");
1562
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001563static PyTypeObject imap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001564 PyObject_HEAD_INIT(NULL)
1565 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001566 "itertools.imap", /* tp_name */
1567 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001568 0, /* tp_itemsize */
1569 /* methods */
1570 (destructor)imap_dealloc, /* tp_dealloc */
1571 0, /* tp_print */
1572 0, /* tp_getattr */
1573 0, /* tp_setattr */
1574 0, /* tp_compare */
1575 0, /* tp_repr */
1576 0, /* tp_as_number */
1577 0, /* tp_as_sequence */
1578 0, /* tp_as_mapping */
1579 0, /* tp_hash */
1580 0, /* tp_call */
1581 0, /* tp_str */
1582 PyObject_GenericGetAttr, /* tp_getattro */
1583 0, /* tp_setattro */
1584 0, /* tp_as_buffer */
1585 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1586 Py_TPFLAGS_BASETYPE, /* tp_flags */
1587 imap_doc, /* tp_doc */
1588 (traverseproc)imap_traverse, /* tp_traverse */
1589 0, /* tp_clear */
1590 0, /* tp_richcompare */
1591 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001592 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001593 (iternextfunc)imap_next, /* tp_iternext */
1594 0, /* tp_methods */
1595 0, /* tp_members */
1596 0, /* tp_getset */
1597 0, /* tp_base */
1598 0, /* tp_dict */
1599 0, /* tp_descr_get */
1600 0, /* tp_descr_set */
1601 0, /* tp_dictoffset */
1602 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001603 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001604 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001605 PyObject_GC_Del, /* tp_free */
1606};
1607
1608
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001609/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001610
1611typedef struct {
1612 PyObject_HEAD
Jack Diederich6c433a92006-05-26 11:15:17 +00001613 Py_ssize_t tuplesize;
1614 Py_ssize_t iternum; /* which iterator is active */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001615 PyObject *ittuple; /* tuple of iterators */
1616} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001617
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001618static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001619
1620static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001621chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001622{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001623 chainobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001624 Py_ssize_t tuplesize = PySequence_Length(args);
Jack Diederich6c433a92006-05-26 11:15:17 +00001625 Py_ssize_t i;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001626 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001627
Georg Brandlb84c1372007-01-21 10:28:43 +00001628 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001629 return NULL;
1630
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001631 /* obtain iterators */
1632 assert(PyTuple_Check(args));
1633 ittuple = PyTuple_New(tuplesize);
Neal Norwitz2f3136b2006-05-27 05:18:57 +00001634 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001635 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001636 for (i=0; i < tuplesize; ++i) {
1637 PyObject *item = PyTuple_GET_ITEM(args, i);
1638 PyObject *it = PyObject_GetIter(item);
1639 if (it == NULL) {
1640 if (PyErr_ExceptionMatches(PyExc_TypeError))
1641 PyErr_Format(PyExc_TypeError,
Neal Norwitz2f3136b2006-05-27 05:18:57 +00001642 "chain argument #%zd must support iteration",
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001643 i+1);
1644 Py_DECREF(ittuple);
1645 return NULL;
1646 }
1647 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001648 }
1649
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001650 /* create chainobject structure */
1651 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001652 if (lz == NULL) {
1653 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001654 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001655 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001656
1657 lz->ittuple = ittuple;
1658 lz->iternum = 0;
1659 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001660
1661 return (PyObject *)lz;
1662}
1663
1664static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001665chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001666{
1667 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001668 Py_XDECREF(lz->ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001669 lz->ob_type->tp_free(lz);
1670}
1671
Raymond Hettinger2012f172003-02-07 05:32:58 +00001672static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001673chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001674{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001675 Py_VISIT(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001676 return 0;
1677}
1678
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001679static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001680chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001681{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001682 PyObject *it;
1683 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001684
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001685 while (lz->iternum < lz->tuplesize) {
1686 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1687 item = PyIter_Next(it);
1688 if (item != NULL)
1689 return item;
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001690 if (PyErr_Occurred()) {
1691 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1692 PyErr_Clear();
1693 else
1694 return NULL;
1695 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001696 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001697 }
1698 return NULL;
1699}
1700
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001701PyDoc_STRVAR(chain_doc,
1702"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001703\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001704Return a chain object whose .next() method returns elements from the\n\
1705first iterable until it is exhausted, then elements from the next\n\
1706iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001707
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001708static PyTypeObject chain_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001709 PyObject_HEAD_INIT(NULL)
1710 0, /* ob_size */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001711 "itertools.chain", /* tp_name */
1712 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001713 0, /* tp_itemsize */
1714 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001715 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001716 0, /* tp_print */
1717 0, /* tp_getattr */
1718 0, /* tp_setattr */
1719 0, /* tp_compare */
1720 0, /* tp_repr */
1721 0, /* tp_as_number */
1722 0, /* tp_as_sequence */
1723 0, /* tp_as_mapping */
1724 0, /* tp_hash */
1725 0, /* tp_call */
1726 0, /* tp_str */
1727 PyObject_GenericGetAttr, /* tp_getattro */
1728 0, /* tp_setattro */
1729 0, /* tp_as_buffer */
1730 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1731 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001732 chain_doc, /* tp_doc */
1733 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001734 0, /* tp_clear */
1735 0, /* tp_richcompare */
1736 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001737 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001738 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001739 0, /* tp_methods */
1740 0, /* tp_members */
1741 0, /* tp_getset */
1742 0, /* tp_base */
1743 0, /* tp_dict */
1744 0, /* tp_descr_get */
1745 0, /* tp_descr_set */
1746 0, /* tp_dictoffset */
1747 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001748 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001749 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001750 PyObject_GC_Del, /* tp_free */
1751};
1752
1753
1754/* ifilter object ************************************************************/
1755
1756typedef struct {
1757 PyObject_HEAD
1758 PyObject *func;
1759 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001760} ifilterobject;
1761
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001762static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001763
1764static PyObject *
1765ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1766{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001767 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001768 PyObject *it;
1769 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001770
Georg Brandlb84c1372007-01-21 10:28:43 +00001771 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001772 return NULL;
1773
Raymond Hettinger60eca932003-02-09 06:40:58 +00001774 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001775 return NULL;
1776
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001777 /* Get iterator. */
1778 it = PyObject_GetIter(seq);
1779 if (it == NULL)
1780 return NULL;
1781
1782 /* create ifilterobject structure */
1783 lz = (ifilterobject *)type->tp_alloc(type, 0);
1784 if (lz == NULL) {
1785 Py_DECREF(it);
1786 return NULL;
1787 }
1788 Py_INCREF(func);
1789 lz->func = func;
1790 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001791
1792 return (PyObject *)lz;
1793}
1794
1795static void
1796ifilter_dealloc(ifilterobject *lz)
1797{
1798 PyObject_GC_UnTrack(lz);
1799 Py_XDECREF(lz->func);
1800 Py_XDECREF(lz->it);
1801 lz->ob_type->tp_free(lz);
1802}
1803
1804static int
1805ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1806{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001807 Py_VISIT(lz->it);
1808 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001809 return 0;
1810}
1811
1812static PyObject *
1813ifilter_next(ifilterobject *lz)
1814{
1815 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001816 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001817 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001818 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001819
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001820 assert(PyIter_Check(it));
1821 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001822 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001823 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001824 if (item == NULL)
1825 return NULL;
1826
1827 if (lz->func == Py_None) {
1828 ok = PyObject_IsTrue(item);
1829 } else {
1830 PyObject *good;
1831 good = PyObject_CallFunctionObjArgs(lz->func,
1832 item, NULL);
1833 if (good == NULL) {
1834 Py_DECREF(item);
1835 return NULL;
1836 }
1837 ok = PyObject_IsTrue(good);
1838 Py_DECREF(good);
1839 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00001840 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001841 return item;
1842 Py_DECREF(item);
1843 }
1844}
1845
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001846PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001847"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001848\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001849Return those items of sequence for which function(item) is true.\n\
1850If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001851
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001852static PyTypeObject ifilter_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001853 PyObject_HEAD_INIT(NULL)
1854 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001855 "itertools.ifilter", /* tp_name */
1856 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001857 0, /* tp_itemsize */
1858 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001859 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001860 0, /* tp_print */
1861 0, /* tp_getattr */
1862 0, /* tp_setattr */
1863 0, /* tp_compare */
1864 0, /* tp_repr */
1865 0, /* tp_as_number */
1866 0, /* tp_as_sequence */
1867 0, /* tp_as_mapping */
1868 0, /* tp_hash */
1869 0, /* tp_call */
1870 0, /* tp_str */
1871 PyObject_GenericGetAttr, /* tp_getattro */
1872 0, /* tp_setattro */
1873 0, /* tp_as_buffer */
1874 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1875 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001876 ifilter_doc, /* tp_doc */
1877 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001878 0, /* tp_clear */
1879 0, /* tp_richcompare */
1880 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001881 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001882 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001883 0, /* tp_methods */
1884 0, /* tp_members */
1885 0, /* tp_getset */
1886 0, /* tp_base */
1887 0, /* tp_dict */
1888 0, /* tp_descr_get */
1889 0, /* tp_descr_set */
1890 0, /* tp_dictoffset */
1891 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001892 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001893 ifilter_new, /* tp_new */
1894 PyObject_GC_Del, /* tp_free */
1895};
1896
1897
1898/* ifilterfalse object ************************************************************/
1899
1900typedef struct {
1901 PyObject_HEAD
1902 PyObject *func;
1903 PyObject *it;
1904} ifilterfalseobject;
1905
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001906static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001907
1908static PyObject *
1909ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1910{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00001911 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001912 PyObject *it;
1913 ifilterfalseobject *lz;
1914
Georg Brandlb84c1372007-01-21 10:28:43 +00001915 if (type == &ifilterfalse_type &&
1916 !_PyArg_NoKeywords("ifilterfalse()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001917 return NULL;
1918
Raymond Hettinger60eca932003-02-09 06:40:58 +00001919 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1920 return NULL;
1921
1922 /* Get iterator. */
1923 it = PyObject_GetIter(seq);
1924 if (it == NULL)
1925 return NULL;
1926
1927 /* create ifilterfalseobject structure */
1928 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1929 if (lz == NULL) {
1930 Py_DECREF(it);
1931 return NULL;
1932 }
1933 Py_INCREF(func);
1934 lz->func = func;
1935 lz->it = it;
1936
1937 return (PyObject *)lz;
1938}
1939
1940static void
1941ifilterfalse_dealloc(ifilterfalseobject *lz)
1942{
1943 PyObject_GC_UnTrack(lz);
1944 Py_XDECREF(lz->func);
1945 Py_XDECREF(lz->it);
1946 lz->ob_type->tp_free(lz);
1947}
1948
1949static int
1950ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1951{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001952 Py_VISIT(lz->it);
1953 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001954 return 0;
1955}
1956
1957static PyObject *
1958ifilterfalse_next(ifilterfalseobject *lz)
1959{
1960 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001961 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001962 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001963 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001964
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001965 assert(PyIter_Check(it));
1966 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001967 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001968 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001969 if (item == NULL)
1970 return NULL;
1971
1972 if (lz->func == Py_None) {
1973 ok = PyObject_IsTrue(item);
1974 } else {
1975 PyObject *good;
1976 good = PyObject_CallFunctionObjArgs(lz->func,
1977 item, NULL);
1978 if (good == NULL) {
1979 Py_DECREF(item);
1980 return NULL;
1981 }
1982 ok = PyObject_IsTrue(good);
1983 Py_DECREF(good);
1984 }
1985 if (!ok)
1986 return item;
1987 Py_DECREF(item);
1988 }
1989}
1990
Raymond Hettinger60eca932003-02-09 06:40:58 +00001991PyDoc_STRVAR(ifilterfalse_doc,
1992"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1993\n\
1994Return those items of sequence for which function(item) is false.\n\
1995If function is None, return the items that are false.");
1996
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001997static PyTypeObject ifilterfalse_type = {
Raymond Hettinger60eca932003-02-09 06:40:58 +00001998 PyObject_HEAD_INIT(NULL)
1999 0, /* ob_size */
2000 "itertools.ifilterfalse", /* tp_name */
2001 sizeof(ifilterfalseobject), /* tp_basicsize */
2002 0, /* tp_itemsize */
2003 /* methods */
2004 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
2005 0, /* tp_print */
2006 0, /* tp_getattr */
2007 0, /* tp_setattr */
2008 0, /* tp_compare */
2009 0, /* tp_repr */
2010 0, /* tp_as_number */
2011 0, /* tp_as_sequence */
2012 0, /* tp_as_mapping */
2013 0, /* tp_hash */
2014 0, /* tp_call */
2015 0, /* tp_str */
2016 PyObject_GenericGetAttr, /* tp_getattro */
2017 0, /* tp_setattro */
2018 0, /* tp_as_buffer */
2019 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2020 Py_TPFLAGS_BASETYPE, /* tp_flags */
2021 ifilterfalse_doc, /* tp_doc */
2022 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
2023 0, /* tp_clear */
2024 0, /* tp_richcompare */
2025 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002026 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002027 (iternextfunc)ifilterfalse_next, /* tp_iternext */
2028 0, /* tp_methods */
2029 0, /* tp_members */
2030 0, /* tp_getset */
2031 0, /* tp_base */
2032 0, /* tp_dict */
2033 0, /* tp_descr_get */
2034 0, /* tp_descr_set */
2035 0, /* tp_dictoffset */
2036 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002037 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002038 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002039 PyObject_GC_Del, /* tp_free */
2040};
2041
2042
2043/* count object ************************************************************/
2044
2045typedef struct {
2046 PyObject_HEAD
Jack Diederich6c433a92006-05-26 11:15:17 +00002047 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002048} countobject;
2049
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002050static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002051
2052static PyObject *
2053count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2054{
2055 countobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00002056 Py_ssize_t cnt = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002057
Georg Brandlb84c1372007-01-21 10:28:43 +00002058 if (type == &count_type && !_PyArg_NoKeywords("count()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002059 return NULL;
2060
Jack Diederich6c433a92006-05-26 11:15:17 +00002061 if (!PyArg_ParseTuple(args, "|n:count", &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002062 return NULL;
2063
2064 /* create countobject structure */
2065 lz = (countobject *)PyObject_New(countobject, &count_type);
2066 if (lz == NULL)
2067 return NULL;
2068 lz->cnt = cnt;
2069
2070 return (PyObject *)lz;
2071}
2072
2073static PyObject *
2074count_next(countobject *lz)
2075{
Raymond Hettinger6d121f12007-02-08 00:07:32 +00002076 if (lz->cnt == LONG_MAX) {
2077 PyErr_SetString(PyExc_OverflowError,
2078 "cannot count beyond LONG_MAX");
2079 return NULL;
2080 }
Jack Diederich36234e82006-09-21 17:50:26 +00002081 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002082}
2083
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002084static PyObject *
2085count_repr(countobject *lz)
2086{
Jack Diederich6c433a92006-05-26 11:15:17 +00002087 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002088}
2089
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002090PyDoc_STRVAR(count_doc,
2091"count([firstval]) --> count object\n\
2092\n\
2093Return a count object whose .next() method returns consecutive\n\
2094integers starting from zero or, if specified, from firstval.");
2095
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002096static PyTypeObject count_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002097 PyObject_HEAD_INIT(NULL)
2098 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002099 "itertools.count", /* tp_name */
2100 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002101 0, /* tp_itemsize */
2102 /* methods */
2103 (destructor)PyObject_Del, /* tp_dealloc */
2104 0, /* tp_print */
2105 0, /* tp_getattr */
2106 0, /* tp_setattr */
2107 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002108 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002109 0, /* tp_as_number */
2110 0, /* tp_as_sequence */
2111 0, /* tp_as_mapping */
2112 0, /* tp_hash */
2113 0, /* tp_call */
2114 0, /* tp_str */
2115 PyObject_GenericGetAttr, /* tp_getattro */
2116 0, /* tp_setattro */
2117 0, /* tp_as_buffer */
2118 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002119 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002120 0, /* tp_traverse */
2121 0, /* tp_clear */
2122 0, /* tp_richcompare */
2123 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002124 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002125 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002126 0, /* tp_methods */
2127 0, /* tp_members */
2128 0, /* tp_getset */
2129 0, /* tp_base */
2130 0, /* tp_dict */
2131 0, /* tp_descr_get */
2132 0, /* tp_descr_set */
2133 0, /* tp_dictoffset */
2134 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002135 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002136 count_new, /* tp_new */
2137};
2138
2139
2140/* izip object ************************************************************/
2141
2142#include "Python.h"
2143
2144typedef struct {
2145 PyObject_HEAD
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002146 Py_ssize_t tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002147 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002148 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002149} izipobject;
2150
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002151static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002152
2153static PyObject *
2154izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2155{
2156 izipobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00002157 Py_ssize_t i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002158 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002159 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002160 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002161
Georg Brandlb84c1372007-01-21 10:28:43 +00002162 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002163 return NULL;
2164
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002165 /* args must be a tuple */
2166 assert(PyTuple_Check(args));
2167
2168 /* obtain iterators */
2169 ittuple = PyTuple_New(tuplesize);
Neal Norwitz2f3136b2006-05-27 05:18:57 +00002170 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002171 return NULL;
2172 for (i=0; i < tuplesize; ++i) {
2173 PyObject *item = PyTuple_GET_ITEM(args, i);
2174 PyObject *it = PyObject_GetIter(item);
2175 if (it == NULL) {
2176 if (PyErr_ExceptionMatches(PyExc_TypeError))
2177 PyErr_Format(PyExc_TypeError,
Neal Norwitz2f3136b2006-05-27 05:18:57 +00002178 "izip argument #%zd must support iteration",
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002179 i+1);
2180 Py_DECREF(ittuple);
2181 return NULL;
2182 }
2183 PyTuple_SET_ITEM(ittuple, i, it);
2184 }
2185
Raymond Hettinger2012f172003-02-07 05:32:58 +00002186 /* create a result holder */
2187 result = PyTuple_New(tuplesize);
2188 if (result == NULL) {
2189 Py_DECREF(ittuple);
2190 return NULL;
2191 }
2192 for (i=0 ; i < tuplesize ; i++) {
2193 Py_INCREF(Py_None);
2194 PyTuple_SET_ITEM(result, i, Py_None);
2195 }
2196
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002197 /* create izipobject structure */
2198 lz = (izipobject *)type->tp_alloc(type, 0);
2199 if (lz == NULL) {
2200 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002201 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002202 return NULL;
2203 }
2204 lz->ittuple = ittuple;
2205 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002206 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002207
2208 return (PyObject *)lz;
2209}
2210
2211static void
2212izip_dealloc(izipobject *lz)
2213{
2214 PyObject_GC_UnTrack(lz);
2215 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002216 Py_XDECREF(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002217 lz->ob_type->tp_free(lz);
2218}
2219
2220static int
2221izip_traverse(izipobject *lz, visitproc visit, void *arg)
2222{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002223 Py_VISIT(lz->ittuple);
2224 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002225 return 0;
2226}
2227
2228static PyObject *
2229izip_next(izipobject *lz)
2230{
Jack Diederich6c433a92006-05-26 11:15:17 +00002231 Py_ssize_t i;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002232 Py_ssize_t tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002233 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002234 PyObject *it;
2235 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002236 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002237
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002238 if (tuplesize == 0)
2239 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002240 if (result->ob_refcnt == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002241 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002242 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002243 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002244 assert(PyIter_Check(it));
2245 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002246 if (item == NULL) {
2247 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002248 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002249 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002250 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002251 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002252 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002253 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002254 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002255 result = PyTuple_New(tuplesize);
2256 if (result == NULL)
2257 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002258 for (i=0 ; i < tuplesize ; i++) {
2259 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002260 assert(PyIter_Check(it));
2261 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002262 if (item == NULL) {
2263 Py_DECREF(result);
2264 return NULL;
2265 }
2266 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002267 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002268 }
2269 return result;
2270}
2271
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002272PyDoc_STRVAR(izip_doc,
2273"izip(iter1 [,iter2 [...]]) --> izip object\n\
2274\n\
2275Return a izip object whose .next() method returns a tuple where\n\
2276the i-th element comes from the i-th iterable argument. The .next()\n\
2277method continues until the shortest iterable in the argument sequence\n\
2278is exhausted and then it raises StopIteration. Works like the zip()\n\
2279function but consumes less memory by returning an iterator instead of\n\
2280a list.");
2281
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002282static PyTypeObject izip_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002283 PyObject_HEAD_INIT(NULL)
2284 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002285 "itertools.izip", /* tp_name */
2286 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002287 0, /* tp_itemsize */
2288 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002289 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002290 0, /* tp_print */
2291 0, /* tp_getattr */
2292 0, /* tp_setattr */
2293 0, /* tp_compare */
2294 0, /* tp_repr */
2295 0, /* tp_as_number */
2296 0, /* tp_as_sequence */
2297 0, /* tp_as_mapping */
2298 0, /* tp_hash */
2299 0, /* tp_call */
2300 0, /* tp_str */
2301 PyObject_GenericGetAttr, /* tp_getattro */
2302 0, /* tp_setattro */
2303 0, /* tp_as_buffer */
2304 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2305 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002306 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002307 (traverseproc)izip_traverse, /* tp_traverse */
2308 0, /* tp_clear */
2309 0, /* tp_richcompare */
2310 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002311 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002312 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002313 0, /* tp_methods */
2314 0, /* tp_members */
2315 0, /* tp_getset */
2316 0, /* tp_base */
2317 0, /* tp_dict */
2318 0, /* tp_descr_get */
2319 0, /* tp_descr_set */
2320 0, /* tp_dictoffset */
2321 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002322 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002323 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002324 PyObject_GC_Del, /* tp_free */
2325};
2326
2327
2328/* repeat object ************************************************************/
2329
2330typedef struct {
2331 PyObject_HEAD
2332 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00002333 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002334} repeatobject;
2335
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002336static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002337
2338static PyObject *
2339repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2340{
2341 repeatobject *ro;
2342 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00002343 Py_ssize_t cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002344
Georg Brandlb84c1372007-01-21 10:28:43 +00002345 if (type == &repeat_type && !_PyArg_NoKeywords("repeat()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002346 return NULL;
2347
Jack Diederich6c433a92006-05-26 11:15:17 +00002348 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002349 return NULL;
2350
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002351 if (PyTuple_Size(args) == 2 && cnt < 0)
2352 cnt = 0;
2353
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002354 ro = (repeatobject *)type->tp_alloc(type, 0);
2355 if (ro == NULL)
2356 return NULL;
2357 Py_INCREF(element);
2358 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002359 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002360 return (PyObject *)ro;
2361}
2362
2363static void
2364repeat_dealloc(repeatobject *ro)
2365{
2366 PyObject_GC_UnTrack(ro);
2367 Py_XDECREF(ro->element);
2368 ro->ob_type->tp_free(ro);
2369}
2370
2371static int
2372repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2373{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002374 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002375 return 0;
2376}
2377
2378static PyObject *
2379repeat_next(repeatobject *ro)
2380{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002381 if (ro->cnt == 0)
2382 return NULL;
2383 if (ro->cnt > 0)
2384 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002385 Py_INCREF(ro->element);
2386 return ro->element;
2387}
2388
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002389static PyObject *
2390repeat_repr(repeatobject *ro)
2391{
2392 PyObject *result, *objrepr;
2393
2394 objrepr = PyObject_Repr(ro->element);
2395 if (objrepr == NULL)
2396 return NULL;
2397
2398 if (ro->cnt == -1)
2399 result = PyString_FromFormat("repeat(%s)",
2400 PyString_AS_STRING(objrepr));
2401 else
Jack Diederich6c433a92006-05-26 11:15:17 +00002402 result = PyString_FromFormat("repeat(%s, %zd)",
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002403 PyString_AS_STRING(objrepr), ro->cnt);
2404 Py_DECREF(objrepr);
2405 return result;
2406}
2407
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002408static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002409repeat_len(repeatobject *ro)
2410{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002411 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002412 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002413 return NULL;
2414 }
Jack Diederich6c433a92006-05-26 11:15:17 +00002415 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002416}
2417
Armin Rigof5b3e362006-02-11 21:32:43 +00002418PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002419
2420static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00002421 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002422 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002423};
2424
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002425PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002426"repeat(element [,times]) -> create an iterator which returns the element\n\
2427for the specified number of times. If not specified, returns the element\n\
2428endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002429
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002430static PyTypeObject repeat_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002431 PyObject_HEAD_INIT(NULL)
2432 0, /* ob_size */
2433 "itertools.repeat", /* tp_name */
2434 sizeof(repeatobject), /* tp_basicsize */
2435 0, /* tp_itemsize */
2436 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002437 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002438 0, /* tp_print */
2439 0, /* tp_getattr */
2440 0, /* tp_setattr */
2441 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002442 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002443 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002444 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002445 0, /* tp_as_mapping */
2446 0, /* tp_hash */
2447 0, /* tp_call */
2448 0, /* tp_str */
2449 PyObject_GenericGetAttr, /* tp_getattro */
2450 0, /* tp_setattro */
2451 0, /* tp_as_buffer */
2452 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2453 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002454 repeat_doc, /* tp_doc */
2455 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002456 0, /* tp_clear */
2457 0, /* tp_richcompare */
2458 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002459 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002460 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002461 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002462 0, /* tp_members */
2463 0, /* tp_getset */
2464 0, /* tp_base */
2465 0, /* tp_dict */
2466 0, /* tp_descr_get */
2467 0, /* tp_descr_set */
2468 0, /* tp_dictoffset */
2469 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002470 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002471 repeat_new, /* tp_new */
2472 PyObject_GC_Del, /* tp_free */
2473};
2474
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002475/* iziplongest object ************************************************************/
2476
2477#include "Python.h"
2478
2479typedef struct {
2480 PyObject_HEAD
2481 Py_ssize_t tuplesize;
2482 Py_ssize_t numactive;
2483 PyObject *ittuple; /* tuple of iterators */
2484 PyObject *result;
2485 PyObject *fillvalue;
2486 PyObject *filler; /* repeat(fillvalue) */
2487} iziplongestobject;
2488
2489static PyTypeObject iziplongest_type;
2490
2491static PyObject *
2492izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2493{
2494 iziplongestobject *lz;
2495 Py_ssize_t i;
2496 PyObject *ittuple; /* tuple of iterators */
2497 PyObject *result;
2498 PyObject *fillvalue = Py_None;
2499 PyObject *filler;
2500 Py_ssize_t tuplesize = PySequence_Length(args);
2501
2502 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
2503 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
2504 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
2505 PyErr_SetString(PyExc_TypeError,
2506 "izip_longest() got an unexpected keyword argument");
2507 return NULL;
2508 }
2509 }
2510
2511 /* args must be a tuple */
2512 assert(PyTuple_Check(args));
2513
2514 /* obtain iterators */
2515 ittuple = PyTuple_New(tuplesize);
2516 if (ittuple == NULL)
2517 return NULL;
2518 for (i=0; i < tuplesize; ++i) {
2519 PyObject *item = PyTuple_GET_ITEM(args, i);
2520 PyObject *it = PyObject_GetIter(item);
2521 if (it == NULL) {
2522 if (PyErr_ExceptionMatches(PyExc_TypeError))
2523 PyErr_Format(PyExc_TypeError,
2524 "izip_longest argument #%zd must support iteration",
2525 i+1);
2526 Py_DECREF(ittuple);
2527 return NULL;
2528 }
2529 PyTuple_SET_ITEM(ittuple, i, it);
2530 }
2531
2532 filler = PyObject_CallFunctionObjArgs((PyObject *)(&repeat_type), fillvalue, NULL);
2533 if (filler == NULL) {
2534 Py_DECREF(ittuple);
2535 return NULL;
2536 }
2537
2538 /* create a result holder */
2539 result = PyTuple_New(tuplesize);
2540 if (result == NULL) {
2541 Py_DECREF(ittuple);
2542 Py_DECREF(filler);
2543 return NULL;
2544 }
2545 for (i=0 ; i < tuplesize ; i++) {
2546 Py_INCREF(Py_None);
2547 PyTuple_SET_ITEM(result, i, Py_None);
2548 }
2549
2550 /* create iziplongestobject structure */
2551 lz = (iziplongestobject *)type->tp_alloc(type, 0);
2552 if (lz == NULL) {
2553 Py_DECREF(ittuple);
2554 Py_DECREF(filler);
2555 Py_DECREF(result);
2556 return NULL;
2557 }
2558 lz->ittuple = ittuple;
2559 lz->tuplesize = tuplesize;
2560 lz->numactive = tuplesize;
2561 lz->result = result;
2562 Py_INCREF(fillvalue);
2563 lz->fillvalue = fillvalue;
2564 Py_INCREF(filler);
2565 lz->filler = filler; /* XXX */
2566 return (PyObject *)lz;
2567}
2568
2569static void
2570izip_longest_dealloc(iziplongestobject *lz)
2571{
2572 PyObject_GC_UnTrack(lz);
2573 Py_XDECREF(lz->ittuple);
2574 Py_XDECREF(lz->result);
2575 Py_XDECREF(lz->fillvalue);
2576 Py_XDECREF(lz->filler);
2577 lz->ob_type->tp_free(lz);
2578}
2579
2580static int
2581izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
2582{
2583 Py_VISIT(lz->ittuple);
2584 Py_VISIT(lz->result);
2585 Py_VISIT(lz->fillvalue);
2586 Py_VISIT(lz->filler);
2587 return 0;
2588}
2589
2590static PyObject *
2591izip_longest_next(iziplongestobject *lz)
2592{
2593 Py_ssize_t i;
2594 Py_ssize_t tuplesize = lz->tuplesize;
2595 PyObject *result = lz->result;
2596 PyObject *it;
2597 PyObject *item;
2598 PyObject *olditem;
2599
2600 if (tuplesize == 0)
2601 return NULL;
2602 if (result->ob_refcnt == 1) {
2603 Py_INCREF(result);
2604 for (i=0 ; i < tuplesize ; i++) {
2605 it = PyTuple_GET_ITEM(lz->ittuple, i);
2606 assert(PyIter_Check(it));
2607 item = (*it->ob_type->tp_iternext)(it);
2608 if (item == NULL) {
2609 if (lz->numactive <= 1) {
2610 Py_DECREF(result);
2611 return NULL;
2612 } else {
2613 Py_INCREF(lz->filler);
2614 PyTuple_SET_ITEM(lz->ittuple, i, lz->filler);
2615 Py_INCREF(lz->fillvalue);
2616 item = lz->fillvalue;
2617 Py_DECREF(it);
2618 lz->numactive -= 1;
2619 }
2620 }
2621 olditem = PyTuple_GET_ITEM(result, i);
2622 PyTuple_SET_ITEM(result, i, item);
2623 Py_DECREF(olditem);
2624 }
2625 } else {
2626 result = PyTuple_New(tuplesize);
2627 if (result == NULL)
2628 return NULL;
2629 for (i=0 ; i < tuplesize ; i++) {
2630 it = PyTuple_GET_ITEM(lz->ittuple, i);
2631 assert(PyIter_Check(it));
2632 item = (*it->ob_type->tp_iternext)(it);
2633 if (item == NULL) {
2634 if (lz->numactive <= 1) {
2635 Py_DECREF(result);
2636 return NULL;
2637 } else {
2638 Py_INCREF(lz->filler);
2639 PyTuple_SET_ITEM(lz->ittuple, i, lz->filler);
2640 Py_INCREF(lz->fillvalue);
2641 item = lz->fillvalue;
2642 Py_DECREF(it);
2643 lz->numactive -= 1;
2644 }
2645 }
2646 PyTuple_SET_ITEM(result, i, item);
2647 }
2648 }
2649 return result;
2650}
2651
2652PyDoc_STRVAR(izip_longest_doc,
2653"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
2654\n\
2655Return an izip_longest object whose .next() method returns a tuple where\n\
2656the i-th element comes from the i-th iterable argument. The .next()\n\
2657method continues until the longest iterable in the argument sequence\n\
2658is exhausted and then it raises StopIteration. When the shorter iterables\n\
2659are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
2660defaults to None or can be specified by a keyword argument.\n\
2661");
2662
2663static PyTypeObject iziplongest_type = {
2664 PyObject_HEAD_INIT(NULL)
2665 0, /* ob_size */
2666 "itertools.izip_longest", /* tp_name */
2667 sizeof(iziplongestobject), /* tp_basicsize */
2668 0, /* tp_itemsize */
2669 /* methods */
2670 (destructor)izip_longest_dealloc, /* tp_dealloc */
2671 0, /* tp_print */
2672 0, /* tp_getattr */
2673 0, /* tp_setattr */
2674 0, /* tp_compare */
2675 0, /* tp_repr */
2676 0, /* tp_as_number */
2677 0, /* tp_as_sequence */
2678 0, /* tp_as_mapping */
2679 0, /* tp_hash */
2680 0, /* tp_call */
2681 0, /* tp_str */
2682 PyObject_GenericGetAttr, /* tp_getattro */
2683 0, /* tp_setattro */
2684 0, /* tp_as_buffer */
2685 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2686 Py_TPFLAGS_BASETYPE, /* tp_flags */
2687 izip_longest_doc, /* tp_doc */
2688 (traverseproc)izip_longest_traverse, /* tp_traverse */
2689 0, /* tp_clear */
2690 0, /* tp_richcompare */
2691 0, /* tp_weaklistoffset */
2692 PyObject_SelfIter, /* tp_iter */
2693 (iternextfunc)izip_longest_next, /* tp_iternext */
2694 0, /* tp_methods */
2695 0, /* tp_members */
2696 0, /* tp_getset */
2697 0, /* tp_base */
2698 0, /* tp_dict */
2699 0, /* tp_descr_get */
2700 0, /* tp_descr_set */
2701 0, /* tp_dictoffset */
2702 0, /* tp_init */
2703 0, /* tp_alloc */
2704 izip_longest_new, /* tp_new */
2705 PyObject_GC_Del, /* tp_free */
2706};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002707
2708/* module level code ********************************************************/
2709
2710PyDoc_STRVAR(module_doc,
2711"Functional tools for creating and using iterators.\n\
2712\n\
2713Infinite iterators:\n\
2714count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002715cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002716repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002717\n\
2718Iterators terminating on the shortest input sequence:\n\
2719izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002720izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002721ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2722ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002723islice(seq, [start,] stop [, step]) --> elements from\n\
2724 seq[start:stop:step]\n\
2725imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2726starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00002727tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002728chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002729takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2730dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002731groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002732");
2733
2734
Raymond Hettingerad983e72003-11-12 14:32:26 +00002735static PyMethodDef module_methods[] = {
2736 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2737 {NULL, NULL} /* sentinel */
2738};
2739
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002740PyMODINIT_FUNC
2741inititertools(void)
2742{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002743 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002744 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002745 char *name;
2746 PyTypeObject *typelist[] = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002747 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002748 &dropwhile_type,
2749 &takewhile_type,
2750 &islice_type,
2751 &starmap_type,
2752 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002753 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002754 &ifilter_type,
2755 &ifilterfalse_type,
2756 &count_type,
2757 &izip_type,
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002758 &iziplongest_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002759 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002760 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002761 NULL
2762 };
2763
Skip Montanarof3938fd2004-02-10 20:27:40 +00002764 teedataobject_type.ob_type = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00002765 m = Py_InitModule3("itertools", module_methods, module_doc);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00002766 if (m == NULL)
2767 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002768
Raymond Hettinger60eca932003-02-09 06:40:58 +00002769 for (i=0 ; typelist[i] != NULL ; i++) {
2770 if (PyType_Ready(typelist[i]) < 0)
2771 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002772 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002773 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002774 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002775 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002776 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00002777
2778 if (PyType_Ready(&teedataobject_type) < 0)
2779 return;
2780 if (PyType_Ready(&tee_type) < 0)
2781 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002782 if (PyType_Ready(&_grouper_type) < 0)
2783 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002784}