blob: 55731eb92f0fea4d861abde0958b211e38bb4f41 [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
Raymond Hettinger3662c902008-03-06 22:58:42 +0000202 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000203 if (igo == NULL)
204 return NULL;
205 igo->parent = (PyObject *)parent;
206 Py_INCREF(parent);
207 igo->tgtkey = tgtkey;
208 Py_INCREF(tgtkey);
209
Raymond Hettinger3662c902008-03-06 22:58:42 +0000210 PyObject_GC_Track(igo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000211 return (PyObject *)igo;
212}
213
214static void
215_grouper_dealloc(_grouperobject *igo)
216{
Raymond Hettinger3662c902008-03-06 22:58:42 +0000217 PyObject_GC_UnTrack(igo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000218 Py_DECREF(igo->parent);
219 Py_DECREF(igo->tgtkey);
Raymond Hettinger3662c902008-03-06 22:58:42 +0000220 PyObject_GC_Del(igo);
221}
222
223static int
224_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
225{
226 Py_VISIT(igo->parent);
227 Py_VISIT(igo->tgtkey);
228 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000229}
230
231static PyObject *
232_grouper_next(_grouperobject *igo)
233{
234 groupbyobject *gbo = (groupbyobject *)igo->parent;
235 PyObject *newvalue, *newkey, *r;
236 int rcmp;
237
238 if (gbo->currvalue == NULL) {
239 newvalue = PyIter_Next(gbo->it);
240 if (newvalue == NULL)
241 return NULL;
242
243 if (gbo->keyfunc == Py_None) {
244 newkey = newvalue;
245 Py_INCREF(newvalue);
246 } else {
247 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
248 newvalue, NULL);
249 if (newkey == NULL) {
250 Py_DECREF(newvalue);
251 return NULL;
252 }
253 }
254
255 assert(gbo->currkey == NULL);
256 gbo->currkey = newkey;
257 gbo->currvalue = newvalue;
258 }
259
260 assert(gbo->currkey != NULL);
261 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
262 if (rcmp <= 0)
263 /* got any error or current group is end */
264 return NULL;
265
266 r = gbo->currvalue;
267 gbo->currvalue = NULL;
Raymond Hettinger75ccea32004-09-01 07:02:44 +0000268 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000269
270 return r;
271}
272
273static PyTypeObject _grouper_type = {
274 PyObject_HEAD_INIT(NULL)
275 0, /* ob_size */
276 "itertools._grouper", /* tp_name */
277 sizeof(_grouperobject), /* tp_basicsize */
278 0, /* tp_itemsize */
279 /* methods */
280 (destructor)_grouper_dealloc, /* tp_dealloc */
281 0, /* tp_print */
282 0, /* tp_getattr */
283 0, /* tp_setattr */
284 0, /* tp_compare */
285 0, /* tp_repr */
286 0, /* tp_as_number */
287 0, /* tp_as_sequence */
288 0, /* tp_as_mapping */
289 0, /* tp_hash */
290 0, /* tp_call */
291 0, /* tp_str */
292 PyObject_GenericGetAttr, /* tp_getattro */
293 0, /* tp_setattro */
294 0, /* tp_as_buffer */
Raymond Hettinger3662c902008-03-06 22:58:42 +0000295 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000296 0, /* tp_doc */
Raymond Hettinger3662c902008-03-06 22:58:42 +0000297 (traverseproc)_grouper_traverse,/* tp_traverse */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000298 0, /* tp_clear */
299 0, /* tp_richcompare */
300 0, /* tp_weaklistoffset */
301 PyObject_SelfIter, /* tp_iter */
302 (iternextfunc)_grouper_next, /* tp_iternext */
303 0, /* tp_methods */
304 0, /* tp_members */
305 0, /* tp_getset */
306 0, /* tp_base */
307 0, /* tp_dict */
308 0, /* tp_descr_get */
309 0, /* tp_descr_set */
310 0, /* tp_dictoffset */
311 0, /* tp_init */
312 0, /* tp_alloc */
313 0, /* tp_new */
Raymond Hettinger3662c902008-03-06 22:58:42 +0000314 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000315};
316
317
318
Raymond Hettingerad983e72003-11-12 14:32:26 +0000319/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000320
Raymond Hettingerad983e72003-11-12 14:32:26 +0000321/* The teedataobject pre-allocates space for LINKCELLS number of objects.
322 To help the object fit neatly inside cache lines (space for 16 to 32
323 pointers), the value should be a multiple of 16 minus space for
324 the other structure members including PyHEAD overhead. The larger the
325 value, the less memory overhead per object and the less time spent
326 allocating/deallocating new links. The smaller the number, the less
327 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000328*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000329#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000330
331typedef struct {
332 PyObject_HEAD
333 PyObject *it;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000334 int numread;
335 PyObject *nextlink;
336 PyObject *(values[LINKCELLS]);
337} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000338
339typedef struct {
340 PyObject_HEAD
Raymond Hettingerad983e72003-11-12 14:32:26 +0000341 teedataobject *dataobj;
342 int index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000343 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000344} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000345
Raymond Hettingerad983e72003-11-12 14:32:26 +0000346static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000347
348static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000349teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000350{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000351 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000352
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000353 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000354 if (tdo == NULL)
Raymond Hettinger45143692003-10-25 06:37:47 +0000355 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000356
357 tdo->numread = 0;
358 tdo->nextlink = NULL;
359 Py_INCREF(it);
360 tdo->it = it;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000361 PyObject_GC_Track(tdo);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000362 return (PyObject *)tdo;
363}
364
365static PyObject *
366teedataobject_jumplink(teedataobject *tdo)
367{
368 if (tdo->nextlink == NULL)
369 tdo->nextlink = teedataobject_new(tdo->it);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000370 Py_XINCREF(tdo->nextlink);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000371 return tdo->nextlink;
372}
373
374static PyObject *
375teedataobject_getitem(teedataobject *tdo, int i)
376{
377 PyObject *value;
378
379 assert(i < LINKCELLS);
380 if (i < tdo->numread)
381 value = tdo->values[i];
382 else {
383 /* this is the lead iterator, so fetch more data */
384 assert(i == tdo->numread);
385 value = PyIter_Next(tdo->it);
386 if (value == NULL)
387 return NULL;
388 tdo->numread++;
389 tdo->values[i] = value;
Raymond Hettinger45143692003-10-25 06:37:47 +0000390 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000391 Py_INCREF(value);
392 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000393}
394
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000395static int
396teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
397{
398 int i;
399 Py_VISIT(tdo->it);
400 for (i = 0; i < tdo->numread; i++)
401 Py_VISIT(tdo->values[i]);
402 Py_VISIT(tdo->nextlink);
403 return 0;
404}
405
406static int
407teedataobject_clear(teedataobject *tdo)
408{
409 int i;
410 Py_CLEAR(tdo->it);
411 for (i=0 ; i<tdo->numread ; i++)
412 Py_CLEAR(tdo->values[i]);
413 Py_CLEAR(tdo->nextlink);
414 return 0;
415}
416
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000418teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000419{
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000420 PyObject_GC_UnTrack(tdo);
421 teedataobject_clear(tdo);
422 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000423}
424
Raymond Hettingerad983e72003-11-12 14:32:26 +0000425PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000426
Raymond Hettingerad983e72003-11-12 14:32:26 +0000427static PyTypeObject teedataobject_type = {
Skip Montanarof3938fd2004-02-10 20:27:40 +0000428 PyObject_HEAD_INIT(0) /* Must fill in type value later */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000429 0, /* ob_size */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000430 "itertools.tee_dataobject", /* tp_name */
431 sizeof(teedataobject), /* tp_basicsize */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000432 0, /* tp_itemsize */
433 /* methods */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000434 (destructor)teedataobject_dealloc, /* tp_dealloc */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000435 0, /* tp_print */
436 0, /* tp_getattr */
437 0, /* tp_setattr */
438 0, /* tp_compare */
439 0, /* tp_repr */
440 0, /* tp_as_number */
441 0, /* tp_as_sequence */
442 0, /* tp_as_mapping */
443 0, /* tp_hash */
444 0, /* tp_call */
445 0, /* tp_str */
446 PyObject_GenericGetAttr, /* tp_getattro */
447 0, /* tp_setattro */
448 0, /* tp_as_buffer */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000449 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000450 teedataobject_doc, /* tp_doc */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000451 (traverseproc)teedataobject_traverse, /* tp_traverse */
452 (inquiry)teedataobject_clear, /* tp_clear */
453 0, /* tp_richcompare */
454 0, /* tp_weaklistoffset */
455 0, /* tp_iter */
456 0, /* tp_iternext */
457 0, /* tp_methods */
458 0, /* tp_members */
459 0, /* tp_getset */
460 0, /* tp_base */
461 0, /* tp_dict */
462 0, /* tp_descr_get */
463 0, /* tp_descr_set */
464 0, /* tp_dictoffset */
465 0, /* tp_init */
466 0, /* tp_alloc */
467 0, /* tp_new */
468 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000469};
470
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000471
472static PyTypeObject tee_type;
473
474static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000475tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000476{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000477 PyObject *value, *link;
478
479 if (to->index >= LINKCELLS) {
480 link = teedataobject_jumplink(to->dataobj);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000481 Py_DECREF(to->dataobj);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000482 to->dataobj = (teedataobject *)link;
483 to->index = 0;
484 }
485 value = teedataobject_getitem(to->dataobj, to->index);
486 if (value == NULL)
487 return NULL;
488 to->index++;
489 return value;
490}
491
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000492static int
493tee_traverse(teeobject *to, visitproc visit, void *arg)
494{
495 Py_VISIT((PyObject *)to->dataobj);
496 return 0;
497}
498
Raymond Hettingerad983e72003-11-12 14:32:26 +0000499static PyObject *
500tee_copy(teeobject *to)
501{
502 teeobject *newto;
503
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000504 newto = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000505 if (newto == NULL)
506 return NULL;
507 Py_INCREF(to->dataobj);
508 newto->dataobj = to->dataobj;
509 newto->index = to->index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000510 newto->weakreflist = NULL;
Thomas Woutersb3deb942006-04-15 22:33:13 +0000511 PyObject_GC_Track(newto);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000512 return (PyObject *)newto;
513}
514
515PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
516
517static PyObject *
518tee_fromiterable(PyObject *iterable)
519{
520 teeobject *to;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000521 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000522
523 it = PyObject_GetIter(iterable);
524 if (it == NULL)
525 return NULL;
526 if (PyObject_TypeCheck(it, &tee_type)) {
527 to = (teeobject *)tee_copy((teeobject *)it);
528 goto done;
529 }
530
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000531 to = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000532 if (to == NULL)
533 goto done;
534 to->dataobj = (teedataobject *)teedataobject_new(it);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000535 if (!to->dataobj) {
536 PyObject_GC_Del(to);
537 to = NULL;
538 goto done;
539 }
540
Raymond Hettingerad983e72003-11-12 14:32:26 +0000541 to->index = 0;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000542 to->weakreflist = NULL;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000543 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000544done:
545 Py_XDECREF(it);
546 return (PyObject *)to;
547}
548
549static PyObject *
550tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
551{
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000552 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000553
554 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
555 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000556 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000557}
558
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000559static int
560tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000561{
Raymond Hettingera9f60922004-10-17 16:40:14 +0000562 if (to->weakreflist != NULL)
563 PyObject_ClearWeakRefs((PyObject *) to);
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000564 Py_CLEAR(to->dataobj);
565 return 0;
566}
567
568static void
569tee_dealloc(teeobject *to)
570{
571 PyObject_GC_UnTrack(to);
572 tee_clear(to);
573 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000574}
575
Raymond Hettingerad983e72003-11-12 14:32:26 +0000576PyDoc_STRVAR(teeobject_doc,
577"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000578
Raymond Hettingerad983e72003-11-12 14:32:26 +0000579static PyMethodDef tee_methods[] = {
580 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
581 {NULL, NULL} /* sentinel */
582};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000583
584static PyTypeObject tee_type = {
585 PyObject_HEAD_INIT(NULL)
586 0, /* ob_size */
587 "itertools.tee", /* tp_name */
588 sizeof(teeobject), /* tp_basicsize */
589 0, /* tp_itemsize */
590 /* methods */
591 (destructor)tee_dealloc, /* tp_dealloc */
592 0, /* tp_print */
593 0, /* tp_getattr */
594 0, /* tp_setattr */
595 0, /* tp_compare */
596 0, /* tp_repr */
597 0, /* tp_as_number */
598 0, /* tp_as_sequence */
599 0, /* tp_as_mapping */
600 0, /* tp_hash */
601 0, /* tp_call */
602 0, /* tp_str */
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000603 0, /* tp_getattro */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000604 0, /* tp_setattro */
605 0, /* tp_as_buffer */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000606 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000607 teeobject_doc, /* tp_doc */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000608 (traverseproc)tee_traverse, /* tp_traverse */
609 (inquiry)tee_clear, /* tp_clear */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000610 0, /* tp_richcompare */
Raymond Hettingera9f60922004-10-17 16:40:14 +0000611 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000612 PyObject_SelfIter, /* tp_iter */
613 (iternextfunc)tee_next, /* tp_iternext */
614 tee_methods, /* tp_methods */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000615 0, /* tp_members */
616 0, /* tp_getset */
617 0, /* tp_base */
618 0, /* tp_dict */
619 0, /* tp_descr_get */
620 0, /* tp_descr_set */
621 0, /* tp_dictoffset */
622 0, /* tp_init */
623 0, /* tp_alloc */
624 tee_new, /* tp_new */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000625 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000626};
627
Raymond Hettingerad983e72003-11-12 14:32:26 +0000628static PyObject *
629tee(PyObject *self, PyObject *args)
630{
Neal Norwitz7ae5f292006-09-05 02:30:10 +0000631 Py_ssize_t i, n=2;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000632 PyObject *it, *iterable, *copyable, *result;
633
Neal Norwitz7ae5f292006-09-05 02:30:10 +0000634 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000635 return NULL;
Neal Norwitz7ae5f292006-09-05 02:30:10 +0000636 if (n < 0) {
637 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
638 return NULL;
639 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000640 result = PyTuple_New(n);
641 if (result == NULL)
642 return NULL;
643 if (n == 0)
644 return result;
645 it = PyObject_GetIter(iterable);
646 if (it == NULL) {
647 Py_DECREF(result);
648 return NULL;
649 }
650 if (!PyObject_HasAttrString(it, "__copy__")) {
651 copyable = tee_fromiterable(it);
652 Py_DECREF(it);
653 if (copyable == NULL) {
654 Py_DECREF(result);
655 return NULL;
656 }
657 } else
658 copyable = it;
659 PyTuple_SET_ITEM(result, 0, copyable);
660 for (i=1 ; i<n ; i++) {
661 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
662 if (copyable == NULL) {
663 Py_DECREF(result);
664 return NULL;
665 }
666 PyTuple_SET_ITEM(result, i, copyable);
667 }
668 return result;
669}
670
671PyDoc_STRVAR(tee_doc,
672"tee(iterable, n=2) --> tuple of n independent iterators.");
673
674
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000675/* cycle object **********************************************************/
676
677typedef struct {
678 PyObject_HEAD
679 PyObject *it;
680 PyObject *saved;
681 int firstpass;
682} cycleobject;
683
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000684static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000685
686static PyObject *
687cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
688{
689 PyObject *it;
690 PyObject *iterable;
691 PyObject *saved;
692 cycleobject *lz;
693
Georg Brandl8e932e72007-01-21 10:28:56 +0000694 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000695 return NULL;
696
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000697 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
698 return NULL;
699
700 /* Get iterator. */
701 it = PyObject_GetIter(iterable);
702 if (it == NULL)
703 return NULL;
704
705 saved = PyList_New(0);
706 if (saved == NULL) {
707 Py_DECREF(it);
708 return NULL;
709 }
710
711 /* create cycleobject structure */
712 lz = (cycleobject *)type->tp_alloc(type, 0);
713 if (lz == NULL) {
714 Py_DECREF(it);
715 Py_DECREF(saved);
716 return NULL;
717 }
718 lz->it = it;
719 lz->saved = saved;
720 lz->firstpass = 0;
721
722 return (PyObject *)lz;
723}
724
725static void
726cycle_dealloc(cycleobject *lz)
727{
728 PyObject_GC_UnTrack(lz);
729 Py_XDECREF(lz->saved);
730 Py_XDECREF(lz->it);
731 lz->ob_type->tp_free(lz);
732}
733
734static int
735cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
736{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000737 Py_VISIT(lz->it);
738 Py_VISIT(lz->saved);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000739 return 0;
740}
741
742static PyObject *
743cycle_next(cycleobject *lz)
744{
745 PyObject *item;
746 PyObject *it;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000747 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000748
749 while (1) {
750 item = PyIter_Next(lz->it);
751 if (item != NULL) {
752 if (!lz->firstpass)
753 PyList_Append(lz->saved, item);
754 return item;
755 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000756 if (PyErr_Occurred()) {
757 if (PyErr_ExceptionMatches(PyExc_StopIteration))
758 PyErr_Clear();
759 else
760 return NULL;
761 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000762 if (PyList_Size(lz->saved) == 0)
763 return NULL;
764 it = PyObject_GetIter(lz->saved);
765 if (it == NULL)
766 return NULL;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000767 tmp = lz->it;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000768 lz->it = it;
769 lz->firstpass = 1;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000770 Py_DECREF(tmp);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000771 }
772}
773
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000774PyDoc_STRVAR(cycle_doc,
775"cycle(iterable) --> cycle object\n\
776\n\
777Return elements from the iterable until it is exhausted.\n\
778Then repeat the sequence indefinitely.");
779
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000780static PyTypeObject cycle_type = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000781 PyObject_HEAD_INIT(NULL)
782 0, /* ob_size */
783 "itertools.cycle", /* tp_name */
784 sizeof(cycleobject), /* tp_basicsize */
785 0, /* tp_itemsize */
786 /* methods */
787 (destructor)cycle_dealloc, /* tp_dealloc */
788 0, /* tp_print */
789 0, /* tp_getattr */
790 0, /* tp_setattr */
791 0, /* tp_compare */
792 0, /* tp_repr */
793 0, /* tp_as_number */
794 0, /* tp_as_sequence */
795 0, /* tp_as_mapping */
796 0, /* tp_hash */
797 0, /* tp_call */
798 0, /* tp_str */
799 PyObject_GenericGetAttr, /* tp_getattro */
800 0, /* tp_setattro */
801 0, /* tp_as_buffer */
802 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
803 Py_TPFLAGS_BASETYPE, /* tp_flags */
804 cycle_doc, /* tp_doc */
805 (traverseproc)cycle_traverse, /* tp_traverse */
806 0, /* tp_clear */
807 0, /* tp_richcompare */
808 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000809 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000810 (iternextfunc)cycle_next, /* tp_iternext */
811 0, /* tp_methods */
812 0, /* tp_members */
813 0, /* tp_getset */
814 0, /* tp_base */
815 0, /* tp_dict */
816 0, /* tp_descr_get */
817 0, /* tp_descr_set */
818 0, /* tp_dictoffset */
819 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000820 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000821 cycle_new, /* tp_new */
822 PyObject_GC_Del, /* tp_free */
823};
824
825
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000826/* dropwhile object **********************************************************/
827
828typedef struct {
829 PyObject_HEAD
830 PyObject *func;
831 PyObject *it;
832 long start;
833} dropwhileobject;
834
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000835static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000836
837static PyObject *
838dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
839{
840 PyObject *func, *seq;
841 PyObject *it;
842 dropwhileobject *lz;
843
Georg Brandl8e932e72007-01-21 10:28:56 +0000844 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000845 return NULL;
846
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000847 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
848 return NULL;
849
850 /* Get iterator. */
851 it = PyObject_GetIter(seq);
852 if (it == NULL)
853 return NULL;
854
855 /* create dropwhileobject structure */
856 lz = (dropwhileobject *)type->tp_alloc(type, 0);
857 if (lz == NULL) {
858 Py_DECREF(it);
859 return NULL;
860 }
861 Py_INCREF(func);
862 lz->func = func;
863 lz->it = it;
864 lz->start = 0;
865
866 return (PyObject *)lz;
867}
868
869static void
870dropwhile_dealloc(dropwhileobject *lz)
871{
872 PyObject_GC_UnTrack(lz);
873 Py_XDECREF(lz->func);
874 Py_XDECREF(lz->it);
875 lz->ob_type->tp_free(lz);
876}
877
878static int
879dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
880{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000881 Py_VISIT(lz->it);
882 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000883 return 0;
884}
885
886static PyObject *
887dropwhile_next(dropwhileobject *lz)
888{
889 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000890 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000891 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000892 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000893
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000894 assert(PyIter_Check(it));
895 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000896 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000897 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000898 if (item == NULL)
899 return NULL;
900 if (lz->start == 1)
901 return item;
902
903 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
904 if (good == NULL) {
905 Py_DECREF(item);
906 return NULL;
907 }
908 ok = PyObject_IsTrue(good);
909 Py_DECREF(good);
910 if (!ok) {
911 lz->start = 1;
912 return item;
913 }
914 Py_DECREF(item);
915 }
916}
917
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000918PyDoc_STRVAR(dropwhile_doc,
919"dropwhile(predicate, iterable) --> dropwhile object\n\
920\n\
921Drop items from the iterable while predicate(item) is true.\n\
922Afterwards, return every element until the iterable is exhausted.");
923
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000924static PyTypeObject dropwhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000925 PyObject_HEAD_INIT(NULL)
926 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000927 "itertools.dropwhile", /* tp_name */
928 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000929 0, /* tp_itemsize */
930 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000931 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000932 0, /* tp_print */
933 0, /* tp_getattr */
934 0, /* tp_setattr */
935 0, /* tp_compare */
936 0, /* tp_repr */
937 0, /* tp_as_number */
938 0, /* tp_as_sequence */
939 0, /* tp_as_mapping */
940 0, /* tp_hash */
941 0, /* tp_call */
942 0, /* tp_str */
943 PyObject_GenericGetAttr, /* tp_getattro */
944 0, /* tp_setattro */
945 0, /* tp_as_buffer */
946 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
947 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000948 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000949 (traverseproc)dropwhile_traverse, /* tp_traverse */
950 0, /* tp_clear */
951 0, /* tp_richcompare */
952 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000953 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000954 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000955 0, /* tp_methods */
956 0, /* tp_members */
957 0, /* tp_getset */
958 0, /* tp_base */
959 0, /* tp_dict */
960 0, /* tp_descr_get */
961 0, /* tp_descr_set */
962 0, /* tp_dictoffset */
963 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000964 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000965 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000966 PyObject_GC_Del, /* tp_free */
967};
968
969
970/* takewhile object **********************************************************/
971
972typedef struct {
973 PyObject_HEAD
974 PyObject *func;
975 PyObject *it;
976 long stop;
977} takewhileobject;
978
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000979static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000980
981static PyObject *
982takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
983{
984 PyObject *func, *seq;
985 PyObject *it;
986 takewhileobject *lz;
987
Georg Brandl8e932e72007-01-21 10:28:56 +0000988 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000989 return NULL;
990
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000991 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
992 return NULL;
993
994 /* Get iterator. */
995 it = PyObject_GetIter(seq);
996 if (it == NULL)
997 return NULL;
998
999 /* create takewhileobject structure */
1000 lz = (takewhileobject *)type->tp_alloc(type, 0);
1001 if (lz == NULL) {
1002 Py_DECREF(it);
1003 return NULL;
1004 }
1005 Py_INCREF(func);
1006 lz->func = func;
1007 lz->it = it;
1008 lz->stop = 0;
1009
1010 return (PyObject *)lz;
1011}
1012
1013static void
1014takewhile_dealloc(takewhileobject *lz)
1015{
1016 PyObject_GC_UnTrack(lz);
1017 Py_XDECREF(lz->func);
1018 Py_XDECREF(lz->it);
1019 lz->ob_type->tp_free(lz);
1020}
1021
1022static int
1023takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1024{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001025 Py_VISIT(lz->it);
1026 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001027 return 0;
1028}
1029
1030static PyObject *
1031takewhile_next(takewhileobject *lz)
1032{
1033 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001034 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001035 long ok;
1036
1037 if (lz->stop == 1)
1038 return NULL;
1039
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001040 assert(PyIter_Check(it));
1041 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001042 if (item == NULL)
1043 return NULL;
1044
1045 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1046 if (good == NULL) {
1047 Py_DECREF(item);
1048 return NULL;
1049 }
1050 ok = PyObject_IsTrue(good);
1051 Py_DECREF(good);
1052 if (ok)
1053 return item;
1054 Py_DECREF(item);
1055 lz->stop = 1;
1056 return NULL;
1057}
1058
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001059PyDoc_STRVAR(takewhile_doc,
1060"takewhile(predicate, iterable) --> takewhile object\n\
1061\n\
1062Return successive entries from an iterable as long as the \n\
1063predicate evaluates to true for each entry.");
1064
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001065static PyTypeObject takewhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001066 PyObject_HEAD_INIT(NULL)
1067 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001068 "itertools.takewhile", /* tp_name */
1069 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001070 0, /* tp_itemsize */
1071 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001072 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001073 0, /* tp_print */
1074 0, /* tp_getattr */
1075 0, /* tp_setattr */
1076 0, /* tp_compare */
1077 0, /* tp_repr */
1078 0, /* tp_as_number */
1079 0, /* tp_as_sequence */
1080 0, /* tp_as_mapping */
1081 0, /* tp_hash */
1082 0, /* tp_call */
1083 0, /* tp_str */
1084 PyObject_GenericGetAttr, /* tp_getattro */
1085 0, /* tp_setattro */
1086 0, /* tp_as_buffer */
1087 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1088 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001089 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001090 (traverseproc)takewhile_traverse, /* tp_traverse */
1091 0, /* tp_clear */
1092 0, /* tp_richcompare */
1093 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001094 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001095 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001096 0, /* tp_methods */
1097 0, /* tp_members */
1098 0, /* tp_getset */
1099 0, /* tp_base */
1100 0, /* tp_dict */
1101 0, /* tp_descr_get */
1102 0, /* tp_descr_set */
1103 0, /* tp_dictoffset */
1104 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001105 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001106 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001107 PyObject_GC_Del, /* tp_free */
1108};
1109
1110
1111/* islice object ************************************************************/
1112
1113typedef struct {
1114 PyObject_HEAD
1115 PyObject *it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001116 Py_ssize_t next;
1117 Py_ssize_t stop;
1118 Py_ssize_t step;
1119 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001120} isliceobject;
1121
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001122static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001123
1124static PyObject *
1125islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1126{
1127 PyObject *seq;
Jack Diederich6c433a92006-05-26 11:15:17 +00001128 Py_ssize_t start=0, stop=-1, step=1;
Raymond Hettingerb2594052004-12-05 09:25:51 +00001129 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001130 Py_ssize_t numargs;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001131 isliceobject *lz;
1132
Georg Brandl8e932e72007-01-21 10:28:56 +00001133 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001134 return NULL;
1135
Raymond Hettingerb2594052004-12-05 09:25:51 +00001136 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001137 return NULL;
1138
Raymond Hettingerb2594052004-12-05 09:25:51 +00001139 numargs = PyTuple_Size(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001140 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001141 if (a1 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001142 stop = PyInt_AsSsize_t(a1);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001143 if (stop == -1) {
1144 if (PyErr_Occurred())
1145 PyErr_Clear();
1146 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001147 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001148 return NULL;
1149 }
1150 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001151 } else {
Raymond Hettingerb2594052004-12-05 09:25:51 +00001152 if (a1 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001153 start = PyInt_AsSsize_t(a1);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001154 if (start == -1 && PyErr_Occurred())
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001155 PyErr_Clear();
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001156 if (a2 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001157 stop = PyInt_AsSsize_t(a2);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001158 if (stop == -1) {
1159 if (PyErr_Occurred())
1160 PyErr_Clear();
1161 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001162 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001163 return NULL;
1164 }
1165 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001166 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001167 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001168 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001169 "Indices for islice() must be non-negative integers or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001170 return NULL;
1171 }
1172
Raymond Hettingerb2594052004-12-05 09:25:51 +00001173 if (a3 != NULL) {
1174 if (a3 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001175 step = PyInt_AsSsize_t(a3);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001176 if (step == -1 && PyErr_Occurred())
1177 PyErr_Clear();
1178 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001179 if (step<1) {
1180 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001181 "Step for islice() must be a positive integer or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001182 return NULL;
1183 }
1184
1185 /* Get iterator. */
1186 it = PyObject_GetIter(seq);
1187 if (it == NULL)
1188 return NULL;
1189
1190 /* create isliceobject structure */
1191 lz = (isliceobject *)type->tp_alloc(type, 0);
1192 if (lz == NULL) {
1193 Py_DECREF(it);
1194 return NULL;
1195 }
1196 lz->it = it;
1197 lz->next = start;
1198 lz->stop = stop;
1199 lz->step = step;
1200 lz->cnt = 0L;
1201
1202 return (PyObject *)lz;
1203}
1204
1205static void
1206islice_dealloc(isliceobject *lz)
1207{
1208 PyObject_GC_UnTrack(lz);
1209 Py_XDECREF(lz->it);
1210 lz->ob_type->tp_free(lz);
1211}
1212
1213static int
1214islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1215{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001216 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001217 return 0;
1218}
1219
1220static PyObject *
1221islice_next(isliceobject *lz)
1222{
1223 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001224 PyObject *it = lz->it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001225 Py_ssize_t oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001226 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001227
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001228 assert(PyIter_Check(it));
1229 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001230 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001231 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001232 if (item == NULL)
1233 return NULL;
1234 Py_DECREF(item);
1235 lz->cnt++;
1236 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001237 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001238 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001239 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001240 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001241 if (item == NULL)
1242 return NULL;
1243 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001244 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001245 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001246 if (lz->next < oldnext) /* Check for overflow */
1247 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001248 return item;
1249}
1250
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001251PyDoc_STRVAR(islice_doc,
1252"islice(iterable, [start,] stop [, step]) --> islice object\n\
1253\n\
1254Return an iterator whose next() method returns selected values from an\n\
1255iterable. If start is specified, will skip all preceding elements;\n\
1256otherwise, start defaults to zero. Step defaults to one. If\n\
1257specified as another value, step determines how many values are \n\
1258skipped between successive calls. Works like a slice() on a list\n\
1259but returns an iterator.");
1260
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001261static PyTypeObject islice_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001262 PyObject_HEAD_INIT(NULL)
1263 0, /* ob_size */
1264 "itertools.islice", /* tp_name */
1265 sizeof(isliceobject), /* tp_basicsize */
1266 0, /* tp_itemsize */
1267 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001268 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001269 0, /* tp_print */
1270 0, /* tp_getattr */
1271 0, /* tp_setattr */
1272 0, /* tp_compare */
1273 0, /* tp_repr */
1274 0, /* tp_as_number */
1275 0, /* tp_as_sequence */
1276 0, /* tp_as_mapping */
1277 0, /* tp_hash */
1278 0, /* tp_call */
1279 0, /* tp_str */
1280 PyObject_GenericGetAttr, /* tp_getattro */
1281 0, /* tp_setattro */
1282 0, /* tp_as_buffer */
1283 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1284 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001285 islice_doc, /* tp_doc */
1286 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001287 0, /* tp_clear */
1288 0, /* tp_richcompare */
1289 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001290 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001291 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001292 0, /* tp_methods */
1293 0, /* tp_members */
1294 0, /* tp_getset */
1295 0, /* tp_base */
1296 0, /* tp_dict */
1297 0, /* tp_descr_get */
1298 0, /* tp_descr_set */
1299 0, /* tp_dictoffset */
1300 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001301 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001302 islice_new, /* tp_new */
1303 PyObject_GC_Del, /* tp_free */
1304};
1305
1306
1307/* starmap object ************************************************************/
1308
1309typedef struct {
1310 PyObject_HEAD
1311 PyObject *func;
1312 PyObject *it;
1313} starmapobject;
1314
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001315static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001316
1317static PyObject *
1318starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1319{
1320 PyObject *func, *seq;
1321 PyObject *it;
1322 starmapobject *lz;
1323
Georg Brandl8e932e72007-01-21 10:28:56 +00001324 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001325 return NULL;
1326
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001327 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1328 return NULL;
1329
1330 /* Get iterator. */
1331 it = PyObject_GetIter(seq);
1332 if (it == NULL)
1333 return NULL;
1334
1335 /* create starmapobject structure */
1336 lz = (starmapobject *)type->tp_alloc(type, 0);
1337 if (lz == NULL) {
1338 Py_DECREF(it);
1339 return NULL;
1340 }
1341 Py_INCREF(func);
1342 lz->func = func;
1343 lz->it = it;
1344
1345 return (PyObject *)lz;
1346}
1347
1348static void
1349starmap_dealloc(starmapobject *lz)
1350{
1351 PyObject_GC_UnTrack(lz);
1352 Py_XDECREF(lz->func);
1353 Py_XDECREF(lz->it);
1354 lz->ob_type->tp_free(lz);
1355}
1356
1357static int
1358starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1359{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001360 Py_VISIT(lz->it);
1361 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001362 return 0;
1363}
1364
1365static PyObject *
1366starmap_next(starmapobject *lz)
1367{
1368 PyObject *args;
1369 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001370 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001371
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001372 assert(PyIter_Check(it));
1373 args = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001374 if (args == NULL)
1375 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001376 if (!PyTuple_CheckExact(args)) {
1377 Py_DECREF(args);
1378 PyErr_SetString(PyExc_TypeError,
1379 "iterator must return a tuple");
1380 return NULL;
1381 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001382 result = PyObject_Call(lz->func, args, NULL);
1383 Py_DECREF(args);
1384 return result;
1385}
1386
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001387PyDoc_STRVAR(starmap_doc,
1388"starmap(function, sequence) --> starmap object\n\
1389\n\
1390Return an iterator whose values are returned from the function evaluated\n\
1391with a argument tuple taken from the given sequence.");
1392
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001393static PyTypeObject starmap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001394 PyObject_HEAD_INIT(NULL)
1395 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001396 "itertools.starmap", /* tp_name */
1397 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001398 0, /* tp_itemsize */
1399 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001400 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001401 0, /* tp_print */
1402 0, /* tp_getattr */
1403 0, /* tp_setattr */
1404 0, /* tp_compare */
1405 0, /* tp_repr */
1406 0, /* tp_as_number */
1407 0, /* tp_as_sequence */
1408 0, /* tp_as_mapping */
1409 0, /* tp_hash */
1410 0, /* tp_call */
1411 0, /* tp_str */
1412 PyObject_GenericGetAttr, /* tp_getattro */
1413 0, /* tp_setattro */
1414 0, /* tp_as_buffer */
1415 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1416 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001417 starmap_doc, /* tp_doc */
1418 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001419 0, /* tp_clear */
1420 0, /* tp_richcompare */
1421 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001422 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001423 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001424 0, /* tp_methods */
1425 0, /* tp_members */
1426 0, /* tp_getset */
1427 0, /* tp_base */
1428 0, /* tp_dict */
1429 0, /* tp_descr_get */
1430 0, /* tp_descr_set */
1431 0, /* tp_dictoffset */
1432 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001433 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001434 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001435 PyObject_GC_Del, /* tp_free */
1436};
1437
1438
1439/* imap object ************************************************************/
1440
1441typedef struct {
1442 PyObject_HEAD
1443 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001444 PyObject *func;
1445} imapobject;
1446
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001447static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001448
1449static PyObject *
1450imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1451{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001452 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001453 imapobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001454 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001455
Georg Brandl8e932e72007-01-21 10:28:56 +00001456 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001457 return NULL;
1458
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459 numargs = PyTuple_Size(args);
1460 if (numargs < 2) {
1461 PyErr_SetString(PyExc_TypeError,
1462 "imap() must have at least two arguments.");
1463 return NULL;
1464 }
1465
1466 iters = PyTuple_New(numargs-1);
1467 if (iters == NULL)
1468 return NULL;
1469
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001470 for (i=1 ; i<numargs ; i++) {
1471 /* Get iterator. */
1472 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1473 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001474 Py_DECREF(iters);
1475 return NULL;
1476 }
1477 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001478 }
1479
1480 /* create imapobject structure */
1481 lz = (imapobject *)type->tp_alloc(type, 0);
1482 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001483 Py_DECREF(iters);
1484 return NULL;
1485 }
1486 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001487 func = PyTuple_GET_ITEM(args, 0);
1488 Py_INCREF(func);
1489 lz->func = func;
1490
1491 return (PyObject *)lz;
1492}
1493
1494static void
1495imap_dealloc(imapobject *lz)
1496{
1497 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001498 Py_XDECREF(lz->iters);
1499 Py_XDECREF(lz->func);
1500 lz->ob_type->tp_free(lz);
1501}
1502
1503static int
1504imap_traverse(imapobject *lz, visitproc visit, void *arg)
1505{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001506 Py_VISIT(lz->iters);
1507 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001508 return 0;
1509}
1510
Raymond Hettinger2012f172003-02-07 05:32:58 +00001511/*
1512imap() is an iterator version of __builtins__.map() except that it does
1513not have the None fill-in feature. That was intentionally left out for
1514the following reasons:
1515
1516 1) Itertools are designed to be easily combined and chained together.
1517 Having all tools stop with the shortest input is a unifying principle
1518 that makes it easier to combine finite iterators (supplying data) with
1519 infinite iterators like count() and repeat() (for supplying sequential
1520 or constant arguments to a function).
1521
1522 2) In typical use cases for combining itertools, having one finite data
1523 supplier run out before another is likely to be an error condition which
1524 should not pass silently by automatically supplying None.
1525
1526 3) The use cases for automatic None fill-in are rare -- not many functions
1527 do something useful when a parameter suddenly switches type and becomes
1528 None.
1529
1530 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001531 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001532
1533 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1534*/
1535
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001536static PyObject *
1537imap_next(imapobject *lz)
1538{
1539 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001540 PyObject *argtuple;
1541 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001542 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001543
1544 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001545 argtuple = PyTuple_New(numargs);
1546 if (argtuple == NULL)
1547 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001548
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001549 for (i=0 ; i<numargs ; i++) {
1550 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1551 if (val == NULL) {
1552 Py_DECREF(argtuple);
1553 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001554 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001555 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001556 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001557 if (lz->func == Py_None)
1558 return argtuple;
1559 result = PyObject_Call(lz->func, argtuple, NULL);
1560 Py_DECREF(argtuple);
1561 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001562}
1563
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001564PyDoc_STRVAR(imap_doc,
1565"imap(func, *iterables) --> imap object\n\
1566\n\
1567Make an iterator that computes the function using arguments from\n\
1568each of the iterables. Like map() except that it returns\n\
1569an iterator instead of a list and that it stops when the shortest\n\
1570iterable is exhausted instead of filling in None for shorter\n\
1571iterables.");
1572
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001573static PyTypeObject imap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001574 PyObject_HEAD_INIT(NULL)
1575 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001576 "itertools.imap", /* tp_name */
1577 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001578 0, /* tp_itemsize */
1579 /* methods */
1580 (destructor)imap_dealloc, /* tp_dealloc */
1581 0, /* tp_print */
1582 0, /* tp_getattr */
1583 0, /* tp_setattr */
1584 0, /* tp_compare */
1585 0, /* tp_repr */
1586 0, /* tp_as_number */
1587 0, /* tp_as_sequence */
1588 0, /* tp_as_mapping */
1589 0, /* tp_hash */
1590 0, /* tp_call */
1591 0, /* tp_str */
1592 PyObject_GenericGetAttr, /* tp_getattro */
1593 0, /* tp_setattro */
1594 0, /* tp_as_buffer */
1595 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1596 Py_TPFLAGS_BASETYPE, /* tp_flags */
1597 imap_doc, /* tp_doc */
1598 (traverseproc)imap_traverse, /* tp_traverse */
1599 0, /* tp_clear */
1600 0, /* tp_richcompare */
1601 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001602 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001603 (iternextfunc)imap_next, /* tp_iternext */
1604 0, /* tp_methods */
1605 0, /* tp_members */
1606 0, /* tp_getset */
1607 0, /* tp_base */
1608 0, /* tp_dict */
1609 0, /* tp_descr_get */
1610 0, /* tp_descr_set */
1611 0, /* tp_dictoffset */
1612 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001613 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001614 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001615 PyObject_GC_Del, /* tp_free */
1616};
1617
1618
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001619/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001620
1621typedef struct {
1622 PyObject_HEAD
Jack Diederich6c433a92006-05-26 11:15:17 +00001623 Py_ssize_t tuplesize;
1624 Py_ssize_t iternum; /* which iterator is active */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001625 PyObject *ittuple; /* tuple of iterators */
1626} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001627
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001628static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001629
1630static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001631chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001632{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001633 chainobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001634 Py_ssize_t tuplesize = PySequence_Length(args);
Jack Diederich6c433a92006-05-26 11:15:17 +00001635 Py_ssize_t i;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001636 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001637
Georg Brandl8e932e72007-01-21 10:28:56 +00001638 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001639 return NULL;
1640
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001641 /* obtain iterators */
1642 assert(PyTuple_Check(args));
1643 ittuple = PyTuple_New(tuplesize);
Neal Norwitz2f3136b2006-05-27 05:18:57 +00001644 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001645 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001646 for (i=0; i < tuplesize; ++i) {
1647 PyObject *item = PyTuple_GET_ITEM(args, i);
1648 PyObject *it = PyObject_GetIter(item);
1649 if (it == NULL) {
1650 if (PyErr_ExceptionMatches(PyExc_TypeError))
1651 PyErr_Format(PyExc_TypeError,
Neal Norwitz2f3136b2006-05-27 05:18:57 +00001652 "chain argument #%zd must support iteration",
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001653 i+1);
1654 Py_DECREF(ittuple);
1655 return NULL;
1656 }
1657 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001658 }
1659
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001660 /* create chainobject structure */
1661 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001662 if (lz == NULL) {
1663 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001664 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001665 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001666
1667 lz->ittuple = ittuple;
1668 lz->iternum = 0;
1669 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001670
1671 return (PyObject *)lz;
1672}
1673
1674static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001675chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001676{
1677 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001678 Py_XDECREF(lz->ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001679 lz->ob_type->tp_free(lz);
1680}
1681
Raymond Hettinger2012f172003-02-07 05:32:58 +00001682static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001683chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001684{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001685 Py_VISIT(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001686 return 0;
1687}
1688
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001689static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001690chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001691{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001692 PyObject *it;
1693 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001694
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001695 while (lz->iternum < lz->tuplesize) {
1696 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1697 item = PyIter_Next(it);
1698 if (item != NULL)
1699 return item;
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001700 if (PyErr_Occurred()) {
1701 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1702 PyErr_Clear();
1703 else
1704 return NULL;
1705 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001706 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001707 }
1708 return NULL;
1709}
1710
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001711PyDoc_STRVAR(chain_doc,
1712"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001713\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001714Return a chain object whose .next() method returns elements from the\n\
1715first iterable until it is exhausted, then elements from the next\n\
1716iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001717
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001718static PyTypeObject chain_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001719 PyObject_HEAD_INIT(NULL)
1720 0, /* ob_size */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001721 "itertools.chain", /* tp_name */
1722 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001723 0, /* tp_itemsize */
1724 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001725 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001726 0, /* tp_print */
1727 0, /* tp_getattr */
1728 0, /* tp_setattr */
1729 0, /* tp_compare */
1730 0, /* tp_repr */
1731 0, /* tp_as_number */
1732 0, /* tp_as_sequence */
1733 0, /* tp_as_mapping */
1734 0, /* tp_hash */
1735 0, /* tp_call */
1736 0, /* tp_str */
1737 PyObject_GenericGetAttr, /* tp_getattro */
1738 0, /* tp_setattro */
1739 0, /* tp_as_buffer */
1740 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1741 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001742 chain_doc, /* tp_doc */
1743 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001744 0, /* tp_clear */
1745 0, /* tp_richcompare */
1746 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001747 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001748 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001749 0, /* tp_methods */
1750 0, /* tp_members */
1751 0, /* tp_getset */
1752 0, /* tp_base */
1753 0, /* tp_dict */
1754 0, /* tp_descr_get */
1755 0, /* tp_descr_set */
1756 0, /* tp_dictoffset */
1757 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001758 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001759 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001760 PyObject_GC_Del, /* tp_free */
1761};
1762
1763
1764/* ifilter object ************************************************************/
1765
1766typedef struct {
1767 PyObject_HEAD
1768 PyObject *func;
1769 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001770} ifilterobject;
1771
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001772static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001773
1774static PyObject *
1775ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1776{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001777 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001778 PyObject *it;
1779 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001780
Georg Brandl8e932e72007-01-21 10:28:56 +00001781 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001782 return NULL;
1783
Raymond Hettinger60eca932003-02-09 06:40:58 +00001784 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001785 return NULL;
1786
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001787 /* Get iterator. */
1788 it = PyObject_GetIter(seq);
1789 if (it == NULL)
1790 return NULL;
1791
1792 /* create ifilterobject structure */
1793 lz = (ifilterobject *)type->tp_alloc(type, 0);
1794 if (lz == NULL) {
1795 Py_DECREF(it);
1796 return NULL;
1797 }
1798 Py_INCREF(func);
1799 lz->func = func;
1800 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001801
1802 return (PyObject *)lz;
1803}
1804
1805static void
1806ifilter_dealloc(ifilterobject *lz)
1807{
1808 PyObject_GC_UnTrack(lz);
1809 Py_XDECREF(lz->func);
1810 Py_XDECREF(lz->it);
1811 lz->ob_type->tp_free(lz);
1812}
1813
1814static int
1815ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1816{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001817 Py_VISIT(lz->it);
1818 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001819 return 0;
1820}
1821
1822static PyObject *
1823ifilter_next(ifilterobject *lz)
1824{
1825 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001826 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001827 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001828 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001829
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001830 assert(PyIter_Check(it));
1831 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001832 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001833 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001834 if (item == NULL)
1835 return NULL;
1836
1837 if (lz->func == Py_None) {
1838 ok = PyObject_IsTrue(item);
1839 } else {
1840 PyObject *good;
1841 good = PyObject_CallFunctionObjArgs(lz->func,
1842 item, NULL);
1843 if (good == NULL) {
1844 Py_DECREF(item);
1845 return NULL;
1846 }
1847 ok = PyObject_IsTrue(good);
1848 Py_DECREF(good);
1849 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00001850 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001851 return item;
1852 Py_DECREF(item);
1853 }
1854}
1855
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001856PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001857"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001858\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001859Return those items of sequence for which function(item) is true.\n\
1860If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001861
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001862static PyTypeObject ifilter_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001863 PyObject_HEAD_INIT(NULL)
1864 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001865 "itertools.ifilter", /* tp_name */
1866 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001867 0, /* tp_itemsize */
1868 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001869 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001870 0, /* tp_print */
1871 0, /* tp_getattr */
1872 0, /* tp_setattr */
1873 0, /* tp_compare */
1874 0, /* tp_repr */
1875 0, /* tp_as_number */
1876 0, /* tp_as_sequence */
1877 0, /* tp_as_mapping */
1878 0, /* tp_hash */
1879 0, /* tp_call */
1880 0, /* tp_str */
1881 PyObject_GenericGetAttr, /* tp_getattro */
1882 0, /* tp_setattro */
1883 0, /* tp_as_buffer */
1884 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1885 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001886 ifilter_doc, /* tp_doc */
1887 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001888 0, /* tp_clear */
1889 0, /* tp_richcompare */
1890 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001891 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001892 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001893 0, /* tp_methods */
1894 0, /* tp_members */
1895 0, /* tp_getset */
1896 0, /* tp_base */
1897 0, /* tp_dict */
1898 0, /* tp_descr_get */
1899 0, /* tp_descr_set */
1900 0, /* tp_dictoffset */
1901 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001902 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001903 ifilter_new, /* tp_new */
1904 PyObject_GC_Del, /* tp_free */
1905};
1906
1907
1908/* ifilterfalse object ************************************************************/
1909
1910typedef struct {
1911 PyObject_HEAD
1912 PyObject *func;
1913 PyObject *it;
1914} ifilterfalseobject;
1915
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001916static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001917
1918static PyObject *
1919ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1920{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00001921 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001922 PyObject *it;
1923 ifilterfalseobject *lz;
1924
Georg Brandl8e932e72007-01-21 10:28:56 +00001925 if (type == &ifilterfalse_type &&
1926 !_PyArg_NoKeywords("ifilterfalse()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001927 return NULL;
1928
Raymond Hettinger60eca932003-02-09 06:40:58 +00001929 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1930 return NULL;
1931
1932 /* Get iterator. */
1933 it = PyObject_GetIter(seq);
1934 if (it == NULL)
1935 return NULL;
1936
1937 /* create ifilterfalseobject structure */
1938 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1939 if (lz == NULL) {
1940 Py_DECREF(it);
1941 return NULL;
1942 }
1943 Py_INCREF(func);
1944 lz->func = func;
1945 lz->it = it;
1946
1947 return (PyObject *)lz;
1948}
1949
1950static void
1951ifilterfalse_dealloc(ifilterfalseobject *lz)
1952{
1953 PyObject_GC_UnTrack(lz);
1954 Py_XDECREF(lz->func);
1955 Py_XDECREF(lz->it);
1956 lz->ob_type->tp_free(lz);
1957}
1958
1959static int
1960ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1961{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001962 Py_VISIT(lz->it);
1963 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001964 return 0;
1965}
1966
1967static PyObject *
1968ifilterfalse_next(ifilterfalseobject *lz)
1969{
1970 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001971 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001972 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001973 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001974
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001975 assert(PyIter_Check(it));
1976 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001977 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001978 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001979 if (item == NULL)
1980 return NULL;
1981
1982 if (lz->func == Py_None) {
1983 ok = PyObject_IsTrue(item);
1984 } else {
1985 PyObject *good;
1986 good = PyObject_CallFunctionObjArgs(lz->func,
1987 item, NULL);
1988 if (good == NULL) {
1989 Py_DECREF(item);
1990 return NULL;
1991 }
1992 ok = PyObject_IsTrue(good);
1993 Py_DECREF(good);
1994 }
1995 if (!ok)
1996 return item;
1997 Py_DECREF(item);
1998 }
1999}
2000
Raymond Hettinger60eca932003-02-09 06:40:58 +00002001PyDoc_STRVAR(ifilterfalse_doc,
2002"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
2003\n\
2004Return those items of sequence for which function(item) is false.\n\
2005If function is None, return the items that are false.");
2006
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002007static PyTypeObject ifilterfalse_type = {
Raymond Hettinger60eca932003-02-09 06:40:58 +00002008 PyObject_HEAD_INIT(NULL)
2009 0, /* ob_size */
2010 "itertools.ifilterfalse", /* tp_name */
2011 sizeof(ifilterfalseobject), /* tp_basicsize */
2012 0, /* tp_itemsize */
2013 /* methods */
2014 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
2015 0, /* tp_print */
2016 0, /* tp_getattr */
2017 0, /* tp_setattr */
2018 0, /* tp_compare */
2019 0, /* tp_repr */
2020 0, /* tp_as_number */
2021 0, /* tp_as_sequence */
2022 0, /* tp_as_mapping */
2023 0, /* tp_hash */
2024 0, /* tp_call */
2025 0, /* tp_str */
2026 PyObject_GenericGetAttr, /* tp_getattro */
2027 0, /* tp_setattro */
2028 0, /* tp_as_buffer */
2029 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2030 Py_TPFLAGS_BASETYPE, /* tp_flags */
2031 ifilterfalse_doc, /* tp_doc */
2032 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
2033 0, /* tp_clear */
2034 0, /* tp_richcompare */
2035 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002036 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002037 (iternextfunc)ifilterfalse_next, /* tp_iternext */
2038 0, /* tp_methods */
2039 0, /* tp_members */
2040 0, /* tp_getset */
2041 0, /* tp_base */
2042 0, /* tp_dict */
2043 0, /* tp_descr_get */
2044 0, /* tp_descr_set */
2045 0, /* tp_dictoffset */
2046 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002047 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002048 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002049 PyObject_GC_Del, /* tp_free */
2050};
2051
2052
2053/* count object ************************************************************/
2054
2055typedef struct {
2056 PyObject_HEAD
Jack Diederich6c433a92006-05-26 11:15:17 +00002057 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002058} countobject;
2059
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002060static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002061
2062static PyObject *
2063count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2064{
2065 countobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00002066 Py_ssize_t cnt = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002067
Georg Brandl8e932e72007-01-21 10:28:56 +00002068 if (type == &count_type && !_PyArg_NoKeywords("count()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002069 return NULL;
2070
Jack Diederich6c433a92006-05-26 11:15:17 +00002071 if (!PyArg_ParseTuple(args, "|n:count", &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002072 return NULL;
2073
2074 /* create countobject structure */
2075 lz = (countobject *)PyObject_New(countobject, &count_type);
2076 if (lz == NULL)
2077 return NULL;
2078 lz->cnt = cnt;
2079
2080 return (PyObject *)lz;
2081}
2082
2083static PyObject *
2084count_next(countobject *lz)
2085{
Kristján Valur Jónssonf4601d82007-05-07 18:30:48 +00002086 if (lz->cnt == PY_SSIZE_T_MAX) {
Raymond Hettinger9f0e1ea2007-02-07 23:57:05 +00002087 PyErr_SetString(PyExc_OverflowError,
Kristján Valur Jónssonf4601d82007-05-07 18:30:48 +00002088 "cannot count beyond PY_SSIZE_T_MAX");
Raymond Hettinger9f0e1ea2007-02-07 23:57:05 +00002089 return NULL;
2090 }
Jack Diederich187e6482006-09-21 18:32:11 +00002091 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002092}
2093
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002094static PyObject *
2095count_repr(countobject *lz)
2096{
Jack Diederich6c433a92006-05-26 11:15:17 +00002097 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002098}
2099
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002100PyDoc_STRVAR(count_doc,
2101"count([firstval]) --> count object\n\
2102\n\
2103Return a count object whose .next() method returns consecutive\n\
2104integers starting from zero or, if specified, from firstval.");
2105
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002106static PyTypeObject count_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002107 PyObject_HEAD_INIT(NULL)
2108 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002109 "itertools.count", /* tp_name */
2110 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002111 0, /* tp_itemsize */
2112 /* methods */
2113 (destructor)PyObject_Del, /* tp_dealloc */
2114 0, /* tp_print */
2115 0, /* tp_getattr */
2116 0, /* tp_setattr */
2117 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002118 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002119 0, /* tp_as_number */
2120 0, /* tp_as_sequence */
2121 0, /* tp_as_mapping */
2122 0, /* tp_hash */
2123 0, /* tp_call */
2124 0, /* tp_str */
2125 PyObject_GenericGetAttr, /* tp_getattro */
2126 0, /* tp_setattro */
2127 0, /* tp_as_buffer */
2128 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002129 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002130 0, /* tp_traverse */
2131 0, /* tp_clear */
2132 0, /* tp_richcompare */
2133 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002134 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002135 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002136 0, /* tp_methods */
2137 0, /* tp_members */
2138 0, /* tp_getset */
2139 0, /* tp_base */
2140 0, /* tp_dict */
2141 0, /* tp_descr_get */
2142 0, /* tp_descr_set */
2143 0, /* tp_dictoffset */
2144 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002145 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002146 count_new, /* tp_new */
2147};
2148
2149
2150/* izip object ************************************************************/
2151
2152#include "Python.h"
2153
2154typedef struct {
2155 PyObject_HEAD
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002156 Py_ssize_t tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002157 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002158 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002159} izipobject;
2160
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002161static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002162
2163static PyObject *
2164izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2165{
2166 izipobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00002167 Py_ssize_t i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002168 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002169 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002170 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002171
Georg Brandl8e932e72007-01-21 10:28:56 +00002172 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002173 return NULL;
2174
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002175 /* args must be a tuple */
2176 assert(PyTuple_Check(args));
2177
2178 /* obtain iterators */
2179 ittuple = PyTuple_New(tuplesize);
Neal Norwitz2f3136b2006-05-27 05:18:57 +00002180 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002181 return NULL;
2182 for (i=0; i < tuplesize; ++i) {
2183 PyObject *item = PyTuple_GET_ITEM(args, i);
2184 PyObject *it = PyObject_GetIter(item);
2185 if (it == NULL) {
2186 if (PyErr_ExceptionMatches(PyExc_TypeError))
2187 PyErr_Format(PyExc_TypeError,
Neal Norwitz2f3136b2006-05-27 05:18:57 +00002188 "izip argument #%zd must support iteration",
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002189 i+1);
2190 Py_DECREF(ittuple);
2191 return NULL;
2192 }
2193 PyTuple_SET_ITEM(ittuple, i, it);
2194 }
2195
Raymond Hettinger2012f172003-02-07 05:32:58 +00002196 /* create a result holder */
2197 result = PyTuple_New(tuplesize);
2198 if (result == NULL) {
2199 Py_DECREF(ittuple);
2200 return NULL;
2201 }
2202 for (i=0 ; i < tuplesize ; i++) {
2203 Py_INCREF(Py_None);
2204 PyTuple_SET_ITEM(result, i, Py_None);
2205 }
2206
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002207 /* create izipobject structure */
2208 lz = (izipobject *)type->tp_alloc(type, 0);
2209 if (lz == NULL) {
2210 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002211 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002212 return NULL;
2213 }
2214 lz->ittuple = ittuple;
2215 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002216 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002217
2218 return (PyObject *)lz;
2219}
2220
2221static void
2222izip_dealloc(izipobject *lz)
2223{
2224 PyObject_GC_UnTrack(lz);
2225 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002226 Py_XDECREF(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002227 lz->ob_type->tp_free(lz);
2228}
2229
2230static int
2231izip_traverse(izipobject *lz, visitproc visit, void *arg)
2232{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002233 Py_VISIT(lz->ittuple);
2234 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002235 return 0;
2236}
2237
2238static PyObject *
2239izip_next(izipobject *lz)
2240{
Jack Diederich6c433a92006-05-26 11:15:17 +00002241 Py_ssize_t i;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002242 Py_ssize_t tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002243 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002244 PyObject *it;
2245 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002246 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002247
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002248 if (tuplesize == 0)
2249 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002250 if (result->ob_refcnt == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002251 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002252 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002253 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002254 assert(PyIter_Check(it));
2255 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002256 if (item == NULL) {
2257 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002258 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002259 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002260 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002261 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002262 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002263 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002264 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002265 result = PyTuple_New(tuplesize);
2266 if (result == NULL)
2267 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002268 for (i=0 ; i < tuplesize ; i++) {
2269 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002270 assert(PyIter_Check(it));
2271 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002272 if (item == NULL) {
2273 Py_DECREF(result);
2274 return NULL;
2275 }
2276 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002277 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002278 }
2279 return result;
2280}
2281
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002282PyDoc_STRVAR(izip_doc,
2283"izip(iter1 [,iter2 [...]]) --> izip object\n\
2284\n\
2285Return a izip object whose .next() method returns a tuple where\n\
2286the i-th element comes from the i-th iterable argument. The .next()\n\
2287method continues until the shortest iterable in the argument sequence\n\
2288is exhausted and then it raises StopIteration. Works like the zip()\n\
2289function but consumes less memory by returning an iterator instead of\n\
2290a list.");
2291
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002292static PyTypeObject izip_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002293 PyObject_HEAD_INIT(NULL)
2294 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002295 "itertools.izip", /* tp_name */
2296 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002297 0, /* tp_itemsize */
2298 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002299 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002300 0, /* tp_print */
2301 0, /* tp_getattr */
2302 0, /* tp_setattr */
2303 0, /* tp_compare */
2304 0, /* tp_repr */
2305 0, /* tp_as_number */
2306 0, /* tp_as_sequence */
2307 0, /* tp_as_mapping */
2308 0, /* tp_hash */
2309 0, /* tp_call */
2310 0, /* tp_str */
2311 PyObject_GenericGetAttr, /* tp_getattro */
2312 0, /* tp_setattro */
2313 0, /* tp_as_buffer */
2314 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2315 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002316 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002317 (traverseproc)izip_traverse, /* tp_traverse */
2318 0, /* tp_clear */
2319 0, /* tp_richcompare */
2320 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002321 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002322 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002323 0, /* tp_methods */
2324 0, /* tp_members */
2325 0, /* tp_getset */
2326 0, /* tp_base */
2327 0, /* tp_dict */
2328 0, /* tp_descr_get */
2329 0, /* tp_descr_set */
2330 0, /* tp_dictoffset */
2331 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002332 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002333 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002334 PyObject_GC_Del, /* tp_free */
2335};
2336
2337
2338/* repeat object ************************************************************/
2339
2340typedef struct {
2341 PyObject_HEAD
2342 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00002343 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002344} repeatobject;
2345
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002346static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002347
2348static PyObject *
2349repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2350{
2351 repeatobject *ro;
2352 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00002353 Py_ssize_t cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002354
Georg Brandl8e932e72007-01-21 10:28:56 +00002355 if (type == &repeat_type && !_PyArg_NoKeywords("repeat()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002356 return NULL;
2357
Jack Diederich6c433a92006-05-26 11:15:17 +00002358 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002359 return NULL;
2360
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002361 if (PyTuple_Size(args) == 2 && cnt < 0)
2362 cnt = 0;
2363
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002364 ro = (repeatobject *)type->tp_alloc(type, 0);
2365 if (ro == NULL)
2366 return NULL;
2367 Py_INCREF(element);
2368 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002369 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002370 return (PyObject *)ro;
2371}
2372
2373static void
2374repeat_dealloc(repeatobject *ro)
2375{
2376 PyObject_GC_UnTrack(ro);
2377 Py_XDECREF(ro->element);
2378 ro->ob_type->tp_free(ro);
2379}
2380
2381static int
2382repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2383{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002384 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002385 return 0;
2386}
2387
2388static PyObject *
2389repeat_next(repeatobject *ro)
2390{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002391 if (ro->cnt == 0)
2392 return NULL;
2393 if (ro->cnt > 0)
2394 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002395 Py_INCREF(ro->element);
2396 return ro->element;
2397}
2398
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002399static PyObject *
2400repeat_repr(repeatobject *ro)
2401{
2402 PyObject *result, *objrepr;
2403
2404 objrepr = PyObject_Repr(ro->element);
2405 if (objrepr == NULL)
2406 return NULL;
2407
2408 if (ro->cnt == -1)
2409 result = PyString_FromFormat("repeat(%s)",
2410 PyString_AS_STRING(objrepr));
2411 else
Jack Diederich6c433a92006-05-26 11:15:17 +00002412 result = PyString_FromFormat("repeat(%s, %zd)",
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002413 PyString_AS_STRING(objrepr), ro->cnt);
2414 Py_DECREF(objrepr);
2415 return result;
2416}
2417
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002418static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002419repeat_len(repeatobject *ro)
2420{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002421 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002422 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002423 return NULL;
2424 }
Jack Diederich6c433a92006-05-26 11:15:17 +00002425 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002426}
2427
Armin Rigof5b3e362006-02-11 21:32:43 +00002428PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002429
2430static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00002431 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002432 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002433};
2434
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002435PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002436"repeat(element [,times]) -> create an iterator which returns the element\n\
2437for the specified number of times. If not specified, returns the element\n\
2438endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002439
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002440static PyTypeObject repeat_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002441 PyObject_HEAD_INIT(NULL)
2442 0, /* ob_size */
2443 "itertools.repeat", /* tp_name */
2444 sizeof(repeatobject), /* tp_basicsize */
2445 0, /* tp_itemsize */
2446 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002447 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002448 0, /* tp_print */
2449 0, /* tp_getattr */
2450 0, /* tp_setattr */
2451 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002452 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002453 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002454 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002455 0, /* tp_as_mapping */
2456 0, /* tp_hash */
2457 0, /* tp_call */
2458 0, /* tp_str */
2459 PyObject_GenericGetAttr, /* tp_getattro */
2460 0, /* tp_setattro */
2461 0, /* tp_as_buffer */
2462 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2463 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002464 repeat_doc, /* tp_doc */
2465 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002466 0, /* tp_clear */
2467 0, /* tp_richcompare */
2468 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002469 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002470 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002471 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002472 0, /* tp_members */
2473 0, /* tp_getset */
2474 0, /* tp_base */
2475 0, /* tp_dict */
2476 0, /* tp_descr_get */
2477 0, /* tp_descr_set */
2478 0, /* tp_dictoffset */
2479 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002480 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002481 repeat_new, /* tp_new */
2482 PyObject_GC_Del, /* tp_free */
2483};
2484
2485
2486/* module level code ********************************************************/
2487
2488PyDoc_STRVAR(module_doc,
2489"Functional tools for creating and using iterators.\n\
2490\n\
2491Infinite iterators:\n\
2492count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002493cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002494repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002495\n\
2496Iterators terminating on the shortest input sequence:\n\
2497izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002498ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2499ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002500islice(seq, [start,] stop [, step]) --> elements from\n\
2501 seq[start:stop:step]\n\
2502imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2503starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00002504tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002505chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002506takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2507dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002508groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002509");
2510
2511
Raymond Hettingerad983e72003-11-12 14:32:26 +00002512static PyMethodDef module_methods[] = {
2513 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2514 {NULL, NULL} /* sentinel */
2515};
2516
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002517PyMODINIT_FUNC
2518inititertools(void)
2519{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002520 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002521 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002522 char *name;
2523 PyTypeObject *typelist[] = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002524 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002525 &dropwhile_type,
2526 &takewhile_type,
2527 &islice_type,
2528 &starmap_type,
2529 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002530 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002531 &ifilter_type,
2532 &ifilterfalse_type,
2533 &count_type,
2534 &izip_type,
2535 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002536 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002537 NULL
2538 };
2539
Skip Montanarof3938fd2004-02-10 20:27:40 +00002540 teedataobject_type.ob_type = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00002541 m = Py_InitModule3("itertools", module_methods, module_doc);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00002542 if (m == NULL)
2543 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002544
Raymond Hettinger60eca932003-02-09 06:40:58 +00002545 for (i=0 ; typelist[i] != NULL ; i++) {
2546 if (PyType_Ready(typelist[i]) < 0)
2547 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002548 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002549 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002550 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002551 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002552 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00002553
2554 if (PyType_Ready(&teedataobject_type) < 0)
2555 return;
2556 if (PyType_Ready(&tee_type) < 0)
2557 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002558 if (PyType_Ready(&_grouper_type) < 0)
2559 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002560}