blob: d91389033ff4ec719105716ef4dafc8b0dbbd16d [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{
621 int i, n=2;
622 PyObject *it, *iterable, *copyable, *result;
623
624 if (!PyArg_ParseTuple(args, "O|i", &iterable, &n))
625 return NULL;
626 result = PyTuple_New(n);
627 if (result == NULL)
628 return NULL;
629 if (n == 0)
630 return result;
631 it = PyObject_GetIter(iterable);
632 if (it == NULL) {
633 Py_DECREF(result);
634 return NULL;
635 }
636 if (!PyObject_HasAttrString(it, "__copy__")) {
637 copyable = tee_fromiterable(it);
638 Py_DECREF(it);
639 if (copyable == NULL) {
640 Py_DECREF(result);
641 return NULL;
642 }
643 } else
644 copyable = it;
645 PyTuple_SET_ITEM(result, 0, copyable);
646 for (i=1 ; i<n ; i++) {
647 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
648 if (copyable == NULL) {
649 Py_DECREF(result);
650 return NULL;
651 }
652 PyTuple_SET_ITEM(result, i, copyable);
653 }
654 return result;
655}
656
657PyDoc_STRVAR(tee_doc,
658"tee(iterable, n=2) --> tuple of n independent iterators.");
659
660
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000661/* cycle object **********************************************************/
662
663typedef struct {
664 PyObject_HEAD
665 PyObject *it;
666 PyObject *saved;
667 int firstpass;
668} cycleobject;
669
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000670static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000671
672static PyObject *
673cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
674{
675 PyObject *it;
676 PyObject *iterable;
677 PyObject *saved;
678 cycleobject *lz;
679
Georg Brandl02c42872005-08-26 06:42:30 +0000680 if (!_PyArg_NoKeywords("cycle()", kwds))
681 return NULL;
682
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000683 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
684 return NULL;
685
686 /* Get iterator. */
687 it = PyObject_GetIter(iterable);
688 if (it == NULL)
689 return NULL;
690
691 saved = PyList_New(0);
692 if (saved == NULL) {
693 Py_DECREF(it);
694 return NULL;
695 }
696
697 /* create cycleobject structure */
698 lz = (cycleobject *)type->tp_alloc(type, 0);
699 if (lz == NULL) {
700 Py_DECREF(it);
701 Py_DECREF(saved);
702 return NULL;
703 }
704 lz->it = it;
705 lz->saved = saved;
706 lz->firstpass = 0;
707
708 return (PyObject *)lz;
709}
710
711static void
712cycle_dealloc(cycleobject *lz)
713{
714 PyObject_GC_UnTrack(lz);
715 Py_XDECREF(lz->saved);
716 Py_XDECREF(lz->it);
717 lz->ob_type->tp_free(lz);
718}
719
720static int
721cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
722{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000723 Py_VISIT(lz->it);
724 Py_VISIT(lz->saved);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000725 return 0;
726}
727
728static PyObject *
729cycle_next(cycleobject *lz)
730{
731 PyObject *item;
732 PyObject *it;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000733 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000734
735 while (1) {
736 item = PyIter_Next(lz->it);
737 if (item != NULL) {
738 if (!lz->firstpass)
739 PyList_Append(lz->saved, item);
740 return item;
741 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000742 if (PyErr_Occurred()) {
743 if (PyErr_ExceptionMatches(PyExc_StopIteration))
744 PyErr_Clear();
745 else
746 return NULL;
747 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000748 if (PyList_Size(lz->saved) == 0)
749 return NULL;
750 it = PyObject_GetIter(lz->saved);
751 if (it == NULL)
752 return NULL;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000753 tmp = lz->it;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000754 lz->it = it;
755 lz->firstpass = 1;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000756 Py_DECREF(tmp);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000757 }
758}
759
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000760PyDoc_STRVAR(cycle_doc,
761"cycle(iterable) --> cycle object\n\
762\n\
763Return elements from the iterable until it is exhausted.\n\
764Then repeat the sequence indefinitely.");
765
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000766static PyTypeObject cycle_type = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000767 PyObject_HEAD_INIT(NULL)
768 0, /* ob_size */
769 "itertools.cycle", /* tp_name */
770 sizeof(cycleobject), /* tp_basicsize */
771 0, /* tp_itemsize */
772 /* methods */
773 (destructor)cycle_dealloc, /* tp_dealloc */
774 0, /* tp_print */
775 0, /* tp_getattr */
776 0, /* tp_setattr */
777 0, /* tp_compare */
778 0, /* tp_repr */
779 0, /* tp_as_number */
780 0, /* tp_as_sequence */
781 0, /* tp_as_mapping */
782 0, /* tp_hash */
783 0, /* tp_call */
784 0, /* tp_str */
785 PyObject_GenericGetAttr, /* tp_getattro */
786 0, /* tp_setattro */
787 0, /* tp_as_buffer */
788 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
789 Py_TPFLAGS_BASETYPE, /* tp_flags */
790 cycle_doc, /* tp_doc */
791 (traverseproc)cycle_traverse, /* tp_traverse */
792 0, /* tp_clear */
793 0, /* tp_richcompare */
794 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000795 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000796 (iternextfunc)cycle_next, /* tp_iternext */
797 0, /* tp_methods */
798 0, /* tp_members */
799 0, /* tp_getset */
800 0, /* tp_base */
801 0, /* tp_dict */
802 0, /* tp_descr_get */
803 0, /* tp_descr_set */
804 0, /* tp_dictoffset */
805 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000806 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000807 cycle_new, /* tp_new */
808 PyObject_GC_Del, /* tp_free */
809};
810
811
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000812/* dropwhile object **********************************************************/
813
814typedef struct {
815 PyObject_HEAD
816 PyObject *func;
817 PyObject *it;
818 long start;
819} dropwhileobject;
820
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000821static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000822
823static PyObject *
824dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
825{
826 PyObject *func, *seq;
827 PyObject *it;
828 dropwhileobject *lz;
829
Georg Brandl02c42872005-08-26 06:42:30 +0000830 if (!_PyArg_NoKeywords("dropwhile()", kwds))
831 return NULL;
832
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000833 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
834 return NULL;
835
836 /* Get iterator. */
837 it = PyObject_GetIter(seq);
838 if (it == NULL)
839 return NULL;
840
841 /* create dropwhileobject structure */
842 lz = (dropwhileobject *)type->tp_alloc(type, 0);
843 if (lz == NULL) {
844 Py_DECREF(it);
845 return NULL;
846 }
847 Py_INCREF(func);
848 lz->func = func;
849 lz->it = it;
850 lz->start = 0;
851
852 return (PyObject *)lz;
853}
854
855static void
856dropwhile_dealloc(dropwhileobject *lz)
857{
858 PyObject_GC_UnTrack(lz);
859 Py_XDECREF(lz->func);
860 Py_XDECREF(lz->it);
861 lz->ob_type->tp_free(lz);
862}
863
864static int
865dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
866{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000867 Py_VISIT(lz->it);
868 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000869 return 0;
870}
871
872static PyObject *
873dropwhile_next(dropwhileobject *lz)
874{
875 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000876 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000877 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000878 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000879
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000880 assert(PyIter_Check(it));
881 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000882 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000883 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000884 if (item == NULL)
885 return NULL;
886 if (lz->start == 1)
887 return item;
888
889 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
890 if (good == NULL) {
891 Py_DECREF(item);
892 return NULL;
893 }
894 ok = PyObject_IsTrue(good);
895 Py_DECREF(good);
896 if (!ok) {
897 lz->start = 1;
898 return item;
899 }
900 Py_DECREF(item);
901 }
902}
903
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000904PyDoc_STRVAR(dropwhile_doc,
905"dropwhile(predicate, iterable) --> dropwhile object\n\
906\n\
907Drop items from the iterable while predicate(item) is true.\n\
908Afterwards, return every element until the iterable is exhausted.");
909
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000910static PyTypeObject dropwhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000911 PyObject_HEAD_INIT(NULL)
912 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000913 "itertools.dropwhile", /* tp_name */
914 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000915 0, /* tp_itemsize */
916 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000917 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000918 0, /* tp_print */
919 0, /* tp_getattr */
920 0, /* tp_setattr */
921 0, /* tp_compare */
922 0, /* tp_repr */
923 0, /* tp_as_number */
924 0, /* tp_as_sequence */
925 0, /* tp_as_mapping */
926 0, /* tp_hash */
927 0, /* tp_call */
928 0, /* tp_str */
929 PyObject_GenericGetAttr, /* tp_getattro */
930 0, /* tp_setattro */
931 0, /* tp_as_buffer */
932 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
933 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000934 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000935 (traverseproc)dropwhile_traverse, /* tp_traverse */
936 0, /* tp_clear */
937 0, /* tp_richcompare */
938 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000939 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000940 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000941 0, /* tp_methods */
942 0, /* tp_members */
943 0, /* tp_getset */
944 0, /* tp_base */
945 0, /* tp_dict */
946 0, /* tp_descr_get */
947 0, /* tp_descr_set */
948 0, /* tp_dictoffset */
949 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000950 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000951 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000952 PyObject_GC_Del, /* tp_free */
953};
954
955
956/* takewhile object **********************************************************/
957
958typedef struct {
959 PyObject_HEAD
960 PyObject *func;
961 PyObject *it;
962 long stop;
963} takewhileobject;
964
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000965static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000966
967static PyObject *
968takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
969{
970 PyObject *func, *seq;
971 PyObject *it;
972 takewhileobject *lz;
973
Georg Brandl02c42872005-08-26 06:42:30 +0000974 if (!_PyArg_NoKeywords("takewhile()", kwds))
975 return NULL;
976
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000977 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
978 return NULL;
979
980 /* Get iterator. */
981 it = PyObject_GetIter(seq);
982 if (it == NULL)
983 return NULL;
984
985 /* create takewhileobject structure */
986 lz = (takewhileobject *)type->tp_alloc(type, 0);
987 if (lz == NULL) {
988 Py_DECREF(it);
989 return NULL;
990 }
991 Py_INCREF(func);
992 lz->func = func;
993 lz->it = it;
994 lz->stop = 0;
995
996 return (PyObject *)lz;
997}
998
999static void
1000takewhile_dealloc(takewhileobject *lz)
1001{
1002 PyObject_GC_UnTrack(lz);
1003 Py_XDECREF(lz->func);
1004 Py_XDECREF(lz->it);
1005 lz->ob_type->tp_free(lz);
1006}
1007
1008static int
1009takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1010{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001011 Py_VISIT(lz->it);
1012 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001013 return 0;
1014}
1015
1016static PyObject *
1017takewhile_next(takewhileobject *lz)
1018{
1019 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001020 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001021 long ok;
1022
1023 if (lz->stop == 1)
1024 return NULL;
1025
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001026 assert(PyIter_Check(it));
1027 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001028 if (item == NULL)
1029 return NULL;
1030
1031 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1032 if (good == NULL) {
1033 Py_DECREF(item);
1034 return NULL;
1035 }
1036 ok = PyObject_IsTrue(good);
1037 Py_DECREF(good);
1038 if (ok)
1039 return item;
1040 Py_DECREF(item);
1041 lz->stop = 1;
1042 return NULL;
1043}
1044
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001045PyDoc_STRVAR(takewhile_doc,
1046"takewhile(predicate, iterable) --> takewhile object\n\
1047\n\
1048Return successive entries from an iterable as long as the \n\
1049predicate evaluates to true for each entry.");
1050
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001051static PyTypeObject takewhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001052 PyObject_HEAD_INIT(NULL)
1053 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001054 "itertools.takewhile", /* tp_name */
1055 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001056 0, /* tp_itemsize */
1057 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001058 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001059 0, /* tp_print */
1060 0, /* tp_getattr */
1061 0, /* tp_setattr */
1062 0, /* tp_compare */
1063 0, /* tp_repr */
1064 0, /* tp_as_number */
1065 0, /* tp_as_sequence */
1066 0, /* tp_as_mapping */
1067 0, /* tp_hash */
1068 0, /* tp_call */
1069 0, /* tp_str */
1070 PyObject_GenericGetAttr, /* tp_getattro */
1071 0, /* tp_setattro */
1072 0, /* tp_as_buffer */
1073 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1074 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001075 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001076 (traverseproc)takewhile_traverse, /* tp_traverse */
1077 0, /* tp_clear */
1078 0, /* tp_richcompare */
1079 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001080 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001081 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001082 0, /* tp_methods */
1083 0, /* tp_members */
1084 0, /* tp_getset */
1085 0, /* tp_base */
1086 0, /* tp_dict */
1087 0, /* tp_descr_get */
1088 0, /* tp_descr_set */
1089 0, /* tp_dictoffset */
1090 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001091 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001092 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001093 PyObject_GC_Del, /* tp_free */
1094};
1095
1096
1097/* islice object ************************************************************/
1098
1099typedef struct {
1100 PyObject_HEAD
1101 PyObject *it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001102 Py_ssize_t next;
1103 Py_ssize_t stop;
1104 Py_ssize_t step;
1105 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001106} isliceobject;
1107
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001108static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001109
1110static PyObject *
1111islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1112{
1113 PyObject *seq;
Jack Diederich6c433a92006-05-26 11:15:17 +00001114 Py_ssize_t start=0, stop=-1, step=1;
Raymond Hettingerb2594052004-12-05 09:25:51 +00001115 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001116 Py_ssize_t numargs;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001117 isliceobject *lz;
1118
Georg Brandl02c42872005-08-26 06:42:30 +00001119 if (!_PyArg_NoKeywords("islice()", kwds))
1120 return NULL;
1121
Raymond Hettingerb2594052004-12-05 09:25:51 +00001122 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001123 return NULL;
1124
Raymond Hettingerb2594052004-12-05 09:25:51 +00001125 numargs = PyTuple_Size(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001126 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001127 if (a1 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001128 stop = PyInt_AsSsize_t(a1);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001129 if (stop == -1) {
1130 if (PyErr_Occurred())
1131 PyErr_Clear();
1132 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001133 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001134 return NULL;
1135 }
1136 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001137 } else {
Raymond Hettingerb2594052004-12-05 09:25:51 +00001138 if (a1 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001139 start = PyInt_AsSsize_t(a1);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001140 if (start == -1 && PyErr_Occurred())
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001141 PyErr_Clear();
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001142 if (a2 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001143 stop = PyInt_AsSsize_t(a2);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001144 if (stop == -1) {
1145 if (PyErr_Occurred())
1146 PyErr_Clear();
1147 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001148 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001149 return NULL;
1150 }
1151 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001152 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001153 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001154 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001155 "Indices for islice() must be non-negative integers or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001156 return NULL;
1157 }
1158
Raymond Hettingerb2594052004-12-05 09:25:51 +00001159 if (a3 != NULL) {
1160 if (a3 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001161 step = PyInt_AsSsize_t(a3);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001162 if (step == -1 && PyErr_Occurred())
1163 PyErr_Clear();
1164 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001165 if (step<1) {
1166 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001167 "Step for islice() must be a positive integer or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001168 return NULL;
1169 }
1170
1171 /* Get iterator. */
1172 it = PyObject_GetIter(seq);
1173 if (it == NULL)
1174 return NULL;
1175
1176 /* create isliceobject structure */
1177 lz = (isliceobject *)type->tp_alloc(type, 0);
1178 if (lz == NULL) {
1179 Py_DECREF(it);
1180 return NULL;
1181 }
1182 lz->it = it;
1183 lz->next = start;
1184 lz->stop = stop;
1185 lz->step = step;
1186 lz->cnt = 0L;
1187
1188 return (PyObject *)lz;
1189}
1190
1191static void
1192islice_dealloc(isliceobject *lz)
1193{
1194 PyObject_GC_UnTrack(lz);
1195 Py_XDECREF(lz->it);
1196 lz->ob_type->tp_free(lz);
1197}
1198
1199static int
1200islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1201{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001202 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001203 return 0;
1204}
1205
1206static PyObject *
1207islice_next(isliceobject *lz)
1208{
1209 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001210 PyObject *it = lz->it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001211 Py_ssize_t oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001212 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001213
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001214 assert(PyIter_Check(it));
1215 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001216 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001217 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001218 if (item == NULL)
1219 return NULL;
1220 Py_DECREF(item);
1221 lz->cnt++;
1222 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001223 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001224 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001225 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001226 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001227 if (item == NULL)
1228 return NULL;
1229 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001230 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001231 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001232 if (lz->next < oldnext) /* Check for overflow */
1233 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001234 return item;
1235}
1236
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001237PyDoc_STRVAR(islice_doc,
1238"islice(iterable, [start,] stop [, step]) --> islice object\n\
1239\n\
1240Return an iterator whose next() method returns selected values from an\n\
1241iterable. If start is specified, will skip all preceding elements;\n\
1242otherwise, start defaults to zero. Step defaults to one. If\n\
1243specified as another value, step determines how many values are \n\
1244skipped between successive calls. Works like a slice() on a list\n\
1245but returns an iterator.");
1246
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001247static PyTypeObject islice_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001248 PyObject_HEAD_INIT(NULL)
1249 0, /* ob_size */
1250 "itertools.islice", /* tp_name */
1251 sizeof(isliceobject), /* tp_basicsize */
1252 0, /* tp_itemsize */
1253 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001254 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001255 0, /* tp_print */
1256 0, /* tp_getattr */
1257 0, /* tp_setattr */
1258 0, /* tp_compare */
1259 0, /* tp_repr */
1260 0, /* tp_as_number */
1261 0, /* tp_as_sequence */
1262 0, /* tp_as_mapping */
1263 0, /* tp_hash */
1264 0, /* tp_call */
1265 0, /* tp_str */
1266 PyObject_GenericGetAttr, /* tp_getattro */
1267 0, /* tp_setattro */
1268 0, /* tp_as_buffer */
1269 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1270 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001271 islice_doc, /* tp_doc */
1272 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001273 0, /* tp_clear */
1274 0, /* tp_richcompare */
1275 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001276 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001277 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001278 0, /* tp_methods */
1279 0, /* tp_members */
1280 0, /* tp_getset */
1281 0, /* tp_base */
1282 0, /* tp_dict */
1283 0, /* tp_descr_get */
1284 0, /* tp_descr_set */
1285 0, /* tp_dictoffset */
1286 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001287 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001288 islice_new, /* tp_new */
1289 PyObject_GC_Del, /* tp_free */
1290};
1291
1292
1293/* starmap object ************************************************************/
1294
1295typedef struct {
1296 PyObject_HEAD
1297 PyObject *func;
1298 PyObject *it;
1299} starmapobject;
1300
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001301static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001302
1303static PyObject *
1304starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1305{
1306 PyObject *func, *seq;
1307 PyObject *it;
1308 starmapobject *lz;
1309
Georg Brandl02c42872005-08-26 06:42:30 +00001310 if (!_PyArg_NoKeywords("starmap()", kwds))
1311 return NULL;
1312
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001313 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1314 return NULL;
1315
1316 /* Get iterator. */
1317 it = PyObject_GetIter(seq);
1318 if (it == NULL)
1319 return NULL;
1320
1321 /* create starmapobject structure */
1322 lz = (starmapobject *)type->tp_alloc(type, 0);
1323 if (lz == NULL) {
1324 Py_DECREF(it);
1325 return NULL;
1326 }
1327 Py_INCREF(func);
1328 lz->func = func;
1329 lz->it = it;
1330
1331 return (PyObject *)lz;
1332}
1333
1334static void
1335starmap_dealloc(starmapobject *lz)
1336{
1337 PyObject_GC_UnTrack(lz);
1338 Py_XDECREF(lz->func);
1339 Py_XDECREF(lz->it);
1340 lz->ob_type->tp_free(lz);
1341}
1342
1343static int
1344starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1345{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001346 Py_VISIT(lz->it);
1347 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001348 return 0;
1349}
1350
1351static PyObject *
1352starmap_next(starmapobject *lz)
1353{
1354 PyObject *args;
1355 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001356 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001357
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001358 assert(PyIter_Check(it));
1359 args = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001360 if (args == NULL)
1361 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001362 if (!PyTuple_CheckExact(args)) {
1363 Py_DECREF(args);
1364 PyErr_SetString(PyExc_TypeError,
1365 "iterator must return a tuple");
1366 return NULL;
1367 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001368 result = PyObject_Call(lz->func, args, NULL);
1369 Py_DECREF(args);
1370 return result;
1371}
1372
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001373PyDoc_STRVAR(starmap_doc,
1374"starmap(function, sequence) --> starmap object\n\
1375\n\
1376Return an iterator whose values are returned from the function evaluated\n\
1377with a argument tuple taken from the given sequence.");
1378
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001379static PyTypeObject starmap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001380 PyObject_HEAD_INIT(NULL)
1381 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001382 "itertools.starmap", /* tp_name */
1383 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001384 0, /* tp_itemsize */
1385 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001386 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001387 0, /* tp_print */
1388 0, /* tp_getattr */
1389 0, /* tp_setattr */
1390 0, /* tp_compare */
1391 0, /* tp_repr */
1392 0, /* tp_as_number */
1393 0, /* tp_as_sequence */
1394 0, /* tp_as_mapping */
1395 0, /* tp_hash */
1396 0, /* tp_call */
1397 0, /* tp_str */
1398 PyObject_GenericGetAttr, /* tp_getattro */
1399 0, /* tp_setattro */
1400 0, /* tp_as_buffer */
1401 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1402 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001403 starmap_doc, /* tp_doc */
1404 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001405 0, /* tp_clear */
1406 0, /* tp_richcompare */
1407 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001408 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001409 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001410 0, /* tp_methods */
1411 0, /* tp_members */
1412 0, /* tp_getset */
1413 0, /* tp_base */
1414 0, /* tp_dict */
1415 0, /* tp_descr_get */
1416 0, /* tp_descr_set */
1417 0, /* tp_dictoffset */
1418 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001419 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001420 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001421 PyObject_GC_Del, /* tp_free */
1422};
1423
1424
1425/* imap object ************************************************************/
1426
1427typedef struct {
1428 PyObject_HEAD
1429 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001430 PyObject *func;
1431} imapobject;
1432
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001433static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001434
1435static PyObject *
1436imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1437{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001438 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001439 imapobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001440 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001441
Georg Brandl02c42872005-08-26 06:42:30 +00001442 if (!_PyArg_NoKeywords("imap()", kwds))
1443 return NULL;
1444
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001445 numargs = PyTuple_Size(args);
1446 if (numargs < 2) {
1447 PyErr_SetString(PyExc_TypeError,
1448 "imap() must have at least two arguments.");
1449 return NULL;
1450 }
1451
1452 iters = PyTuple_New(numargs-1);
1453 if (iters == NULL)
1454 return NULL;
1455
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001456 for (i=1 ; i<numargs ; i++) {
1457 /* Get iterator. */
1458 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1459 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001460 Py_DECREF(iters);
1461 return NULL;
1462 }
1463 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001464 }
1465
1466 /* create imapobject structure */
1467 lz = (imapobject *)type->tp_alloc(type, 0);
1468 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469 Py_DECREF(iters);
1470 return NULL;
1471 }
1472 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001473 func = PyTuple_GET_ITEM(args, 0);
1474 Py_INCREF(func);
1475 lz->func = func;
1476
1477 return (PyObject *)lz;
1478}
1479
1480static void
1481imap_dealloc(imapobject *lz)
1482{
1483 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001484 Py_XDECREF(lz->iters);
1485 Py_XDECREF(lz->func);
1486 lz->ob_type->tp_free(lz);
1487}
1488
1489static int
1490imap_traverse(imapobject *lz, visitproc visit, void *arg)
1491{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001492 Py_VISIT(lz->iters);
1493 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001494 return 0;
1495}
1496
Raymond Hettinger2012f172003-02-07 05:32:58 +00001497/*
1498imap() is an iterator version of __builtins__.map() except that it does
1499not have the None fill-in feature. That was intentionally left out for
1500the following reasons:
1501
1502 1) Itertools are designed to be easily combined and chained together.
1503 Having all tools stop with the shortest input is a unifying principle
1504 that makes it easier to combine finite iterators (supplying data) with
1505 infinite iterators like count() and repeat() (for supplying sequential
1506 or constant arguments to a function).
1507
1508 2) In typical use cases for combining itertools, having one finite data
1509 supplier run out before another is likely to be an error condition which
1510 should not pass silently by automatically supplying None.
1511
1512 3) The use cases for automatic None fill-in are rare -- not many functions
1513 do something useful when a parameter suddenly switches type and becomes
1514 None.
1515
1516 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001517 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001518
1519 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1520*/
1521
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001522static PyObject *
1523imap_next(imapobject *lz)
1524{
1525 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001526 PyObject *argtuple;
1527 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001528 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001529
1530 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001531 argtuple = PyTuple_New(numargs);
1532 if (argtuple == NULL)
1533 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001534
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001535 for (i=0 ; i<numargs ; i++) {
1536 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1537 if (val == NULL) {
1538 Py_DECREF(argtuple);
1539 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001540 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001541 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001542 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001543 if (lz->func == Py_None)
1544 return argtuple;
1545 result = PyObject_Call(lz->func, argtuple, NULL);
1546 Py_DECREF(argtuple);
1547 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001548}
1549
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001550PyDoc_STRVAR(imap_doc,
1551"imap(func, *iterables) --> imap object\n\
1552\n\
1553Make an iterator that computes the function using arguments from\n\
1554each of the iterables. Like map() except that it returns\n\
1555an iterator instead of a list and that it stops when the shortest\n\
1556iterable is exhausted instead of filling in None for shorter\n\
1557iterables.");
1558
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001559static PyTypeObject imap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001560 PyObject_HEAD_INIT(NULL)
1561 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001562 "itertools.imap", /* tp_name */
1563 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001564 0, /* tp_itemsize */
1565 /* methods */
1566 (destructor)imap_dealloc, /* tp_dealloc */
1567 0, /* tp_print */
1568 0, /* tp_getattr */
1569 0, /* tp_setattr */
1570 0, /* tp_compare */
1571 0, /* tp_repr */
1572 0, /* tp_as_number */
1573 0, /* tp_as_sequence */
1574 0, /* tp_as_mapping */
1575 0, /* tp_hash */
1576 0, /* tp_call */
1577 0, /* tp_str */
1578 PyObject_GenericGetAttr, /* tp_getattro */
1579 0, /* tp_setattro */
1580 0, /* tp_as_buffer */
1581 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1582 Py_TPFLAGS_BASETYPE, /* tp_flags */
1583 imap_doc, /* tp_doc */
1584 (traverseproc)imap_traverse, /* tp_traverse */
1585 0, /* tp_clear */
1586 0, /* tp_richcompare */
1587 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001588 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001589 (iternextfunc)imap_next, /* tp_iternext */
1590 0, /* tp_methods */
1591 0, /* tp_members */
1592 0, /* tp_getset */
1593 0, /* tp_base */
1594 0, /* tp_dict */
1595 0, /* tp_descr_get */
1596 0, /* tp_descr_set */
1597 0, /* tp_dictoffset */
1598 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001599 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001600 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001601 PyObject_GC_Del, /* tp_free */
1602};
1603
1604
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001605/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001606
1607typedef struct {
1608 PyObject_HEAD
Jack Diederich6c433a92006-05-26 11:15:17 +00001609 Py_ssize_t tuplesize;
1610 Py_ssize_t iternum; /* which iterator is active */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001611 PyObject *ittuple; /* tuple of iterators */
1612} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001613
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001614static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001615
1616static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001617chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001618{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001619 chainobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001620 Py_ssize_t tuplesize = PySequence_Length(args);
Jack Diederich6c433a92006-05-26 11:15:17 +00001621 Py_ssize_t i;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001622 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001623
Georg Brandl02c42872005-08-26 06:42:30 +00001624 if (!_PyArg_NoKeywords("chain()", kwds))
1625 return NULL;
1626
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001627 /* obtain iterators */
1628 assert(PyTuple_Check(args));
1629 ittuple = PyTuple_New(tuplesize);
Neal Norwitz2f3136b2006-05-27 05:18:57 +00001630 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001631 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001632 for (i=0; i < tuplesize; ++i) {
1633 PyObject *item = PyTuple_GET_ITEM(args, i);
1634 PyObject *it = PyObject_GetIter(item);
1635 if (it == NULL) {
1636 if (PyErr_ExceptionMatches(PyExc_TypeError))
1637 PyErr_Format(PyExc_TypeError,
Neal Norwitz2f3136b2006-05-27 05:18:57 +00001638 "chain argument #%zd must support iteration",
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001639 i+1);
1640 Py_DECREF(ittuple);
1641 return NULL;
1642 }
1643 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001644 }
1645
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001646 /* create chainobject structure */
1647 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001648 if (lz == NULL) {
1649 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001650 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001651 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001652
1653 lz->ittuple = ittuple;
1654 lz->iternum = 0;
1655 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001656
1657 return (PyObject *)lz;
1658}
1659
1660static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001661chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001662{
1663 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001664 Py_XDECREF(lz->ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001665 lz->ob_type->tp_free(lz);
1666}
1667
Raymond Hettinger2012f172003-02-07 05:32:58 +00001668static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001669chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001670{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001671 Py_VISIT(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001672 return 0;
1673}
1674
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001675static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001676chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001677{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001678 PyObject *it;
1679 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001680
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001681 while (lz->iternum < lz->tuplesize) {
1682 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1683 item = PyIter_Next(it);
1684 if (item != NULL)
1685 return item;
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001686 if (PyErr_Occurred()) {
1687 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1688 PyErr_Clear();
1689 else
1690 return NULL;
1691 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001692 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001693 }
1694 return NULL;
1695}
1696
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001697PyDoc_STRVAR(chain_doc,
1698"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001699\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001700Return a chain object whose .next() method returns elements from the\n\
1701first iterable until it is exhausted, then elements from the next\n\
1702iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001703
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001704static PyTypeObject chain_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001705 PyObject_HEAD_INIT(NULL)
1706 0, /* ob_size */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001707 "itertools.chain", /* tp_name */
1708 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001709 0, /* tp_itemsize */
1710 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001711 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001712 0, /* tp_print */
1713 0, /* tp_getattr */
1714 0, /* tp_setattr */
1715 0, /* tp_compare */
1716 0, /* tp_repr */
1717 0, /* tp_as_number */
1718 0, /* tp_as_sequence */
1719 0, /* tp_as_mapping */
1720 0, /* tp_hash */
1721 0, /* tp_call */
1722 0, /* tp_str */
1723 PyObject_GenericGetAttr, /* tp_getattro */
1724 0, /* tp_setattro */
1725 0, /* tp_as_buffer */
1726 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1727 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001728 chain_doc, /* tp_doc */
1729 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001730 0, /* tp_clear */
1731 0, /* tp_richcompare */
1732 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001733 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001734 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001735 0, /* tp_methods */
1736 0, /* tp_members */
1737 0, /* tp_getset */
1738 0, /* tp_base */
1739 0, /* tp_dict */
1740 0, /* tp_descr_get */
1741 0, /* tp_descr_set */
1742 0, /* tp_dictoffset */
1743 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001744 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001745 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001746 PyObject_GC_Del, /* tp_free */
1747};
1748
1749
1750/* ifilter object ************************************************************/
1751
1752typedef struct {
1753 PyObject_HEAD
1754 PyObject *func;
1755 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001756} ifilterobject;
1757
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001758static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001759
1760static PyObject *
1761ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1762{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001763 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001764 PyObject *it;
1765 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001766
Georg Brandl02c42872005-08-26 06:42:30 +00001767 if (!_PyArg_NoKeywords("ifilter()", kwds))
1768 return NULL;
1769
Raymond Hettinger60eca932003-02-09 06:40:58 +00001770 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001771 return NULL;
1772
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001773 /* Get iterator. */
1774 it = PyObject_GetIter(seq);
1775 if (it == NULL)
1776 return NULL;
1777
1778 /* create ifilterobject structure */
1779 lz = (ifilterobject *)type->tp_alloc(type, 0);
1780 if (lz == NULL) {
1781 Py_DECREF(it);
1782 return NULL;
1783 }
1784 Py_INCREF(func);
1785 lz->func = func;
1786 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001787
1788 return (PyObject *)lz;
1789}
1790
1791static void
1792ifilter_dealloc(ifilterobject *lz)
1793{
1794 PyObject_GC_UnTrack(lz);
1795 Py_XDECREF(lz->func);
1796 Py_XDECREF(lz->it);
1797 lz->ob_type->tp_free(lz);
1798}
1799
1800static int
1801ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1802{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001803 Py_VISIT(lz->it);
1804 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001805 return 0;
1806}
1807
1808static PyObject *
1809ifilter_next(ifilterobject *lz)
1810{
1811 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001812 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001813 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001814 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001815
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001816 assert(PyIter_Check(it));
1817 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001818 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001819 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001820 if (item == NULL)
1821 return NULL;
1822
1823 if (lz->func == Py_None) {
1824 ok = PyObject_IsTrue(item);
1825 } else {
1826 PyObject *good;
1827 good = PyObject_CallFunctionObjArgs(lz->func,
1828 item, NULL);
1829 if (good == NULL) {
1830 Py_DECREF(item);
1831 return NULL;
1832 }
1833 ok = PyObject_IsTrue(good);
1834 Py_DECREF(good);
1835 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00001836 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001837 return item;
1838 Py_DECREF(item);
1839 }
1840}
1841
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001842PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001843"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001844\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001845Return those items of sequence for which function(item) is true.\n\
1846If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001847
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001848static PyTypeObject ifilter_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001849 PyObject_HEAD_INIT(NULL)
1850 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001851 "itertools.ifilter", /* tp_name */
1852 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001853 0, /* tp_itemsize */
1854 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001855 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001856 0, /* tp_print */
1857 0, /* tp_getattr */
1858 0, /* tp_setattr */
1859 0, /* tp_compare */
1860 0, /* tp_repr */
1861 0, /* tp_as_number */
1862 0, /* tp_as_sequence */
1863 0, /* tp_as_mapping */
1864 0, /* tp_hash */
1865 0, /* tp_call */
1866 0, /* tp_str */
1867 PyObject_GenericGetAttr, /* tp_getattro */
1868 0, /* tp_setattro */
1869 0, /* tp_as_buffer */
1870 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1871 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001872 ifilter_doc, /* tp_doc */
1873 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001874 0, /* tp_clear */
1875 0, /* tp_richcompare */
1876 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001877 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001878 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001879 0, /* tp_methods */
1880 0, /* tp_members */
1881 0, /* tp_getset */
1882 0, /* tp_base */
1883 0, /* tp_dict */
1884 0, /* tp_descr_get */
1885 0, /* tp_descr_set */
1886 0, /* tp_dictoffset */
1887 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001888 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001889 ifilter_new, /* tp_new */
1890 PyObject_GC_Del, /* tp_free */
1891};
1892
1893
1894/* ifilterfalse object ************************************************************/
1895
1896typedef struct {
1897 PyObject_HEAD
1898 PyObject *func;
1899 PyObject *it;
1900} ifilterfalseobject;
1901
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001902static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001903
1904static PyObject *
1905ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1906{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00001907 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001908 PyObject *it;
1909 ifilterfalseobject *lz;
1910
Georg Brandl02c42872005-08-26 06:42:30 +00001911 if (!_PyArg_NoKeywords("ifilterfalse()", kwds))
1912 return NULL;
1913
Raymond Hettinger60eca932003-02-09 06:40:58 +00001914 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1915 return NULL;
1916
1917 /* Get iterator. */
1918 it = PyObject_GetIter(seq);
1919 if (it == NULL)
1920 return NULL;
1921
1922 /* create ifilterfalseobject structure */
1923 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1924 if (lz == NULL) {
1925 Py_DECREF(it);
1926 return NULL;
1927 }
1928 Py_INCREF(func);
1929 lz->func = func;
1930 lz->it = it;
1931
1932 return (PyObject *)lz;
1933}
1934
1935static void
1936ifilterfalse_dealloc(ifilterfalseobject *lz)
1937{
1938 PyObject_GC_UnTrack(lz);
1939 Py_XDECREF(lz->func);
1940 Py_XDECREF(lz->it);
1941 lz->ob_type->tp_free(lz);
1942}
1943
1944static int
1945ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1946{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001947 Py_VISIT(lz->it);
1948 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001949 return 0;
1950}
1951
1952static PyObject *
1953ifilterfalse_next(ifilterfalseobject *lz)
1954{
1955 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001956 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001957 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001958 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001959
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001960 assert(PyIter_Check(it));
1961 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001962 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001963 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001964 if (item == NULL)
1965 return NULL;
1966
1967 if (lz->func == Py_None) {
1968 ok = PyObject_IsTrue(item);
1969 } else {
1970 PyObject *good;
1971 good = PyObject_CallFunctionObjArgs(lz->func,
1972 item, NULL);
1973 if (good == NULL) {
1974 Py_DECREF(item);
1975 return NULL;
1976 }
1977 ok = PyObject_IsTrue(good);
1978 Py_DECREF(good);
1979 }
1980 if (!ok)
1981 return item;
1982 Py_DECREF(item);
1983 }
1984}
1985
Raymond Hettinger60eca932003-02-09 06:40:58 +00001986PyDoc_STRVAR(ifilterfalse_doc,
1987"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1988\n\
1989Return those items of sequence for which function(item) is false.\n\
1990If function is None, return the items that are false.");
1991
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001992static PyTypeObject ifilterfalse_type = {
Raymond Hettinger60eca932003-02-09 06:40:58 +00001993 PyObject_HEAD_INIT(NULL)
1994 0, /* ob_size */
1995 "itertools.ifilterfalse", /* tp_name */
1996 sizeof(ifilterfalseobject), /* tp_basicsize */
1997 0, /* tp_itemsize */
1998 /* methods */
1999 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
2000 0, /* tp_print */
2001 0, /* tp_getattr */
2002 0, /* tp_setattr */
2003 0, /* tp_compare */
2004 0, /* tp_repr */
2005 0, /* tp_as_number */
2006 0, /* tp_as_sequence */
2007 0, /* tp_as_mapping */
2008 0, /* tp_hash */
2009 0, /* tp_call */
2010 0, /* tp_str */
2011 PyObject_GenericGetAttr, /* tp_getattro */
2012 0, /* tp_setattro */
2013 0, /* tp_as_buffer */
2014 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2015 Py_TPFLAGS_BASETYPE, /* tp_flags */
2016 ifilterfalse_doc, /* tp_doc */
2017 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
2018 0, /* tp_clear */
2019 0, /* tp_richcompare */
2020 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002021 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002022 (iternextfunc)ifilterfalse_next, /* tp_iternext */
2023 0, /* tp_methods */
2024 0, /* tp_members */
2025 0, /* tp_getset */
2026 0, /* tp_base */
2027 0, /* tp_dict */
2028 0, /* tp_descr_get */
2029 0, /* tp_descr_set */
2030 0, /* tp_dictoffset */
2031 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002032 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002033 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002034 PyObject_GC_Del, /* tp_free */
2035};
2036
2037
2038/* count object ************************************************************/
2039
2040typedef struct {
2041 PyObject_HEAD
Jack Diederich6c433a92006-05-26 11:15:17 +00002042 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002043} countobject;
2044
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002045static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002046
2047static PyObject *
2048count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2049{
2050 countobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00002051 Py_ssize_t cnt = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002052
Georg Brandl02c42872005-08-26 06:42:30 +00002053 if (!_PyArg_NoKeywords("count()", kwds))
2054 return NULL;
2055
Jack Diederich6c433a92006-05-26 11:15:17 +00002056 if (!PyArg_ParseTuple(args, "|n:count", &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002057 return NULL;
2058
2059 /* create countobject structure */
2060 lz = (countobject *)PyObject_New(countobject, &count_type);
2061 if (lz == NULL)
2062 return NULL;
2063 lz->cnt = cnt;
2064
2065 return (PyObject *)lz;
2066}
2067
2068static PyObject *
2069count_next(countobject *lz)
2070{
Jack Diederich6c433a92006-05-26 11:15:17 +00002071 return PyInt_FromSize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002072}
2073
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002074static PyObject *
2075count_repr(countobject *lz)
2076{
Jack Diederich6c433a92006-05-26 11:15:17 +00002077 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002078}
2079
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002080PyDoc_STRVAR(count_doc,
2081"count([firstval]) --> count object\n\
2082\n\
2083Return a count object whose .next() method returns consecutive\n\
2084integers starting from zero or, if specified, from firstval.");
2085
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002086static PyTypeObject count_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002087 PyObject_HEAD_INIT(NULL)
2088 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002089 "itertools.count", /* tp_name */
2090 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002091 0, /* tp_itemsize */
2092 /* methods */
2093 (destructor)PyObject_Del, /* tp_dealloc */
2094 0, /* tp_print */
2095 0, /* tp_getattr */
2096 0, /* tp_setattr */
2097 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002098 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002099 0, /* tp_as_number */
2100 0, /* tp_as_sequence */
2101 0, /* tp_as_mapping */
2102 0, /* tp_hash */
2103 0, /* tp_call */
2104 0, /* tp_str */
2105 PyObject_GenericGetAttr, /* tp_getattro */
2106 0, /* tp_setattro */
2107 0, /* tp_as_buffer */
2108 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002109 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002110 0, /* tp_traverse */
2111 0, /* tp_clear */
2112 0, /* tp_richcompare */
2113 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002114 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002115 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002116 0, /* tp_methods */
2117 0, /* tp_members */
2118 0, /* tp_getset */
2119 0, /* tp_base */
2120 0, /* tp_dict */
2121 0, /* tp_descr_get */
2122 0, /* tp_descr_set */
2123 0, /* tp_dictoffset */
2124 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002125 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002126 count_new, /* tp_new */
2127};
2128
2129
2130/* izip object ************************************************************/
2131
2132#include "Python.h"
2133
2134typedef struct {
2135 PyObject_HEAD
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002136 Py_ssize_t tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002137 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002138 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002139} izipobject;
2140
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002141static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002142
2143static PyObject *
2144izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2145{
2146 izipobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00002147 Py_ssize_t i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002148 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002149 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002150 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002151
Georg Brandl02c42872005-08-26 06:42:30 +00002152 if (!_PyArg_NoKeywords("izip()", kwds))
2153 return NULL;
2154
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002155 /* args must be a tuple */
2156 assert(PyTuple_Check(args));
2157
2158 /* obtain iterators */
2159 ittuple = PyTuple_New(tuplesize);
Neal Norwitz2f3136b2006-05-27 05:18:57 +00002160 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002161 return NULL;
2162 for (i=0; i < tuplesize; ++i) {
2163 PyObject *item = PyTuple_GET_ITEM(args, i);
2164 PyObject *it = PyObject_GetIter(item);
2165 if (it == NULL) {
2166 if (PyErr_ExceptionMatches(PyExc_TypeError))
2167 PyErr_Format(PyExc_TypeError,
Neal Norwitz2f3136b2006-05-27 05:18:57 +00002168 "izip argument #%zd must support iteration",
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002169 i+1);
2170 Py_DECREF(ittuple);
2171 return NULL;
2172 }
2173 PyTuple_SET_ITEM(ittuple, i, it);
2174 }
2175
Raymond Hettinger2012f172003-02-07 05:32:58 +00002176 /* create a result holder */
2177 result = PyTuple_New(tuplesize);
2178 if (result == NULL) {
2179 Py_DECREF(ittuple);
2180 return NULL;
2181 }
2182 for (i=0 ; i < tuplesize ; i++) {
2183 Py_INCREF(Py_None);
2184 PyTuple_SET_ITEM(result, i, Py_None);
2185 }
2186
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002187 /* create izipobject structure */
2188 lz = (izipobject *)type->tp_alloc(type, 0);
2189 if (lz == NULL) {
2190 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002191 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002192 return NULL;
2193 }
2194 lz->ittuple = ittuple;
2195 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002196 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002197
2198 return (PyObject *)lz;
2199}
2200
2201static void
2202izip_dealloc(izipobject *lz)
2203{
2204 PyObject_GC_UnTrack(lz);
2205 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002206 Py_XDECREF(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002207 lz->ob_type->tp_free(lz);
2208}
2209
2210static int
2211izip_traverse(izipobject *lz, visitproc visit, void *arg)
2212{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002213 Py_VISIT(lz->ittuple);
2214 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002215 return 0;
2216}
2217
2218static PyObject *
2219izip_next(izipobject *lz)
2220{
Jack Diederich6c433a92006-05-26 11:15:17 +00002221 Py_ssize_t i;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002222 Py_ssize_t tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002223 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002224 PyObject *it;
2225 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002226 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002227
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002228 if (tuplesize == 0)
2229 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002230 if (result->ob_refcnt == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002231 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002232 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002233 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002234 assert(PyIter_Check(it));
2235 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002236 if (item == NULL) {
2237 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002238 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002239 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002240 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002241 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002242 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002243 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002244 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002245 result = PyTuple_New(tuplesize);
2246 if (result == NULL)
2247 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002248 for (i=0 ; i < tuplesize ; i++) {
2249 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002250 assert(PyIter_Check(it));
2251 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002252 if (item == NULL) {
2253 Py_DECREF(result);
2254 return NULL;
2255 }
2256 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002257 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002258 }
2259 return result;
2260}
2261
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002262PyDoc_STRVAR(izip_doc,
2263"izip(iter1 [,iter2 [...]]) --> izip object\n\
2264\n\
2265Return a izip object whose .next() method returns a tuple where\n\
2266the i-th element comes from the i-th iterable argument. The .next()\n\
2267method continues until the shortest iterable in the argument sequence\n\
2268is exhausted and then it raises StopIteration. Works like the zip()\n\
2269function but consumes less memory by returning an iterator instead of\n\
2270a list.");
2271
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002272static PyTypeObject izip_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002273 PyObject_HEAD_INIT(NULL)
2274 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002275 "itertools.izip", /* tp_name */
2276 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002277 0, /* tp_itemsize */
2278 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002279 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002280 0, /* tp_print */
2281 0, /* tp_getattr */
2282 0, /* tp_setattr */
2283 0, /* tp_compare */
2284 0, /* tp_repr */
2285 0, /* tp_as_number */
2286 0, /* tp_as_sequence */
2287 0, /* tp_as_mapping */
2288 0, /* tp_hash */
2289 0, /* tp_call */
2290 0, /* tp_str */
2291 PyObject_GenericGetAttr, /* tp_getattro */
2292 0, /* tp_setattro */
2293 0, /* tp_as_buffer */
2294 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2295 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002296 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002297 (traverseproc)izip_traverse, /* tp_traverse */
2298 0, /* tp_clear */
2299 0, /* tp_richcompare */
2300 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002301 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002302 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002303 0, /* tp_methods */
2304 0, /* tp_members */
2305 0, /* tp_getset */
2306 0, /* tp_base */
2307 0, /* tp_dict */
2308 0, /* tp_descr_get */
2309 0, /* tp_descr_set */
2310 0, /* tp_dictoffset */
2311 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002312 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002313 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002314 PyObject_GC_Del, /* tp_free */
2315};
2316
2317
2318/* repeat object ************************************************************/
2319
2320typedef struct {
2321 PyObject_HEAD
2322 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00002323 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002324} repeatobject;
2325
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002326static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002327
2328static PyObject *
2329repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2330{
2331 repeatobject *ro;
2332 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00002333 Py_ssize_t cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002334
Georg Brandl02c42872005-08-26 06:42:30 +00002335 if (!_PyArg_NoKeywords("repeat()", kwds))
2336 return NULL;
2337
Jack Diederich6c433a92006-05-26 11:15:17 +00002338 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002339 return NULL;
2340
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002341 if (PyTuple_Size(args) == 2 && cnt < 0)
2342 cnt = 0;
2343
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002344 ro = (repeatobject *)type->tp_alloc(type, 0);
2345 if (ro == NULL)
2346 return NULL;
2347 Py_INCREF(element);
2348 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002349 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002350 return (PyObject *)ro;
2351}
2352
2353static void
2354repeat_dealloc(repeatobject *ro)
2355{
2356 PyObject_GC_UnTrack(ro);
2357 Py_XDECREF(ro->element);
2358 ro->ob_type->tp_free(ro);
2359}
2360
2361static int
2362repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2363{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002364 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002365 return 0;
2366}
2367
2368static PyObject *
2369repeat_next(repeatobject *ro)
2370{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002371 if (ro->cnt == 0)
2372 return NULL;
2373 if (ro->cnt > 0)
2374 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002375 Py_INCREF(ro->element);
2376 return ro->element;
2377}
2378
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002379static PyObject *
2380repeat_repr(repeatobject *ro)
2381{
2382 PyObject *result, *objrepr;
2383
2384 objrepr = PyObject_Repr(ro->element);
2385 if (objrepr == NULL)
2386 return NULL;
2387
2388 if (ro->cnt == -1)
2389 result = PyString_FromFormat("repeat(%s)",
2390 PyString_AS_STRING(objrepr));
2391 else
Jack Diederich6c433a92006-05-26 11:15:17 +00002392 result = PyString_FromFormat("repeat(%s, %zd)",
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002393 PyString_AS_STRING(objrepr), ro->cnt);
2394 Py_DECREF(objrepr);
2395 return result;
2396}
2397
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002398static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002399repeat_len(repeatobject *ro)
2400{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002401 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002402 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002403 return NULL;
2404 }
Jack Diederich6c433a92006-05-26 11:15:17 +00002405 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002406}
2407
Armin Rigof5b3e362006-02-11 21:32:43 +00002408PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002409
2410static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00002411 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002412 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002413};
2414
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002415PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002416"repeat(element [,times]) -> create an iterator which returns the element\n\
2417for the specified number of times. If not specified, returns the element\n\
2418endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002419
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002420static PyTypeObject repeat_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002421 PyObject_HEAD_INIT(NULL)
2422 0, /* ob_size */
2423 "itertools.repeat", /* tp_name */
2424 sizeof(repeatobject), /* tp_basicsize */
2425 0, /* tp_itemsize */
2426 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002427 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002428 0, /* tp_print */
2429 0, /* tp_getattr */
2430 0, /* tp_setattr */
2431 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002432 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002433 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002434 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002435 0, /* tp_as_mapping */
2436 0, /* tp_hash */
2437 0, /* tp_call */
2438 0, /* tp_str */
2439 PyObject_GenericGetAttr, /* tp_getattro */
2440 0, /* tp_setattro */
2441 0, /* tp_as_buffer */
2442 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2443 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002444 repeat_doc, /* tp_doc */
2445 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002446 0, /* tp_clear */
2447 0, /* tp_richcompare */
2448 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002449 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002450 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002451 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002452 0, /* tp_members */
2453 0, /* tp_getset */
2454 0, /* tp_base */
2455 0, /* tp_dict */
2456 0, /* tp_descr_get */
2457 0, /* tp_descr_set */
2458 0, /* tp_dictoffset */
2459 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002460 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002461 repeat_new, /* tp_new */
2462 PyObject_GC_Del, /* tp_free */
2463};
2464
2465
2466/* module level code ********************************************************/
2467
2468PyDoc_STRVAR(module_doc,
2469"Functional tools for creating and using iterators.\n\
2470\n\
2471Infinite iterators:\n\
2472count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002473cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002474repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002475\n\
2476Iterators terminating on the shortest input sequence:\n\
2477izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002478ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2479ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002480islice(seq, [start,] stop [, step]) --> elements from\n\
2481 seq[start:stop:step]\n\
2482imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2483starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00002484tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002485chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002486takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2487dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002488groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002489");
2490
2491
Raymond Hettingerad983e72003-11-12 14:32:26 +00002492static PyMethodDef module_methods[] = {
2493 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2494 {NULL, NULL} /* sentinel */
2495};
2496
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002497PyMODINIT_FUNC
2498inititertools(void)
2499{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002500 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002501 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002502 char *name;
2503 PyTypeObject *typelist[] = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002504 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002505 &dropwhile_type,
2506 &takewhile_type,
2507 &islice_type,
2508 &starmap_type,
2509 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002510 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002511 &ifilter_type,
2512 &ifilterfalse_type,
2513 &count_type,
2514 &izip_type,
2515 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002516 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002517 NULL
2518 };
2519
Skip Montanarof3938fd2004-02-10 20:27:40 +00002520 teedataobject_type.ob_type = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00002521 m = Py_InitModule3("itertools", module_methods, module_doc);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00002522 if (m == NULL)
2523 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002524
Raymond Hettinger60eca932003-02-09 06:40:58 +00002525 for (i=0 ; typelist[i] != NULL ; i++) {
2526 if (PyType_Ready(typelist[i]) < 0)
2527 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002528 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002529 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002530 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002531 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002532 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00002533
2534 if (PyType_Ready(&teedataobject_type) < 0)
2535 return;
2536 if (PyType_Ready(&tee_type) < 0)
2537 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002538 if (PyType_Ready(&_grouper_type) < 0)
2539 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002540}