blob: 47db7affb0598288d96edc8801e55120e9172c9e [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
Antoine Pitrouc83ea132010-05-09 14:46:46 +00005/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00006 by Raymond D. Hettinger <python@rcn.com>
Raymond Hettinger96ef8112003-02-01 00:10:11 +00007*/
8
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00009
10/* groupby object ***********************************************************/
11
12typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +000013 PyObject_HEAD
14 PyObject *it;
15 PyObject *keyfunc;
16 PyObject *tgtkey;
17 PyObject *currkey;
18 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000019} groupbyobject;
20
21static PyTypeObject groupby_type;
22static PyObject *_grouper_create(groupbyobject *, PyObject *);
23
24static PyObject *
25groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
26{
Antoine Pitrouc83ea132010-05-09 14:46:46 +000027 static char *kwargs[] = {"iterable", "key", NULL};
28 groupbyobject *gbo;
29 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000030
Antoine Pitrouc83ea132010-05-09 14:46:46 +000031 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
32 &it, &keyfunc))
33 return NULL;
34
35 gbo = (groupbyobject *)type->tp_alloc(type, 0);
36 if (gbo == NULL)
37 return NULL;
38 gbo->tgtkey = NULL;
39 gbo->currkey = NULL;
40 gbo->currvalue = NULL;
41 gbo->keyfunc = keyfunc;
42 Py_INCREF(keyfunc);
43 gbo->it = PyObject_GetIter(it);
44 if (gbo->it == NULL) {
45 Py_DECREF(gbo);
46 return NULL;
47 }
48 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000049}
50
51static void
52groupby_dealloc(groupbyobject *gbo)
53{
Antoine Pitrouc83ea132010-05-09 14:46:46 +000054 PyObject_GC_UnTrack(gbo);
55 Py_XDECREF(gbo->it);
56 Py_XDECREF(gbo->keyfunc);
57 Py_XDECREF(gbo->tgtkey);
58 Py_XDECREF(gbo->currkey);
59 Py_XDECREF(gbo->currvalue);
60 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000061}
62
63static int
64groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
65{
Antoine Pitrouc83ea132010-05-09 14:46:46 +000066 Py_VISIT(gbo->it);
67 Py_VISIT(gbo->keyfunc);
68 Py_VISIT(gbo->tgtkey);
69 Py_VISIT(gbo->currkey);
70 Py_VISIT(gbo->currvalue);
71 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000072}
73
Serhiy Storchakad0b9dc32017-09-26 23:15:36 +030074Py_LOCAL_INLINE(int)
75groupby_step(groupbyobject *gbo)
76{
77 PyObject *newvalue, *newkey, *oldvalue;
78
79 newvalue = PyIter_Next(gbo->it);
80 if (newvalue == NULL)
81 return -1;
82
83 if (gbo->keyfunc == Py_None) {
84 newkey = newvalue;
85 Py_INCREF(newvalue);
86 } else {
87 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL);
88 if (newkey == NULL) {
89 Py_DECREF(newvalue);
90 return -1;
91 }
92 }
93
94 oldvalue = gbo->currvalue;
95 gbo->currvalue = newvalue;
96 Py_XSETREF(gbo->currkey, newkey);
97 Py_XDECREF(oldvalue);
98 return 0;
99}
100
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000101static PyObject *
102groupby_next(groupbyobject *gbo)
103{
Serhiy Storchakad0b9dc32017-09-26 23:15:36 +0300104 PyObject *r, *grouper;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000105
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000106 /* skip to next iteration group */
107 for (;;) {
108 if (gbo->currkey == NULL)
109 /* pass */;
110 else if (gbo->tgtkey == NULL)
111 break;
112 else {
113 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000114
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000115 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
116 gbo->currkey, Py_EQ);
117 if (rcmp == -1)
118 return NULL;
119 else if (rcmp == 0)
120 break;
121 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122
Serhiy Storchakad0b9dc32017-09-26 23:15:36 +0300123 if (groupby_step(gbo) < 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000124 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000125 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000126 Py_INCREF(gbo->currkey);
Serhiy Storchakad0b9dc32017-09-26 23:15:36 +0300127 Py_XSETREF(gbo->tgtkey, gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000128
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000129 grouper = _grouper_create(gbo, gbo->tgtkey);
130 if (grouper == NULL)
131 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000132
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000133 r = PyTuple_Pack(2, gbo->currkey, grouper);
134 Py_DECREF(grouper);
135 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000136}
137
138PyDoc_STRVAR(groupby_doc,
139"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
140(key, sub-iterator) grouped by each value of key(value).\n");
141
142static PyTypeObject groupby_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000143 PyVarObject_HEAD_INIT(NULL, 0)
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 */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000184};
185
186
187/* _grouper object (internal) ************************************************/
188
189typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000190 PyObject_HEAD
191 PyObject *parent;
192 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000193} _grouperobject;
194
195static PyTypeObject _grouper_type;
196
197static PyObject *
198_grouper_create(groupbyobject *parent, PyObject *tgtkey)
199{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000200 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000201
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000202 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
203 if (igo == NULL)
204 return NULL;
205 igo->parent = (PyObject *)parent;
206 Py_INCREF(parent);
207 igo->tgtkey = tgtkey;
208 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000209
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000210 PyObject_GC_Track(igo);
211 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000212}
213
214static void
215_grouper_dealloc(_grouperobject *igo)
216{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000217 PyObject_GC_UnTrack(igo);
218 Py_DECREF(igo->parent);
219 Py_DECREF(igo->tgtkey);
220 PyObject_GC_Del(igo);
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000221}
222
223static int
224_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
225{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000226 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{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000234 groupbyobject *gbo = (groupbyobject *)igo->parent;
Serhiy Storchakad0b9dc32017-09-26 23:15:36 +0300235 PyObject *r;
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000236 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000237
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000238 if (gbo->currvalue == NULL) {
Serhiy Storchakad0b9dc32017-09-26 23:15:36 +0300239 if (groupby_step(gbo) < 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000240 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000241 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000242
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000243 assert(gbo->currkey != NULL);
244 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
245 if (rcmp <= 0)
246 /* got any error or current group is end */
247 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000248
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000249 r = gbo->currvalue;
250 gbo->currvalue = NULL;
251 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000252
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000253 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000254}
255
256static PyTypeObject _grouper_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000257 PyVarObject_HEAD_INIT(NULL, 0)
258 "itertools._grouper", /* tp_name */
259 sizeof(_grouperobject), /* tp_basicsize */
260 0, /* tp_itemsize */
261 /* methods */
262 (destructor)_grouper_dealloc, /* tp_dealloc */
263 0, /* tp_print */
264 0, /* tp_getattr */
265 0, /* tp_setattr */
266 0, /* tp_compare */
267 0, /* tp_repr */
268 0, /* tp_as_number */
269 0, /* tp_as_sequence */
270 0, /* tp_as_mapping */
271 0, /* tp_hash */
272 0, /* tp_call */
273 0, /* tp_str */
274 PyObject_GenericGetAttr, /* tp_getattro */
275 0, /* tp_setattro */
276 0, /* tp_as_buffer */
277 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
278 0, /* tp_doc */
279 (traverseproc)_grouper_traverse,/* tp_traverse */
280 0, /* tp_clear */
281 0, /* tp_richcompare */
282 0, /* tp_weaklistoffset */
283 PyObject_SelfIter, /* tp_iter */
284 (iternextfunc)_grouper_next, /* tp_iternext */
285 0, /* tp_methods */
286 0, /* tp_members */
287 0, /* tp_getset */
288 0, /* tp_base */
289 0, /* tp_dict */
290 0, /* tp_descr_get */
291 0, /* tp_descr_set */
292 0, /* tp_dictoffset */
293 0, /* tp_init */
294 0, /* tp_alloc */
295 0, /* tp_new */
296 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000297};
298
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000299
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000300
Raymond Hettingerad983e72003-11-12 14:32:26 +0000301/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000302
Raymond Hettingerad983e72003-11-12 14:32:26 +0000303/* The teedataobject pre-allocates space for LINKCELLS number of objects.
304 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000305 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000306 the other structure members including PyHEAD overhead. The larger the
307 value, the less memory overhead per object and the less time spent
308 allocating/deallocating new links. The smaller the number, the less
309 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000310*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000311#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000312
313typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000314 PyObject_HEAD
315 PyObject *it;
316 int numread;
317 PyObject *nextlink;
318 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000319} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000320
321typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000322 PyObject_HEAD
323 teedataobject *dataobj;
324 int index;
325 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000326} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000327
Raymond Hettingerad983e72003-11-12 14:32:26 +0000328static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000329
330static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000331teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000332{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000333 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000334
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000335 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
336 if (tdo == NULL)
337 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000338
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000339 tdo->numread = 0;
340 tdo->nextlink = NULL;
341 Py_INCREF(it);
342 tdo->it = it;
343 PyObject_GC_Track(tdo);
344 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000345}
346
347static PyObject *
348teedataobject_jumplink(teedataobject *tdo)
349{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000350 if (tdo->nextlink == NULL)
351 tdo->nextlink = teedataobject_new(tdo->it);
352 Py_XINCREF(tdo->nextlink);
353 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000354}
355
356static PyObject *
357teedataobject_getitem(teedataobject *tdo, int i)
358{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000359 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000360
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000361 assert(i < LINKCELLS);
362 if (i < tdo->numread)
363 value = tdo->values[i];
364 else {
365 /* this is the lead iterator, so fetch more data */
366 assert(i == tdo->numread);
367 value = PyIter_Next(tdo->it);
368 if (value == NULL)
369 return NULL;
370 tdo->numread++;
371 tdo->values[i] = value;
372 }
373 Py_INCREF(value);
374 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000375}
376
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000377static int
378teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
379{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000380 int i;
381 Py_VISIT(tdo->it);
382 for (i = 0; i < tdo->numread; i++)
383 Py_VISIT(tdo->values[i]);
384 Py_VISIT(tdo->nextlink);
385 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000386}
387
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200388static void
389teedataobject_safe_decref(PyObject *obj)
390{
391 while (obj && Py_TYPE(obj) == &teedataobject_type &&
392 Py_REFCNT(obj) == 1) {
393 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
394 ((teedataobject *)obj)->nextlink = NULL;
395 Py_DECREF(obj);
396 obj = nextlink;
397 }
398 Py_XDECREF(obj);
399}
400
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000401static int
402teedataobject_clear(teedataobject *tdo)
403{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000404 int i;
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200405 PyObject *tmp;
406
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000407 Py_CLEAR(tdo->it);
408 for (i=0 ; i<tdo->numread ; i++)
409 Py_CLEAR(tdo->values[i]);
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200410 tmp = tdo->nextlink;
411 tdo->nextlink = NULL;
412 teedataobject_safe_decref(tmp);
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000413 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000414}
415
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000416static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000417teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000418{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000419 PyObject_GC_UnTrack(tdo);
420 teedataobject_clear(tdo);
421 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000422}
423
Raymond Hettingerad983e72003-11-12 14:32:26 +0000424PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000425
Raymond Hettingerad983e72003-11-12 14:32:26 +0000426static PyTypeObject teedataobject_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000427 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
428 "itertools.tee_dataobject", /* tp_name */
429 sizeof(teedataobject), /* tp_basicsize */
430 0, /* tp_itemsize */
431 /* methods */
432 (destructor)teedataobject_dealloc, /* tp_dealloc */
433 0, /* tp_print */
434 0, /* tp_getattr */
435 0, /* tp_setattr */
436 0, /* tp_compare */
437 0, /* tp_repr */
438 0, /* tp_as_number */
439 0, /* tp_as_sequence */
440 0, /* tp_as_mapping */
441 0, /* tp_hash */
442 0, /* tp_call */
443 0, /* tp_str */
444 PyObject_GenericGetAttr, /* tp_getattro */
445 0, /* tp_setattro */
446 0, /* tp_as_buffer */
447 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
448 teedataobject_doc, /* tp_doc */
449 (traverseproc)teedataobject_traverse, /* tp_traverse */
450 (inquiry)teedataobject_clear, /* tp_clear */
451 0, /* tp_richcompare */
452 0, /* tp_weaklistoffset */
453 0, /* tp_iter */
454 0, /* tp_iternext */
455 0, /* tp_methods */
456 0, /* tp_members */
457 0, /* tp_getset */
458 0, /* tp_base */
459 0, /* tp_dict */
460 0, /* tp_descr_get */
461 0, /* tp_descr_set */
462 0, /* tp_dictoffset */
463 0, /* tp_init */
464 0, /* tp_alloc */
465 0, /* tp_new */
466 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000467};
468
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000469
470static PyTypeObject tee_type;
471
472static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000473tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000474{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000475 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000476
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000477 if (to->index >= LINKCELLS) {
478 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200479 if (link == NULL)
480 return NULL;
Serhiy Storchaka763a61c2016-04-10 18:05:12 +0300481 Py_SETREF(to->dataobj, (teedataobject *)link);
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000482 to->index = 0;
483 }
484 value = teedataobject_getitem(to->dataobj, to->index);
485 if (value == NULL)
486 return NULL;
487 to->index++;
488 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000489}
490
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000491static int
492tee_traverse(teeobject *to, visitproc visit, void *arg)
493{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000494 Py_VISIT((PyObject *)to->dataobj);
495 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000496}
497
Raymond Hettingerad983e72003-11-12 14:32:26 +0000498static PyObject *
499tee_copy(teeobject *to)
500{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000501 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000502
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000503 newto = PyObject_GC_New(teeobject, &tee_type);
504 if (newto == NULL)
505 return NULL;
506 Py_INCREF(to->dataobj);
507 newto->dataobj = to->dataobj;
508 newto->index = to->index;
509 newto->weakreflist = NULL;
510 PyObject_GC_Track(newto);
511 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000512}
513
514PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
515
516static PyObject *
517tee_fromiterable(PyObject *iterable)
518{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000519 teeobject *to;
520 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000521
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000522 it = PyObject_GetIter(iterable);
523 if (it == NULL)
524 return NULL;
525 if (PyObject_TypeCheck(it, &tee_type)) {
526 to = (teeobject *)tee_copy((teeobject *)it);
527 goto done;
528 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000529
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000530 to = PyObject_GC_New(teeobject, &tee_type);
531 if (to == NULL)
532 goto done;
533 to->dataobj = (teedataobject *)teedataobject_new(it);
534 if (!to->dataobj) {
535 PyObject_GC_Del(to);
536 to = NULL;
537 goto done;
538 }
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000539
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000540 to->index = 0;
541 to->weakreflist = NULL;
542 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000543done:
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000544 Py_XDECREF(it);
545 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000546}
547
548static PyObject *
549tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
550{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000551 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000552
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000553 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
554 return NULL;
555 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000556}
557
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000558static int
559tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000560{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000561 if (to->weakreflist != NULL)
562 PyObject_ClearWeakRefs((PyObject *) to);
563 Py_CLEAR(to->dataobj);
564 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000565}
566
567static void
568tee_dealloc(teeobject *to)
569{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000570 PyObject_GC_UnTrack(to);
571 tee_clear(to);
572 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000573}
574
Raymond Hettingerad983e72003-11-12 14:32:26 +0000575PyDoc_STRVAR(teeobject_doc,
576"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000577
Raymond Hettingerad983e72003-11-12 14:32:26 +0000578static PyMethodDef tee_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000579 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
580 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000581};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000582
583static PyTypeObject tee_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000584 PyVarObject_HEAD_INIT(NULL, 0)
585 "itertools.tee", /* tp_name */
586 sizeof(teeobject), /* tp_basicsize */
587 0, /* tp_itemsize */
588 /* methods */
589 (destructor)tee_dealloc, /* tp_dealloc */
590 0, /* tp_print */
591 0, /* tp_getattr */
592 0, /* tp_setattr */
593 0, /* tp_compare */
594 0, /* tp_repr */
595 0, /* tp_as_number */
596 0, /* tp_as_sequence */
597 0, /* tp_as_mapping */
598 0, /* tp_hash */
599 0, /* tp_call */
600 0, /* tp_str */
601 0, /* tp_getattro */
602 0, /* tp_setattro */
603 0, /* tp_as_buffer */
604 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
605 teeobject_doc, /* tp_doc */
606 (traverseproc)tee_traverse, /* tp_traverse */
607 (inquiry)tee_clear, /* tp_clear */
608 0, /* tp_richcompare */
609 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
610 PyObject_SelfIter, /* tp_iter */
611 (iternextfunc)tee_next, /* tp_iternext */
612 tee_methods, /* tp_methods */
613 0, /* tp_members */
614 0, /* tp_getset */
615 0, /* tp_base */
616 0, /* tp_dict */
617 0, /* tp_descr_get */
618 0, /* tp_descr_set */
619 0, /* tp_dictoffset */
620 0, /* tp_init */
621 0, /* tp_alloc */
622 tee_new, /* tp_new */
623 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000624};
625
Raymond Hettingerad983e72003-11-12 14:32:26 +0000626static PyObject *
627tee(PyObject *self, PyObject *args)
628{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000629 Py_ssize_t i, n=2;
630 PyObject *it, *iterable, *copyable, *result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000631
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000632 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
633 return NULL;
634 if (n < 0) {
635 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
636 return NULL;
637 }
638 result = PyTuple_New(n);
639 if (result == NULL)
640 return NULL;
641 if (n == 0)
642 return result;
643 it = PyObject_GetIter(iterable);
644 if (it == NULL) {
645 Py_DECREF(result);
646 return NULL;
647 }
648 if (!PyObject_HasAttrString(it, "__copy__")) {
649 copyable = tee_fromiterable(it);
650 Py_DECREF(it);
651 if (copyable == NULL) {
652 Py_DECREF(result);
653 return NULL;
654 }
655 } else
656 copyable = it;
657 PyTuple_SET_ITEM(result, 0, copyable);
658 for (i=1 ; i<n ; i++) {
659 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
660 if (copyable == NULL) {
661 Py_DECREF(result);
662 return NULL;
663 }
664 PyTuple_SET_ITEM(result, i, copyable);
665 }
666 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000667}
668
669PyDoc_STRVAR(tee_doc,
670"tee(iterable, n=2) --> tuple of n independent iterators.");
671
672
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000673/* cycle object **********************************************************/
674
675typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000676 PyObject_HEAD
677 PyObject *it;
678 PyObject *saved;
679 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000680} cycleobject;
681
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000682static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000683
684static PyObject *
685cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
686{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000687 PyObject *it;
688 PyObject *iterable;
689 PyObject *saved;
690 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000691
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000692 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
693 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000694
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000695 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
696 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000697
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000698 /* Get iterator. */
699 it = PyObject_GetIter(iterable);
700 if (it == NULL)
701 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000702
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000703 saved = PyList_New(0);
704 if (saved == NULL) {
705 Py_DECREF(it);
706 return NULL;
707 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000708
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000709 /* create cycleobject structure */
710 lz = (cycleobject *)type->tp_alloc(type, 0);
711 if (lz == NULL) {
712 Py_DECREF(it);
713 Py_DECREF(saved);
714 return NULL;
715 }
716 lz->it = it;
717 lz->saved = saved;
718 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000719
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000720 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000721}
722
723static void
724cycle_dealloc(cycleobject *lz)
725{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000726 PyObject_GC_UnTrack(lz);
727 Py_XDECREF(lz->saved);
728 Py_XDECREF(lz->it);
729 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000730}
731
732static int
733cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
734{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000735 Py_VISIT(lz->it);
736 Py_VISIT(lz->saved);
737 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000738}
739
740static PyObject *
741cycle_next(cycleobject *lz)
742{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000743 PyObject *item;
744 PyObject *it;
745 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000746
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000747 while (1) {
748 item = PyIter_Next(lz->it);
749 if (item != NULL) {
750 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
751 Py_DECREF(item);
752 return NULL;
753 }
754 return item;
755 }
756 if (PyErr_Occurred()) {
757 if (PyErr_ExceptionMatches(PyExc_StopIteration))
758 PyErr_Clear();
759 else
760 return NULL;
761 }
762 if (PyList_Size(lz->saved) == 0)
763 return NULL;
764 it = PyObject_GetIter(lz->saved);
765 if (it == NULL)
766 return NULL;
767 tmp = lz->it;
768 lz->it = it;
769 lz->firstpass = 1;
770 Py_DECREF(tmp);
771 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000772}
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 = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000781 PyVarObject_HEAD_INIT(NULL, 0)
782 "itertools.cycle", /* tp_name */
783 sizeof(cycleobject), /* tp_basicsize */
784 0, /* tp_itemsize */
785 /* methods */
786 (destructor)cycle_dealloc, /* tp_dealloc */
787 0, /* tp_print */
788 0, /* tp_getattr */
789 0, /* tp_setattr */
790 0, /* tp_compare */
791 0, /* tp_repr */
792 0, /* tp_as_number */
793 0, /* tp_as_sequence */
794 0, /* tp_as_mapping */
795 0, /* tp_hash */
796 0, /* tp_call */
797 0, /* tp_str */
798 PyObject_GenericGetAttr, /* tp_getattro */
799 0, /* tp_setattro */
800 0, /* tp_as_buffer */
801 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
802 Py_TPFLAGS_BASETYPE, /* tp_flags */
803 cycle_doc, /* tp_doc */
804 (traverseproc)cycle_traverse, /* tp_traverse */
805 0, /* tp_clear */
806 0, /* tp_richcompare */
807 0, /* tp_weaklistoffset */
808 PyObject_SelfIter, /* tp_iter */
809 (iternextfunc)cycle_next, /* tp_iternext */
810 0, /* tp_methods */
811 0, /* tp_members */
812 0, /* tp_getset */
813 0, /* tp_base */
814 0, /* tp_dict */
815 0, /* tp_descr_get */
816 0, /* tp_descr_set */
817 0, /* tp_dictoffset */
818 0, /* tp_init */
819 0, /* tp_alloc */
820 cycle_new, /* tp_new */
821 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000822};
823
824
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000825/* dropwhile object **********************************************************/
826
827typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000828 PyObject_HEAD
829 PyObject *func;
830 PyObject *it;
831 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000832} dropwhileobject;
833
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000834static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000835
836static PyObject *
837dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
838{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000839 PyObject *func, *seq;
840 PyObject *it;
841 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000842
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000843 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
844 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000845
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000846 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
847 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000848
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000849 /* Get iterator. */
850 it = PyObject_GetIter(seq);
851 if (it == NULL)
852 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000853
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000854 /* create dropwhileobject structure */
855 lz = (dropwhileobject *)type->tp_alloc(type, 0);
856 if (lz == NULL) {
857 Py_DECREF(it);
858 return NULL;
859 }
860 Py_INCREF(func);
861 lz->func = func;
862 lz->it = it;
863 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000864
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000865 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000866}
867
868static void
869dropwhile_dealloc(dropwhileobject *lz)
870{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000871 PyObject_GC_UnTrack(lz);
872 Py_XDECREF(lz->func);
873 Py_XDECREF(lz->it);
874 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000875}
876
877static int
878dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
879{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000880 Py_VISIT(lz->it);
881 Py_VISIT(lz->func);
882 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000883}
884
885static PyObject *
886dropwhile_next(dropwhileobject *lz)
887{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000888 PyObject *item, *good;
889 PyObject *it = lz->it;
890 long ok;
891 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000892
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000893 iternext = *Py_TYPE(it)->tp_iternext;
894 for (;;) {
895 item = iternext(it);
896 if (item == NULL)
897 return NULL;
898 if (lz->start == 1)
899 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000900
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000901 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
902 if (good == NULL) {
903 Py_DECREF(item);
904 return NULL;
905 }
906 ok = PyObject_IsTrue(good);
907 Py_DECREF(good);
Antoine Pitrouc5bef752012-08-15 23:16:51 +0200908 if (ok == 0) {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000909 lz->start = 1;
910 return item;
911 }
912 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +0200913 if (ok < 0)
914 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000915 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000916}
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 = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000925 PyVarObject_HEAD_INIT(NULL, 0)
926 "itertools.dropwhile", /* tp_name */
927 sizeof(dropwhileobject), /* tp_basicsize */
928 0, /* tp_itemsize */
929 /* methods */
930 (destructor)dropwhile_dealloc, /* tp_dealloc */
931 0, /* tp_print */
932 0, /* tp_getattr */
933 0, /* tp_setattr */
934 0, /* tp_compare */
935 0, /* tp_repr */
936 0, /* tp_as_number */
937 0, /* tp_as_sequence */
938 0, /* tp_as_mapping */
939 0, /* tp_hash */
940 0, /* tp_call */
941 0, /* tp_str */
942 PyObject_GenericGetAttr, /* tp_getattro */
943 0, /* tp_setattro */
944 0, /* tp_as_buffer */
945 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
946 Py_TPFLAGS_BASETYPE, /* tp_flags */
947 dropwhile_doc, /* tp_doc */
948 (traverseproc)dropwhile_traverse, /* tp_traverse */
949 0, /* tp_clear */
950 0, /* tp_richcompare */
951 0, /* tp_weaklistoffset */
952 PyObject_SelfIter, /* tp_iter */
953 (iternextfunc)dropwhile_next, /* tp_iternext */
954 0, /* tp_methods */
955 0, /* tp_members */
956 0, /* tp_getset */
957 0, /* tp_base */
958 0, /* tp_dict */
959 0, /* tp_descr_get */
960 0, /* tp_descr_set */
961 0, /* tp_dictoffset */
962 0, /* tp_init */
963 0, /* tp_alloc */
964 dropwhile_new, /* tp_new */
965 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000966};
967
968
969/* takewhile object **********************************************************/
970
971typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000972 PyObject_HEAD
973 PyObject *func;
974 PyObject *it;
975 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000976} takewhileobject;
977
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000978static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000979
980static PyObject *
981takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
982{
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000983 PyObject *func, *seq;
984 PyObject *it;
985 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000986
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000987 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
988 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000989
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000990 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
991 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000992
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000993 /* Get iterator. */
994 it = PyObject_GetIter(seq);
995 if (it == NULL)
996 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000997
Antoine Pitrouc83ea132010-05-09 14:46:46 +0000998 /* create takewhileobject structure */
999 lz = (takewhileobject *)type->tp_alloc(type, 0);
1000 if (lz == NULL) {
1001 Py_DECREF(it);
1002 return NULL;
1003 }
1004 Py_INCREF(func);
1005 lz->func = func;
1006 lz->it = it;
1007 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001008
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001009 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001010}
1011
1012static void
1013takewhile_dealloc(takewhileobject *lz)
1014{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001015 PyObject_GC_UnTrack(lz);
1016 Py_XDECREF(lz->func);
1017 Py_XDECREF(lz->it);
1018 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001019}
1020
1021static int
1022takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1023{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001024 Py_VISIT(lz->it);
1025 Py_VISIT(lz->func);
1026 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001027}
1028
1029static PyObject *
1030takewhile_next(takewhileobject *lz)
1031{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001032 PyObject *item, *good;
1033 PyObject *it = lz->it;
1034 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001035
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001036 if (lz->stop == 1)
1037 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001038
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001039 item = (*Py_TYPE(it)->tp_iternext)(it);
1040 if (item == NULL)
1041 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001042
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001043 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1044 if (good == NULL) {
1045 Py_DECREF(item);
1046 return NULL;
1047 }
1048 ok = PyObject_IsTrue(good);
1049 Py_DECREF(good);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02001050 if (ok > 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001051 return item;
1052 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02001053 if (ok == 0)
1054 lz->stop = 1;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001055 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001056}
1057
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001058PyDoc_STRVAR(takewhile_doc,
1059"takewhile(predicate, iterable) --> takewhile object\n\
1060\n\
1061Return successive entries from an iterable as long as the \n\
1062predicate evaluates to true for each entry.");
1063
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001064static PyTypeObject takewhile_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001065 PyVarObject_HEAD_INIT(NULL, 0)
1066 "itertools.takewhile", /* tp_name */
1067 sizeof(takewhileobject), /* tp_basicsize */
1068 0, /* tp_itemsize */
1069 /* methods */
1070 (destructor)takewhile_dealloc, /* tp_dealloc */
1071 0, /* tp_print */
1072 0, /* tp_getattr */
1073 0, /* tp_setattr */
1074 0, /* tp_compare */
1075 0, /* tp_repr */
1076 0, /* tp_as_number */
1077 0, /* tp_as_sequence */
1078 0, /* tp_as_mapping */
1079 0, /* tp_hash */
1080 0, /* tp_call */
1081 0, /* tp_str */
1082 PyObject_GenericGetAttr, /* tp_getattro */
1083 0, /* tp_setattro */
1084 0, /* tp_as_buffer */
1085 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1086 Py_TPFLAGS_BASETYPE, /* tp_flags */
1087 takewhile_doc, /* tp_doc */
1088 (traverseproc)takewhile_traverse, /* tp_traverse */
1089 0, /* tp_clear */
1090 0, /* tp_richcompare */
1091 0, /* tp_weaklistoffset */
1092 PyObject_SelfIter, /* tp_iter */
1093 (iternextfunc)takewhile_next, /* tp_iternext */
1094 0, /* tp_methods */
1095 0, /* tp_members */
1096 0, /* tp_getset */
1097 0, /* tp_base */
1098 0, /* tp_dict */
1099 0, /* tp_descr_get */
1100 0, /* tp_descr_set */
1101 0, /* tp_dictoffset */
1102 0, /* tp_init */
1103 0, /* tp_alloc */
1104 takewhile_new, /* tp_new */
1105 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001106};
1107
1108
1109/* islice object ************************************************************/
1110
1111typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001112 PyObject_HEAD
1113 PyObject *it;
1114 Py_ssize_t next;
1115 Py_ssize_t stop;
1116 Py_ssize_t step;
1117 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001118} isliceobject;
1119
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001120static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001121
1122static PyObject *
1123islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1124{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001125 PyObject *seq;
1126 Py_ssize_t start=0, stop=-1, step=1;
1127 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1128 Py_ssize_t numargs;
1129 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001130
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001131 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1132 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001133
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001134 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1135 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001136
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001137 numargs = PyTuple_Size(args);
1138 if (numargs == 2) {
1139 if (a1 != Py_None) {
1140 stop = PyInt_AsSsize_t(a1);
1141 if (stop == -1) {
1142 if (PyErr_Occurred())
1143 PyErr_Clear();
1144 PyErr_SetString(PyExc_ValueError,
1145 "Stop argument for islice() must be None or an integer: 0 <= x <= maxint.");
1146 return NULL;
1147 }
1148 }
1149 } else {
1150 if (a1 != Py_None)
1151 start = PyInt_AsSsize_t(a1);
1152 if (start == -1 && PyErr_Occurred())
1153 PyErr_Clear();
1154 if (a2 != Py_None) {
1155 stop = PyInt_AsSsize_t(a2);
1156 if (stop == -1) {
1157 if (PyErr_Occurred())
1158 PyErr_Clear();
1159 PyErr_SetString(PyExc_ValueError,
1160 "Stop argument for islice() must be None or an integer: 0 <= x <= maxint.");
1161 return NULL;
1162 }
1163 }
1164 }
1165 if (start<0 || stop<-1) {
1166 PyErr_SetString(PyExc_ValueError,
1167 "Indices for islice() must be None or an integer: 0 <= x <= maxint.");
1168 return NULL;
1169 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001170
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001171 if (a3 != NULL) {
1172 if (a3 != Py_None)
1173 step = PyInt_AsSsize_t(a3);
1174 if (step == -1 && PyErr_Occurred())
1175 PyErr_Clear();
1176 }
1177 if (step<1) {
1178 PyErr_SetString(PyExc_ValueError,
1179 "Step for islice() must be a positive integer or None.");
1180 return NULL;
1181 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001182
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001183 /* Get iterator. */
1184 it = PyObject_GetIter(seq);
1185 if (it == NULL)
1186 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001187
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001188 /* create isliceobject structure */
1189 lz = (isliceobject *)type->tp_alloc(type, 0);
1190 if (lz == NULL) {
1191 Py_DECREF(it);
1192 return NULL;
1193 }
1194 lz->it = it;
1195 lz->next = start;
1196 lz->stop = stop;
1197 lz->step = step;
1198 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001199
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001200 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001201}
1202
1203static void
1204islice_dealloc(isliceobject *lz)
1205{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001206 PyObject_GC_UnTrack(lz);
1207 Py_XDECREF(lz->it);
1208 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001209}
1210
1211static int
1212islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1213{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001214 Py_VISIT(lz->it);
1215 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001216}
1217
1218static PyObject *
1219islice_next(isliceobject *lz)
1220{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001221 PyObject *item;
1222 PyObject *it = lz->it;
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001223 Py_ssize_t stop = lz->stop;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001224 Py_ssize_t oldnext;
1225 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001226
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001227 if (it == NULL)
1228 return NULL;
1229
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001230 iternext = *Py_TYPE(it)->tp_iternext;
1231 while (lz->cnt < lz->next) {
1232 item = iternext(it);
1233 if (item == NULL)
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001234 goto empty;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001235 Py_DECREF(item);
1236 lz->cnt++;
1237 }
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001238 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001239 goto empty;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001240 item = iternext(it);
1241 if (item == NULL)
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001242 goto empty;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001243 lz->cnt++;
1244 oldnext = lz->next;
Mark Dickinsona96b0d12011-09-24 09:01:16 +01001245 /* The (size_t) cast below avoids the danger of undefined
1246 behaviour from signed integer overflow. */
1247 lz->next += (size_t)lz->step;
Raymond Hettinger061bf7a2010-11-30 03:15:35 +00001248 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1249 lz->next = stop;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001250 return item;
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02001251
1252empty:
1253 Py_CLEAR(lz->it);
1254 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001255}
1256
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001257PyDoc_STRVAR(islice_doc,
1258"islice(iterable, [start,] stop [, step]) --> islice object\n\
1259\n\
1260Return an iterator whose next() method returns selected values from an\n\
1261iterable. If start is specified, will skip all preceding elements;\n\
1262otherwise, start defaults to zero. Step defaults to one. If\n\
1263specified as another value, step determines how many values are \n\
1264skipped between successive calls. Works like a slice() on a list\n\
1265but returns an iterator.");
1266
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001267static PyTypeObject islice_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001268 PyVarObject_HEAD_INIT(NULL, 0)
1269 "itertools.islice", /* tp_name */
1270 sizeof(isliceobject), /* tp_basicsize */
1271 0, /* tp_itemsize */
1272 /* methods */
1273 (destructor)islice_dealloc, /* tp_dealloc */
1274 0, /* tp_print */
1275 0, /* tp_getattr */
1276 0, /* tp_setattr */
1277 0, /* tp_compare */
1278 0, /* tp_repr */
1279 0, /* tp_as_number */
1280 0, /* tp_as_sequence */
1281 0, /* tp_as_mapping */
1282 0, /* tp_hash */
1283 0, /* tp_call */
1284 0, /* tp_str */
1285 PyObject_GenericGetAttr, /* tp_getattro */
1286 0, /* tp_setattro */
1287 0, /* tp_as_buffer */
1288 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1289 Py_TPFLAGS_BASETYPE, /* tp_flags */
1290 islice_doc, /* tp_doc */
1291 (traverseproc)islice_traverse, /* tp_traverse */
1292 0, /* tp_clear */
1293 0, /* tp_richcompare */
1294 0, /* tp_weaklistoffset */
1295 PyObject_SelfIter, /* tp_iter */
1296 (iternextfunc)islice_next, /* tp_iternext */
1297 0, /* tp_methods */
1298 0, /* tp_members */
1299 0, /* tp_getset */
1300 0, /* tp_base */
1301 0, /* tp_dict */
1302 0, /* tp_descr_get */
1303 0, /* tp_descr_set */
1304 0, /* tp_dictoffset */
1305 0, /* tp_init */
1306 0, /* tp_alloc */
1307 islice_new, /* tp_new */
1308 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001309};
1310
1311
1312/* starmap object ************************************************************/
1313
1314typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001315 PyObject_HEAD
1316 PyObject *func;
1317 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001318} starmapobject;
1319
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001320static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001321
1322static PyObject *
1323starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1324{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001325 PyObject *func, *seq;
1326 PyObject *it;
1327 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001328
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001329 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1330 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001331
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001332 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1333 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001334
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001335 /* Get iterator. */
1336 it = PyObject_GetIter(seq);
1337 if (it == NULL)
1338 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001339
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001340 /* create starmapobject structure */
1341 lz = (starmapobject *)type->tp_alloc(type, 0);
1342 if (lz == NULL) {
1343 Py_DECREF(it);
1344 return NULL;
1345 }
1346 Py_INCREF(func);
1347 lz->func = func;
1348 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001349
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001350 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001351}
1352
1353static void
1354starmap_dealloc(starmapobject *lz)
1355{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001356 PyObject_GC_UnTrack(lz);
1357 Py_XDECREF(lz->func);
1358 Py_XDECREF(lz->it);
1359 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001360}
1361
1362static int
1363starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1364{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001365 Py_VISIT(lz->it);
1366 Py_VISIT(lz->func);
1367 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001368}
1369
1370static PyObject *
1371starmap_next(starmapobject *lz)
1372{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001373 PyObject *args;
1374 PyObject *result;
1375 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001376
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001377 args = (*Py_TYPE(it)->tp_iternext)(it);
1378 if (args == NULL)
1379 return NULL;
1380 if (!PyTuple_CheckExact(args)) {
1381 PyObject *newargs = PySequence_Tuple(args);
1382 Py_DECREF(args);
1383 if (newargs == NULL)
1384 return NULL;
1385 args = newargs;
1386 }
1387 result = PyObject_Call(lz->func, args, NULL);
1388 Py_DECREF(args);
1389 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001390}
1391
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001392PyDoc_STRVAR(starmap_doc,
1393"starmap(function, sequence) --> starmap object\n\
1394\n\
1395Return an iterator whose values are returned from the function evaluated\n\
Serhiy Storchakac72e66a2015-11-02 15:06:09 +02001396with an argument tuple taken from the given sequence.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001397
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001398static PyTypeObject starmap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001399 PyVarObject_HEAD_INIT(NULL, 0)
1400 "itertools.starmap", /* tp_name */
1401 sizeof(starmapobject), /* tp_basicsize */
1402 0, /* tp_itemsize */
1403 /* methods */
1404 (destructor)starmap_dealloc, /* tp_dealloc */
1405 0, /* tp_print */
1406 0, /* tp_getattr */
1407 0, /* tp_setattr */
1408 0, /* tp_compare */
1409 0, /* tp_repr */
1410 0, /* tp_as_number */
1411 0, /* tp_as_sequence */
1412 0, /* tp_as_mapping */
1413 0, /* tp_hash */
1414 0, /* tp_call */
1415 0, /* tp_str */
1416 PyObject_GenericGetAttr, /* tp_getattro */
1417 0, /* tp_setattro */
1418 0, /* tp_as_buffer */
1419 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1420 Py_TPFLAGS_BASETYPE, /* tp_flags */
1421 starmap_doc, /* tp_doc */
1422 (traverseproc)starmap_traverse, /* tp_traverse */
1423 0, /* tp_clear */
1424 0, /* tp_richcompare */
1425 0, /* tp_weaklistoffset */
1426 PyObject_SelfIter, /* tp_iter */
1427 (iternextfunc)starmap_next, /* tp_iternext */
1428 0, /* tp_methods */
1429 0, /* tp_members */
1430 0, /* tp_getset */
1431 0, /* tp_base */
1432 0, /* tp_dict */
1433 0, /* tp_descr_get */
1434 0, /* tp_descr_set */
1435 0, /* tp_dictoffset */
1436 0, /* tp_init */
1437 0, /* tp_alloc */
1438 starmap_new, /* tp_new */
1439 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001440};
1441
1442
1443/* imap object ************************************************************/
1444
1445typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001446 PyObject_HEAD
1447 PyObject *iters;
1448 PyObject *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001449} imapobject;
1450
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001451static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001452
1453static PyObject *
1454imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1455{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001456 PyObject *it, *iters, *func;
1457 imapobject *lz;
1458 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001460 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
1461 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001462
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001463 numargs = PyTuple_Size(args);
1464 if (numargs < 2) {
1465 PyErr_SetString(PyExc_TypeError,
1466 "imap() must have at least two arguments.");
1467 return NULL;
1468 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001470 iters = PyTuple_New(numargs-1);
1471 if (iters == NULL)
1472 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001473
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001474 for (i=1 ; i<numargs ; i++) {
1475 /* Get iterator. */
1476 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1477 if (it == NULL) {
1478 Py_DECREF(iters);
1479 return NULL;
1480 }
1481 PyTuple_SET_ITEM(iters, i-1, it);
1482 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001483
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001484 /* create imapobject structure */
1485 lz = (imapobject *)type->tp_alloc(type, 0);
1486 if (lz == NULL) {
1487 Py_DECREF(iters);
1488 return NULL;
1489 }
1490 lz->iters = iters;
1491 func = PyTuple_GET_ITEM(args, 0);
1492 Py_INCREF(func);
1493 lz->func = func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001494
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001495 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001496}
1497
1498static void
1499imap_dealloc(imapobject *lz)
1500{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001501 PyObject_GC_UnTrack(lz);
1502 Py_XDECREF(lz->iters);
1503 Py_XDECREF(lz->func);
1504 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001505}
1506
1507static int
1508imap_traverse(imapobject *lz, visitproc visit, void *arg)
1509{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001510 Py_VISIT(lz->iters);
1511 Py_VISIT(lz->func);
1512 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001513}
1514
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001515/*
Raymond Hettinger2012f172003-02-07 05:32:58 +00001516imap() is an iterator version of __builtins__.map() except that it does
1517not have the None fill-in feature. That was intentionally left out for
1518the following reasons:
1519
1520 1) Itertools are designed to be easily combined and chained together.
1521 Having all tools stop with the shortest input is a unifying principle
1522 that makes it easier to combine finite iterators (supplying data) with
1523 infinite iterators like count() and repeat() (for supplying sequential
1524 or constant arguments to a function).
1525
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001526 2) In typical use cases for combining itertools, having one finite data
Raymond Hettinger2012f172003-02-07 05:32:58 +00001527 supplier run out before another is likely to be an error condition which
1528 should not pass silently by automatically supplying None.
1529
1530 3) The use cases for automatic None fill-in are rare -- not many functions
1531 do something useful when a parameter suddenly switches type and becomes
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001532 None.
Raymond Hettinger2012f172003-02-07 05:32:58 +00001533
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001534 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001535 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001536
1537 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1538*/
1539
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001540static PyObject *
1541imap_next(imapobject *lz)
1542{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001543 PyObject *val;
1544 PyObject *argtuple;
1545 PyObject *result;
1546 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001547
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001548 numargs = PyTuple_Size(lz->iters);
1549 argtuple = PyTuple_New(numargs);
1550 if (argtuple == NULL)
1551 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001552
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001553 for (i=0 ; i<numargs ; i++) {
1554 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1555 if (val == NULL) {
1556 Py_DECREF(argtuple);
1557 return NULL;
1558 }
1559 PyTuple_SET_ITEM(argtuple, i, val);
1560 }
1561 if (lz->func == Py_None)
1562 return argtuple;
1563 result = PyObject_Call(lz->func, argtuple, NULL);
1564 Py_DECREF(argtuple);
1565 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001566}
1567
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001568PyDoc_STRVAR(imap_doc,
1569"imap(func, *iterables) --> imap object\n\
1570\n\
1571Make an iterator that computes the function using arguments from\n\
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001572each of the iterables. Like map() except that it returns\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001573an iterator instead of a list and that it stops when the shortest\n\
1574iterable is exhausted instead of filling in None for shorter\n\
1575iterables.");
1576
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001577static PyTypeObject imap_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001578 PyVarObject_HEAD_INIT(NULL, 0)
1579 "itertools.imap", /* tp_name */
1580 sizeof(imapobject), /* tp_basicsize */
1581 0, /* tp_itemsize */
1582 /* methods */
1583 (destructor)imap_dealloc, /* tp_dealloc */
1584 0, /* tp_print */
1585 0, /* tp_getattr */
1586 0, /* tp_setattr */
1587 0, /* tp_compare */
1588 0, /* tp_repr */
1589 0, /* tp_as_number */
1590 0, /* tp_as_sequence */
1591 0, /* tp_as_mapping */
1592 0, /* tp_hash */
1593 0, /* tp_call */
1594 0, /* tp_str */
1595 PyObject_GenericGetAttr, /* tp_getattro */
1596 0, /* tp_setattro */
1597 0, /* tp_as_buffer */
1598 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1599 Py_TPFLAGS_BASETYPE, /* tp_flags */
1600 imap_doc, /* tp_doc */
1601 (traverseproc)imap_traverse, /* tp_traverse */
1602 0, /* tp_clear */
1603 0, /* tp_richcompare */
1604 0, /* tp_weaklistoffset */
1605 PyObject_SelfIter, /* tp_iter */
1606 (iternextfunc)imap_next, /* tp_iternext */
1607 0, /* tp_methods */
1608 0, /* tp_members */
1609 0, /* tp_getset */
1610 0, /* tp_base */
1611 0, /* tp_dict */
1612 0, /* tp_descr_get */
1613 0, /* tp_descr_set */
1614 0, /* tp_dictoffset */
1615 0, /* tp_init */
1616 0, /* tp_alloc */
1617 imap_new, /* tp_new */
1618 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001619};
1620
1621
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001622/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001623
1624typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001625 PyObject_HEAD
1626 PyObject *source; /* Iterator over input iterables */
1627 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001628} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001629
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001630static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001631
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001632static PyObject *
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001633chain_new_internal(PyTypeObject *type, PyObject *source)
1634{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001635 chainobject *lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001636
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001637 lz = (chainobject *)type->tp_alloc(type, 0);
1638 if (lz == NULL) {
1639 Py_DECREF(source);
1640 return NULL;
1641 }
1642
1643 lz->source = source;
1644 lz->active = NULL;
1645 return (PyObject *)lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001646}
1647
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001648static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001649chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001650{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001651 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001652
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001653 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1654 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001655
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001656 source = PyObject_GetIter(args);
1657 if (source == NULL)
1658 return NULL;
1659
1660 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001661}
1662
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001663static PyObject *
1664chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1665{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001666 PyObject *source;
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001667
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001668 source = PyObject_GetIter(arg);
1669 if (source == NULL)
1670 return NULL;
1671
1672 return chain_new_internal(type, source);
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001673}
1674
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001675static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001676chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001677{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001678 PyObject_GC_UnTrack(lz);
1679 Py_XDECREF(lz->active);
1680 Py_XDECREF(lz->source);
1681 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001682}
1683
Raymond Hettinger2012f172003-02-07 05:32:58 +00001684static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001685chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001686{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001687 Py_VISIT(lz->source);
1688 Py_VISIT(lz->active);
1689 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001690}
1691
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001692static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001693chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001694{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001695 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001696
T. Woutersd694a062017-03-30 12:49:22 -07001697 /* lz->source is the iterator of iterables. If it's NULL, we've already
1698 * consumed them all. lz->active is the current iterator. If it's NULL,
1699 * we should grab a new one from lz->source. */
1700 while (lz->source != NULL) {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001701 if (lz->active == NULL) {
T. Woutersd694a062017-03-30 12:49:22 -07001702 PyObject *iterable = PyIter_Next(lz->source);
1703 if (iterable == NULL) {
1704 Py_CLEAR(lz->source);
1705 return NULL; /* no more input sources */
1706 }
1707 lz->active = PyObject_GetIter(iterable);
1708 Py_DECREF(iterable);
1709 if (lz->active == NULL) {
1710 Py_CLEAR(lz->source);
1711 return NULL; /* input not iterable */
1712 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001713 }
T. Woutersd694a062017-03-30 12:49:22 -07001714 item = PyIter_Next(lz->active);
1715 if (item != NULL)
1716 return item;
1717 if (PyErr_Occurred()) {
1718 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1719 PyErr_Clear();
1720 else
1721 return NULL; /* input raised an exception */
1722 }
1723 /* lz->active is consumed, try with the next iterable. */
1724 Py_CLEAR(lz->active);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001725 }
T. Woutersd694a062017-03-30 12:49:22 -07001726 /* Everything had been consumed already. */
1727 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001728}
1729
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001730PyDoc_STRVAR(chain_doc,
1731"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001732\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001733Return a chain object whose .next() method returns elements from the\n\
1734first iterable until it is exhausted, then elements from the next\n\
1735iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001736
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001737PyDoc_STRVAR(chain_from_iterable_doc,
1738"chain.from_iterable(iterable) --> chain object\n\
1739\n\
Martin Pantera850ef62016-07-28 01:11:04 +00001740Alternate chain() constructor taking a single iterable argument\n\
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001741that evaluates lazily.");
1742
1743static PyMethodDef chain_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001744 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1745 chain_from_iterable_doc},
1746 {NULL, NULL} /* sentinel */
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001747};
1748
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001749static PyTypeObject chain_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001750 PyVarObject_HEAD_INIT(NULL, 0)
1751 "itertools.chain", /* tp_name */
1752 sizeof(chainobject), /* tp_basicsize */
1753 0, /* tp_itemsize */
1754 /* methods */
1755 (destructor)chain_dealloc, /* tp_dealloc */
1756 0, /* tp_print */
1757 0, /* tp_getattr */
1758 0, /* tp_setattr */
1759 0, /* tp_compare */
1760 0, /* tp_repr */
1761 0, /* tp_as_number */
1762 0, /* tp_as_sequence */
1763 0, /* tp_as_mapping */
1764 0, /* tp_hash */
1765 0, /* tp_call */
1766 0, /* tp_str */
1767 PyObject_GenericGetAttr, /* tp_getattro */
1768 0, /* tp_setattro */
1769 0, /* tp_as_buffer */
1770 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1771 Py_TPFLAGS_BASETYPE, /* tp_flags */
1772 chain_doc, /* tp_doc */
1773 (traverseproc)chain_traverse, /* tp_traverse */
1774 0, /* tp_clear */
1775 0, /* tp_richcompare */
1776 0, /* tp_weaklistoffset */
1777 PyObject_SelfIter, /* tp_iter */
1778 (iternextfunc)chain_next, /* tp_iternext */
1779 chain_methods, /* tp_methods */
1780 0, /* tp_members */
1781 0, /* tp_getset */
1782 0, /* tp_base */
1783 0, /* tp_dict */
1784 0, /* tp_descr_get */
1785 0, /* tp_descr_set */
1786 0, /* tp_dictoffset */
1787 0, /* tp_init */
1788 0, /* tp_alloc */
1789 chain_new, /* tp_new */
1790 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001791};
1792
1793
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001794/* product object ************************************************************/
1795
1796typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001797 PyObject_HEAD
1798 PyObject *pools; /* tuple of pool tuples */
1799 Py_ssize_t *indices; /* one index per pool */
1800 PyObject *result; /* most recently returned result tuple */
1801 int stopped; /* set to 1 when the product iterator is exhausted */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001802} productobject;
1803
1804static PyTypeObject product_type;
1805
1806static PyObject *
1807product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1808{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001809 productobject *lz;
1810 Py_ssize_t nargs, npools, repeat=1;
1811 PyObject *pools = NULL;
1812 Py_ssize_t *indices = NULL;
1813 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001814
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001815 if (kwds != NULL) {
1816 char *kwlist[] = {"repeat", 0};
1817 PyObject *tmpargs = PyTuple_New(0);
1818 if (tmpargs == NULL)
1819 return NULL;
1820 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1821 Py_DECREF(tmpargs);
1822 return NULL;
1823 }
1824 Py_DECREF(tmpargs);
1825 if (repeat < 0) {
1826 PyErr_SetString(PyExc_ValueError,
1827 "repeat argument cannot be negative");
1828 return NULL;
1829 }
1830 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001831
Benjamin Petersondda91212015-02-01 21:34:07 -05001832 assert(PyTuple_CheckExact(args));
1833 if (repeat == 0) {
1834 nargs = 0;
1835 } else {
1836 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02001837 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Petersondda91212015-02-01 21:34:07 -05001838 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
1839 return NULL;
1840 }
1841 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001842 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001843
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02001844 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001845 if (indices == NULL) {
1846 PyErr_NoMemory();
1847 goto error;
1848 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001849
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001850 pools = PyTuple_New(npools);
1851 if (pools == NULL)
1852 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001853
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001854 for (i=0; i < nargs ; ++i) {
1855 PyObject *item = PyTuple_GET_ITEM(args, i);
1856 PyObject *pool = PySequence_Tuple(item);
1857 if (pool == NULL)
1858 goto error;
1859 PyTuple_SET_ITEM(pools, i, pool);
1860 indices[i] = 0;
1861 }
1862 for ( ; i < npools; ++i) {
1863 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1864 Py_INCREF(pool);
1865 PyTuple_SET_ITEM(pools, i, pool);
1866 indices[i] = 0;
1867 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001868
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001869 /* create productobject structure */
1870 lz = (productobject *)type->tp_alloc(type, 0);
1871 if (lz == NULL)
1872 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001873
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001874 lz->pools = pools;
1875 lz->indices = indices;
1876 lz->result = NULL;
1877 lz->stopped = 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001878
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001879 return (PyObject *)lz;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001880
1881error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001882 if (indices != NULL)
1883 PyMem_Free(indices);
1884 Py_XDECREF(pools);
1885 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001886}
1887
1888static void
1889product_dealloc(productobject *lz)
1890{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001891 PyObject_GC_UnTrack(lz);
1892 Py_XDECREF(lz->pools);
1893 Py_XDECREF(lz->result);
1894 if (lz->indices != NULL)
1895 PyMem_Free(lz->indices);
1896 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001897}
1898
1899static int
1900product_traverse(productobject *lz, visitproc visit, void *arg)
1901{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001902 Py_VISIT(lz->pools);
1903 Py_VISIT(lz->result);
1904 return 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001905}
1906
1907static PyObject *
1908product_next(productobject *lz)
1909{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001910 PyObject *pool;
1911 PyObject *elem;
1912 PyObject *oldelem;
1913 PyObject *pools = lz->pools;
1914 PyObject *result = lz->result;
1915 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1916 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001917
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001918 if (lz->stopped)
1919 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001920
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001921 if (result == NULL) {
1922 /* On the first pass, return an initial tuple filled with the
1923 first element from each pool. */
1924 result = PyTuple_New(npools);
1925 if (result == NULL)
1926 goto empty;
1927 lz->result = result;
1928 for (i=0; i < npools; i++) {
1929 pool = PyTuple_GET_ITEM(pools, i);
1930 if (PyTuple_GET_SIZE(pool) == 0)
1931 goto empty;
1932 elem = PyTuple_GET_ITEM(pool, 0);
1933 Py_INCREF(elem);
1934 PyTuple_SET_ITEM(result, i, elem);
1935 }
1936 } else {
1937 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001938
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001939 /* Copy the previous result tuple or re-use it if available */
1940 if (Py_REFCNT(result) > 1) {
1941 PyObject *old_result = result;
1942 result = PyTuple_New(npools);
1943 if (result == NULL)
1944 goto empty;
1945 lz->result = result;
1946 for (i=0; i < npools; i++) {
1947 elem = PyTuple_GET_ITEM(old_result, i);
1948 Py_INCREF(elem);
1949 PyTuple_SET_ITEM(result, i, elem);
1950 }
1951 Py_DECREF(old_result);
1952 }
1953 /* Now, we've got the only copy so we can update it in-place */
1954 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001955
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001956 /* Update the pool indices right-to-left. Only advance to the
1957 next pool when the previous one rolls-over */
1958 for (i=npools-1 ; i >= 0 ; i--) {
1959 pool = PyTuple_GET_ITEM(pools, i);
1960 indices[i]++;
1961 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1962 /* Roll-over and advance to next pool */
1963 indices[i] = 0;
1964 elem = PyTuple_GET_ITEM(pool, 0);
1965 Py_INCREF(elem);
1966 oldelem = PyTuple_GET_ITEM(result, i);
1967 PyTuple_SET_ITEM(result, i, elem);
1968 Py_DECREF(oldelem);
1969 } else {
1970 /* No rollover. Just increment and stop here. */
1971 elem = PyTuple_GET_ITEM(pool, indices[i]);
1972 Py_INCREF(elem);
1973 oldelem = PyTuple_GET_ITEM(result, i);
1974 PyTuple_SET_ITEM(result, i, elem);
1975 Py_DECREF(oldelem);
1976 break;
1977 }
1978 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00001979
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001980 /* If i is negative, then the indices have all rolled-over
1981 and we're done. */
1982 if (i < 0)
1983 goto empty;
1984 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001985
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001986 Py_INCREF(result);
1987 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001988
1989empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00001990 lz->stopped = 1;
1991 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001992}
1993
1994PyDoc_STRVAR(product_doc,
1995"product(*iterables) --> product object\n\
1996\n\
Raymond Hettinger73d79632008-02-23 02:20:41 +00001997Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001998For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1999The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2000cycle in a manner similar to an odometer (with the rightmost element changing\n\
2001on every iteration).\n\n\
Raymond Hettinger94034ea2009-02-09 18:39:41 +00002002To compute the product of an iterable with itself, specify the number\n\
2003of repetitions with the optional repeat keyword argument. For example,\n\
2004product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002005product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2006product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2007
2008static PyTypeObject product_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002009 PyVarObject_HEAD_INIT(NULL, 0)
2010 "itertools.product", /* tp_name */
2011 sizeof(productobject), /* tp_basicsize */
2012 0, /* tp_itemsize */
2013 /* methods */
2014 (destructor)product_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 product_doc, /* tp_doc */
2032 (traverseproc)product_traverse, /* tp_traverse */
2033 0, /* tp_clear */
2034 0, /* tp_richcompare */
2035 0, /* tp_weaklistoffset */
2036 PyObject_SelfIter, /* tp_iter */
2037 (iternextfunc)product_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 */
2047 0, /* tp_alloc */
2048 product_new, /* tp_new */
2049 PyObject_GC_Del, /* tp_free */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002050};
2051
2052
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002053/* combinations object ************************************************************/
2054
2055typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002056 PyObject_HEAD
2057 PyObject *pool; /* input converted to a tuple */
2058 Py_ssize_t *indices; /* one index per result element */
2059 PyObject *result; /* most recently returned result tuple */
2060 Py_ssize_t r; /* size of result tuple */
2061 int stopped; /* set to 1 when the combinations iterator is exhausted */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002062} combinationsobject;
2063
2064static PyTypeObject combinations_type;
2065
2066static PyObject *
2067combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2068{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002069 combinationsobject *co;
2070 Py_ssize_t n;
2071 Py_ssize_t r;
2072 PyObject *pool = NULL;
2073 PyObject *iterable = NULL;
2074 Py_ssize_t *indices = NULL;
2075 Py_ssize_t i;
2076 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002077
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002078 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2079 &iterable, &r))
2080 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002081
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002082 pool = PySequence_Tuple(iterable);
2083 if (pool == NULL)
2084 goto error;
2085 n = PyTuple_GET_SIZE(pool);
2086 if (r < 0) {
2087 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2088 goto error;
2089 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002090
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02002091 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002092 if (indices == NULL) {
2093 PyErr_NoMemory();
2094 goto error;
2095 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002096
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002097 for (i=0 ; i<r ; i++)
2098 indices[i] = i;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002099
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002100 /* create combinationsobject structure */
2101 co = (combinationsobject *)type->tp_alloc(type, 0);
2102 if (co == NULL)
2103 goto error;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002104
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002105 co->pool = pool;
2106 co->indices = indices;
2107 co->result = NULL;
2108 co->r = r;
2109 co->stopped = r > n ? 1 : 0;
2110
2111 return (PyObject *)co;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002112
2113error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002114 if (indices != NULL)
2115 PyMem_Free(indices);
2116 Py_XDECREF(pool);
2117 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002118}
2119
2120static void
2121combinations_dealloc(combinationsobject *co)
2122{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002123 PyObject_GC_UnTrack(co);
2124 Py_XDECREF(co->pool);
2125 Py_XDECREF(co->result);
2126 if (co->indices != NULL)
2127 PyMem_Free(co->indices);
2128 Py_TYPE(co)->tp_free(co);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002129}
2130
2131static int
2132combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2133{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002134 Py_VISIT(co->pool);
2135 Py_VISIT(co->result);
2136 return 0;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002137}
2138
2139static PyObject *
2140combinations_next(combinationsobject *co)
2141{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002142 PyObject *elem;
2143 PyObject *oldelem;
2144 PyObject *pool = co->pool;
2145 Py_ssize_t *indices = co->indices;
2146 PyObject *result = co->result;
2147 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2148 Py_ssize_t r = co->r;
2149 Py_ssize_t i, j, index;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002150
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002151 if (co->stopped)
2152 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002153
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002154 if (result == NULL) {
2155 /* On the first pass, initialize result tuple using the indices */
2156 result = PyTuple_New(r);
2157 if (result == NULL)
2158 goto empty;
2159 co->result = result;
2160 for (i=0; i<r ; i++) {
2161 index = indices[i];
2162 elem = PyTuple_GET_ITEM(pool, index);
2163 Py_INCREF(elem);
2164 PyTuple_SET_ITEM(result, i, elem);
2165 }
2166 } else {
2167 /* Copy the previous result tuple or re-use it if available */
2168 if (Py_REFCNT(result) > 1) {
2169 PyObject *old_result = result;
2170 result = PyTuple_New(r);
2171 if (result == NULL)
2172 goto empty;
2173 co->result = result;
2174 for (i=0; i<r ; i++) {
2175 elem = PyTuple_GET_ITEM(old_result, i);
2176 Py_INCREF(elem);
2177 PyTuple_SET_ITEM(result, i, elem);
2178 }
2179 Py_DECREF(old_result);
2180 }
2181 /* Now, we've got the only copy so we can update it in-place
2182 * CPython's empty tuple is a singleton and cached in
2183 * PyTuple's freelist.
2184 */
2185 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002186
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002187 /* Scan indices right-to-left until finding one that is not
2188 at its maximum (i + n - r). */
2189 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2190 ;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002191
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002192 /* If i is negative, then the indices are all at
2193 their maximum value and we're done. */
2194 if (i < 0)
2195 goto empty;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002196
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002197 /* Increment the current index which we know is not at its
2198 maximum. Then move back to the right setting each index
2199 to its lowest possible value (one higher than the index
2200 to its left -- this maintains the sort order invariant). */
2201 indices[i]++;
2202 for (j=i+1 ; j<r ; j++)
2203 indices[j] = indices[j-1] + 1;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002204
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002205 /* Update the result tuple for the new indices
2206 starting with i, the leftmost index that changed */
2207 for ( ; i<r ; i++) {
2208 index = indices[i];
2209 elem = PyTuple_GET_ITEM(pool, index);
2210 Py_INCREF(elem);
2211 oldelem = PyTuple_GET_ITEM(result, i);
2212 PyTuple_SET_ITEM(result, i, elem);
2213 Py_DECREF(oldelem);
2214 }
2215 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002216
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002217 Py_INCREF(result);
2218 return result;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002219
2220empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002221 co->stopped = 1;
2222 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002223}
2224
2225PyDoc_STRVAR(combinations_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002226"combinations(iterable, r) --> combinations object\n\
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002227\n\
2228Return successive r-length combinations of elements in the iterable.\n\n\
2229combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2230
2231static PyTypeObject combinations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002232 PyVarObject_HEAD_INIT(NULL, 0)
2233 "itertools.combinations", /* tp_name */
2234 sizeof(combinationsobject), /* tp_basicsize */
2235 0, /* tp_itemsize */
2236 /* methods */
2237 (destructor)combinations_dealloc, /* tp_dealloc */
2238 0, /* tp_print */
2239 0, /* tp_getattr */
2240 0, /* tp_setattr */
2241 0, /* tp_compare */
2242 0, /* tp_repr */
2243 0, /* tp_as_number */
2244 0, /* tp_as_sequence */
2245 0, /* tp_as_mapping */
2246 0, /* tp_hash */
2247 0, /* tp_call */
2248 0, /* tp_str */
2249 PyObject_GenericGetAttr, /* tp_getattro */
2250 0, /* tp_setattro */
2251 0, /* tp_as_buffer */
2252 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2253 Py_TPFLAGS_BASETYPE, /* tp_flags */
2254 combinations_doc, /* tp_doc */
2255 (traverseproc)combinations_traverse, /* tp_traverse */
2256 0, /* tp_clear */
2257 0, /* tp_richcompare */
2258 0, /* tp_weaklistoffset */
2259 PyObject_SelfIter, /* tp_iter */
2260 (iternextfunc)combinations_next, /* tp_iternext */
2261 0, /* tp_methods */
2262 0, /* tp_members */
2263 0, /* tp_getset */
2264 0, /* tp_base */
2265 0, /* tp_dict */
2266 0, /* tp_descr_get */
2267 0, /* tp_descr_set */
2268 0, /* tp_dictoffset */
2269 0, /* tp_init */
2270 0, /* tp_alloc */
2271 combinations_new, /* tp_new */
2272 PyObject_GC_Del, /* tp_free */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002273};
2274
2275
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002276/* combinations with replacement object *******************************************/
2277
2278/* Equivalent to:
2279
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002280 def combinations_with_replacement(iterable, r):
2281 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2282 # number items returned: (n+r-1)! / r! / (n-1)!
2283 pool = tuple(iterable)
2284 n = len(pool)
2285 indices = [0] * r
2286 yield tuple(pool[i] for i in indices)
2287 while 1:
2288 for i in reversed(range(r)):
2289 if indices[i] != n - 1:
2290 break
2291 else:
2292 return
2293 indices[i:] = [indices[i] + 1] * (r - i)
2294 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002295
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002296 def combinations_with_replacement2(iterable, r):
2297 'Alternate version that filters from product()'
2298 pool = tuple(iterable)
2299 n = len(pool)
2300 for indices in product(range(n), repeat=r):
2301 if sorted(indices) == list(indices):
2302 yield tuple(pool[i] for i in indices)
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002303*/
2304typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002305 PyObject_HEAD
2306 PyObject *pool; /* input converted to a tuple */
2307 Py_ssize_t *indices; /* one index per result element */
2308 PyObject *result; /* most recently returned result tuple */
2309 Py_ssize_t r; /* size of result tuple */
2310 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002311} cwrobject;
2312
2313static PyTypeObject cwr_type;
2314
2315static PyObject *
2316cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2317{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002318 cwrobject *co;
2319 Py_ssize_t n;
2320 Py_ssize_t r;
2321 PyObject *pool = NULL;
2322 PyObject *iterable = NULL;
2323 Py_ssize_t *indices = NULL;
2324 Py_ssize_t i;
2325 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002326
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002327 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2328 &iterable, &r))
2329 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002330
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002331 pool = PySequence_Tuple(iterable);
2332 if (pool == NULL)
2333 goto error;
2334 n = PyTuple_GET_SIZE(pool);
2335 if (r < 0) {
2336 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2337 goto error;
2338 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002339
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02002340 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002341 if (indices == NULL) {
2342 PyErr_NoMemory();
2343 goto error;
2344 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002345
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002346 for (i=0 ; i<r ; i++)
2347 indices[i] = 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002348
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002349 /* create cwrobject structure */
2350 co = (cwrobject *)type->tp_alloc(type, 0);
2351 if (co == NULL)
2352 goto error;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002353
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002354 co->pool = pool;
2355 co->indices = indices;
2356 co->result = NULL;
2357 co->r = r;
2358 co->stopped = !n && r;
2359
2360 return (PyObject *)co;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002361
2362error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002363 if (indices != NULL)
2364 PyMem_Free(indices);
2365 Py_XDECREF(pool);
2366 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002367}
2368
2369static void
2370cwr_dealloc(cwrobject *co)
2371{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002372 PyObject_GC_UnTrack(co);
2373 Py_XDECREF(co->pool);
2374 Py_XDECREF(co->result);
2375 if (co->indices != NULL)
2376 PyMem_Free(co->indices);
2377 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002378}
2379
2380static int
2381cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2382{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002383 Py_VISIT(co->pool);
2384 Py_VISIT(co->result);
2385 return 0;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002386}
2387
2388static PyObject *
2389cwr_next(cwrobject *co)
2390{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002391 PyObject *elem;
2392 PyObject *oldelem;
2393 PyObject *pool = co->pool;
2394 Py_ssize_t *indices = co->indices;
2395 PyObject *result = co->result;
2396 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2397 Py_ssize_t r = co->r;
2398 Py_ssize_t i, j, index;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002399
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002400 if (co->stopped)
2401 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002402
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002403 if (result == NULL) {
2404 /* On the first pass, initialize result tuple using the indices */
2405 result = PyTuple_New(r);
2406 if (result == NULL)
2407 goto empty;
2408 co->result = result;
2409 for (i=0; i<r ; i++) {
2410 index = indices[i];
2411 elem = PyTuple_GET_ITEM(pool, index);
2412 Py_INCREF(elem);
2413 PyTuple_SET_ITEM(result, i, elem);
2414 }
2415 } else {
2416 /* Copy the previous result tuple or re-use it if available */
2417 if (Py_REFCNT(result) > 1) {
2418 PyObject *old_result = result;
2419 result = PyTuple_New(r);
2420 if (result == NULL)
2421 goto empty;
2422 co->result = result;
2423 for (i=0; i<r ; i++) {
2424 elem = PyTuple_GET_ITEM(old_result, i);
2425 Py_INCREF(elem);
2426 PyTuple_SET_ITEM(result, i, elem);
2427 }
2428 Py_DECREF(old_result);
2429 }
2430 /* Now, we've got the only copy so we can update it in-place CPython's
2431 empty tuple is a singleton and cached in PyTuple's freelist. */
2432 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002433
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002434 /* Scan indices right-to-left until finding one that is not
2435 * at its maximum (n-1). */
2436 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2437 ;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002438
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002439 /* If i is negative, then the indices are all at
2440 their maximum value and we're done. */
2441 if (i < 0)
2442 goto empty;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002443
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002444 /* Increment the current index which we know is not at its
2445 maximum. Then set all to the right to the same value. */
2446 indices[i]++;
2447 for (j=i+1 ; j<r ; j++)
2448 indices[j] = indices[j-1];
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002449
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002450 /* Update the result tuple for the new indices
2451 starting with i, the leftmost index that changed */
2452 for ( ; i<r ; i++) {
2453 index = indices[i];
2454 elem = PyTuple_GET_ITEM(pool, index);
2455 Py_INCREF(elem);
2456 oldelem = PyTuple_GET_ITEM(result, i);
2457 PyTuple_SET_ITEM(result, i, elem);
2458 Py_DECREF(oldelem);
2459 }
2460 }
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002461
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002462 Py_INCREF(result);
2463 return result;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002464
2465empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002466 co->stopped = 1;
2467 return NULL;
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002468}
2469
2470PyDoc_STRVAR(cwr_doc,
Raymond Hettinger9eac1192009-11-19 01:22:04 +00002471"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002472\n\
2473Return successive r-length combinations of elements in the iterable\n\
2474allowing individual elements to have successive repeats.\n\
2475combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2476
2477static PyTypeObject cwr_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002478 PyVarObject_HEAD_INIT(NULL, 0)
2479 "itertools.combinations_with_replacement", /* tp_name */
2480 sizeof(cwrobject), /* tp_basicsize */
2481 0, /* tp_itemsize */
2482 /* methods */
2483 (destructor)cwr_dealloc, /* tp_dealloc */
2484 0, /* tp_print */
2485 0, /* tp_getattr */
2486 0, /* tp_setattr */
2487 0, /* tp_compare */
2488 0, /* tp_repr */
2489 0, /* tp_as_number */
2490 0, /* tp_as_sequence */
2491 0, /* tp_as_mapping */
2492 0, /* tp_hash */
2493 0, /* tp_call */
2494 0, /* tp_str */
2495 PyObject_GenericGetAttr, /* tp_getattro */
2496 0, /* tp_setattro */
2497 0, /* tp_as_buffer */
2498 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2499 Py_TPFLAGS_BASETYPE, /* tp_flags */
2500 cwr_doc, /* tp_doc */
2501 (traverseproc)cwr_traverse, /* tp_traverse */
2502 0, /* tp_clear */
2503 0, /* tp_richcompare */
2504 0, /* tp_weaklistoffset */
2505 PyObject_SelfIter, /* tp_iter */
2506 (iternextfunc)cwr_next, /* tp_iternext */
2507 0, /* tp_methods */
2508 0, /* tp_members */
2509 0, /* tp_getset */
2510 0, /* tp_base */
2511 0, /* tp_dict */
2512 0, /* tp_descr_get */
2513 0, /* tp_descr_set */
2514 0, /* tp_dictoffset */
2515 0, /* tp_init */
2516 0, /* tp_alloc */
2517 cwr_new, /* tp_new */
2518 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd081abc2009-01-27 02:58:49 +00002519};
2520
2521
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002522/* permutations object ************************************************************
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002523
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002524def permutations(iterable, r=None):
2525 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2526 pool = tuple(iterable)
2527 n = len(pool)
2528 r = n if r is None else r
2529 indices = range(n)
2530 cycles = range(n-r+1, n+1)[::-1]
2531 yield tuple(pool[i] for i in indices[:r])
2532 while n:
Serhiy Storchakaf0aa88f2015-06-11 00:06:27 +03002533 for i in reversed(range(r)):
2534 cycles[i] -= 1
2535 if cycles[i] == 0:
2536 indices[i:] = indices[i+1:] + indices[i:i+1]
2537 cycles[i] = n - i
2538 else:
2539 j = cycles[i]
2540 indices[i], indices[-j] = indices[-j], indices[i]
2541 yield tuple(pool[i] for i in indices[:r])
2542 break
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002543 else:
Serhiy Storchakaf0aa88f2015-06-11 00:06:27 +03002544 return
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002545*/
2546
2547typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002548 PyObject_HEAD
2549 PyObject *pool; /* input converted to a tuple */
2550 Py_ssize_t *indices; /* one index per element in the pool */
2551 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2552 PyObject *result; /* most recently returned result tuple */
2553 Py_ssize_t r; /* size of result tuple */
2554 int stopped; /* set to 1 when the permutations iterator is exhausted */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002555} permutationsobject;
2556
2557static PyTypeObject permutations_type;
2558
2559static PyObject *
2560permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2561{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002562 permutationsobject *po;
2563 Py_ssize_t n;
2564 Py_ssize_t r;
2565 PyObject *robj = Py_None;
2566 PyObject *pool = NULL;
2567 PyObject *iterable = NULL;
2568 Py_ssize_t *indices = NULL;
2569 Py_ssize_t *cycles = NULL;
2570 Py_ssize_t i;
2571 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002572
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002573 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2574 &iterable, &robj))
2575 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002576
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002577 pool = PySequence_Tuple(iterable);
2578 if (pool == NULL)
2579 goto error;
2580 n = PyTuple_GET_SIZE(pool);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002581
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002582 r = n;
2583 if (robj != Py_None) {
2584 r = PyInt_AsSsize_t(robj);
2585 if (r == -1 && PyErr_Occurred())
2586 goto error;
2587 }
2588 if (r < 0) {
2589 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2590 goto error;
2591 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002592
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +02002593 indices = PyMem_New(Py_ssize_t, n);
2594 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002595 if (indices == NULL || cycles == NULL) {
2596 PyErr_NoMemory();
2597 goto error;
2598 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002599
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002600 for (i=0 ; i<n ; i++)
2601 indices[i] = i;
2602 for (i=0 ; i<r ; i++)
2603 cycles[i] = n - i;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002604
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002605 /* create permutationsobject structure */
2606 po = (permutationsobject *)type->tp_alloc(type, 0);
2607 if (po == NULL)
2608 goto error;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002609
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002610 po->pool = pool;
2611 po->indices = indices;
2612 po->cycles = cycles;
2613 po->result = NULL;
2614 po->r = r;
2615 po->stopped = r > n ? 1 : 0;
2616
2617 return (PyObject *)po;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002618
2619error:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002620 if (indices != NULL)
2621 PyMem_Free(indices);
2622 if (cycles != NULL)
2623 PyMem_Free(cycles);
2624 Py_XDECREF(pool);
2625 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002626}
2627
2628static void
2629permutations_dealloc(permutationsobject *po)
2630{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002631 PyObject_GC_UnTrack(po);
2632 Py_XDECREF(po->pool);
2633 Py_XDECREF(po->result);
2634 PyMem_Free(po->indices);
2635 PyMem_Free(po->cycles);
2636 Py_TYPE(po)->tp_free(po);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002637}
2638
2639static int
2640permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2641{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002642 Py_VISIT(po->pool);
2643 Py_VISIT(po->result);
2644 return 0;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002645}
2646
2647static PyObject *
2648permutations_next(permutationsobject *po)
2649{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002650 PyObject *elem;
2651 PyObject *oldelem;
2652 PyObject *pool = po->pool;
2653 Py_ssize_t *indices = po->indices;
2654 Py_ssize_t *cycles = po->cycles;
2655 PyObject *result = po->result;
2656 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2657 Py_ssize_t r = po->r;
2658 Py_ssize_t i, j, k, index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002659
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002660 if (po->stopped)
2661 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002662
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002663 if (result == NULL) {
2664 /* On the first pass, initialize result tuple using the indices */
2665 result = PyTuple_New(r);
2666 if (result == NULL)
2667 goto empty;
2668 po->result = result;
2669 for (i=0; i<r ; i++) {
2670 index = indices[i];
2671 elem = PyTuple_GET_ITEM(pool, index);
2672 Py_INCREF(elem);
2673 PyTuple_SET_ITEM(result, i, elem);
2674 }
2675 } else {
2676 if (n == 0)
2677 goto empty;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002678
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002679 /* Copy the previous result tuple or re-use it if available */
2680 if (Py_REFCNT(result) > 1) {
2681 PyObject *old_result = result;
2682 result = PyTuple_New(r);
2683 if (result == NULL)
2684 goto empty;
2685 po->result = result;
2686 for (i=0; i<r ; i++) {
2687 elem = PyTuple_GET_ITEM(old_result, i);
2688 Py_INCREF(elem);
2689 PyTuple_SET_ITEM(result, i, elem);
2690 }
2691 Py_DECREF(old_result);
2692 }
2693 /* Now, we've got the only copy so we can update it in-place */
2694 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002695
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002696 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2697 for (i=r-1 ; i>=0 ; i--) {
2698 cycles[i] -= 1;
2699 if (cycles[i] == 0) {
2700 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2701 index = indices[i];
2702 for (j=i ; j<n-1 ; j++)
2703 indices[j] = indices[j+1];
2704 indices[n-1] = index;
2705 cycles[i] = n - i;
2706 } else {
2707 j = cycles[i];
2708 index = indices[i];
2709 indices[i] = indices[n-j];
2710 indices[n-j] = index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002711
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002712 for (k=i; k<r ; k++) {
2713 /* start with i, the leftmost element that changed */
2714 /* yield tuple(pool[k] for k in indices[:r]) */
2715 index = indices[k];
2716 elem = PyTuple_GET_ITEM(pool, index);
2717 Py_INCREF(elem);
2718 oldelem = PyTuple_GET_ITEM(result, k);
2719 PyTuple_SET_ITEM(result, k, elem);
2720 Py_DECREF(oldelem);
2721 }
2722 break;
2723 }
2724 }
2725 /* If i is negative, then the cycles have all
2726 rolled-over and we're done. */
2727 if (i < 0)
2728 goto empty;
2729 }
2730 Py_INCREF(result);
2731 return result;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002732
2733empty:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002734 po->stopped = 1;
2735 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002736}
2737
2738PyDoc_STRVAR(permutations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +00002739"permutations(iterable[, r]) --> permutations object\n\
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002740\n\
2741Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl8550d692008-06-10 12:46:39 +00002742permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002743
2744static PyTypeObject permutations_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002745 PyVarObject_HEAD_INIT(NULL, 0)
2746 "itertools.permutations", /* tp_name */
2747 sizeof(permutationsobject), /* tp_basicsize */
2748 0, /* tp_itemsize */
2749 /* methods */
2750 (destructor)permutations_dealloc, /* tp_dealloc */
2751 0, /* tp_print */
2752 0, /* tp_getattr */
2753 0, /* tp_setattr */
2754 0, /* tp_compare */
2755 0, /* tp_repr */
2756 0, /* tp_as_number */
2757 0, /* tp_as_sequence */
2758 0, /* tp_as_mapping */
2759 0, /* tp_hash */
2760 0, /* tp_call */
2761 0, /* tp_str */
2762 PyObject_GenericGetAttr, /* tp_getattro */
2763 0, /* tp_setattro */
2764 0, /* tp_as_buffer */
2765 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2766 Py_TPFLAGS_BASETYPE, /* tp_flags */
2767 permutations_doc, /* tp_doc */
2768 (traverseproc)permutations_traverse, /* tp_traverse */
2769 0, /* tp_clear */
2770 0, /* tp_richcompare */
2771 0, /* tp_weaklistoffset */
2772 PyObject_SelfIter, /* tp_iter */
2773 (iternextfunc)permutations_next, /* tp_iternext */
2774 0, /* tp_methods */
2775 0, /* tp_members */
2776 0, /* tp_getset */
2777 0, /* tp_base */
2778 0, /* tp_dict */
2779 0, /* tp_descr_get */
2780 0, /* tp_descr_set */
2781 0, /* tp_dictoffset */
2782 0, /* tp_init */
2783 0, /* tp_alloc */
2784 permutations_new, /* tp_new */
2785 PyObject_GC_Del, /* tp_free */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002786};
2787
2788
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002789/* compress object ************************************************************/
2790
2791/* Equivalent to:
2792
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002793 def compress(data, selectors):
2794 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2795 return (d for d, s in izip(data, selectors) if s)
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002796*/
2797
2798typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002799 PyObject_HEAD
2800 PyObject *data;
2801 PyObject *selectors;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002802} compressobject;
2803
2804static PyTypeObject compress_type;
2805
2806static PyObject *
2807compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2808{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002809 PyObject *seq1, *seq2;
2810 PyObject *data=NULL, *selectors=NULL;
2811 compressobject *lz;
2812 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002813
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002814 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2815 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002816
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002817 data = PyObject_GetIter(seq1);
2818 if (data == NULL)
2819 goto fail;
2820 selectors = PyObject_GetIter(seq2);
2821 if (selectors == NULL)
2822 goto fail;
2823
2824 /* create compressobject structure */
2825 lz = (compressobject *)type->tp_alloc(type, 0);
2826 if (lz == NULL)
2827 goto fail;
2828 lz->data = data;
2829 lz->selectors = selectors;
2830 return (PyObject *)lz;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002831
2832fail:
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002833 Py_XDECREF(data);
2834 Py_XDECREF(selectors);
2835 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002836}
2837
2838static void
2839compress_dealloc(compressobject *lz)
2840{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002841 PyObject_GC_UnTrack(lz);
2842 Py_XDECREF(lz->data);
2843 Py_XDECREF(lz->selectors);
2844 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002845}
2846
2847static int
2848compress_traverse(compressobject *lz, visitproc visit, void *arg)
2849{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002850 Py_VISIT(lz->data);
2851 Py_VISIT(lz->selectors);
2852 return 0;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002853}
2854
2855static PyObject *
2856compress_next(compressobject *lz)
2857{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002858 PyObject *data = lz->data, *selectors = lz->selectors;
2859 PyObject *datum, *selector;
2860 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2861 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2862 int ok;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002863
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002864 while (1) {
2865 /* Steps: get datum, get selector, evaluate selector.
2866 Order is important (to match the pure python version
2867 in terms of which input gets a chance to raise an
2868 exception first).
2869 */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002870
Serhiy Storchakabb845652013-04-06 22:04:10 +03002871 datum = datanext(data);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03002872 if (datum == NULL)
Serhiy Storchakabb845652013-04-06 22:04:10 +03002873 return NULL;
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002874
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002875 selector = selectornext(selectors);
2876 if (selector == NULL) {
2877 Py_DECREF(datum);
2878 return NULL;
2879 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002880
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002881 ok = PyObject_IsTrue(selector);
2882 Py_DECREF(selector);
2883 if (ok == 1)
2884 return datum;
2885 Py_DECREF(datum);
2886 if (ok == -1)
2887 return NULL;
2888 }
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002889}
2890
2891PyDoc_STRVAR(compress_doc,
Raymond Hettinger2e2909f2009-02-19 02:15:14 +00002892"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002893\n\
2894Return data elements corresponding to true selector elements.\n\
2895Forms a shorter iterator from selected data elements using the\n\
2896selectors to choose the data elements.");
2897
2898static PyTypeObject compress_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002899 PyVarObject_HEAD_INIT(NULL, 0)
2900 "itertools.compress", /* tp_name */
2901 sizeof(compressobject), /* tp_basicsize */
2902 0, /* tp_itemsize */
2903 /* methods */
2904 (destructor)compress_dealloc, /* tp_dealloc */
2905 0, /* tp_print */
2906 0, /* tp_getattr */
2907 0, /* tp_setattr */
2908 0, /* tp_compare */
2909 0, /* tp_repr */
2910 0, /* tp_as_number */
2911 0, /* tp_as_sequence */
2912 0, /* tp_as_mapping */
2913 0, /* tp_hash */
2914 0, /* tp_call */
2915 0, /* tp_str */
2916 PyObject_GenericGetAttr, /* tp_getattro */
2917 0, /* tp_setattro */
2918 0, /* tp_as_buffer */
2919 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2920 Py_TPFLAGS_BASETYPE, /* tp_flags */
2921 compress_doc, /* tp_doc */
2922 (traverseproc)compress_traverse, /* tp_traverse */
2923 0, /* tp_clear */
2924 0, /* tp_richcompare */
2925 0, /* tp_weaklistoffset */
2926 PyObject_SelfIter, /* tp_iter */
2927 (iternextfunc)compress_next, /* tp_iternext */
2928 0, /* tp_methods */
2929 0, /* tp_members */
2930 0, /* tp_getset */
2931 0, /* tp_base */
2932 0, /* tp_dict */
2933 0, /* tp_descr_get */
2934 0, /* tp_descr_set */
2935 0, /* tp_dictoffset */
2936 0, /* tp_init */
2937 0, /* tp_alloc */
2938 compress_new, /* tp_new */
2939 PyObject_GC_Del, /* tp_free */
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00002940};
2941
2942
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002943/* ifilter object ************************************************************/
2944
2945typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002946 PyObject_HEAD
2947 PyObject *func;
2948 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002949} ifilterobject;
2950
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002951static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002952
2953static PyObject *
2954ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2955{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002956 PyObject *func, *seq;
2957 PyObject *it;
2958 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002959
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002960 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
2961 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002962
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002963 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
2964 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002965
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002966 /* Get iterator. */
2967 it = PyObject_GetIter(seq);
2968 if (it == NULL)
2969 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002970
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002971 /* create ifilterobject structure */
2972 lz = (ifilterobject *)type->tp_alloc(type, 0);
2973 if (lz == NULL) {
2974 Py_DECREF(it);
2975 return NULL;
2976 }
2977 Py_INCREF(func);
2978 lz->func = func;
2979 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002980
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002981 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002982}
2983
2984static void
2985ifilter_dealloc(ifilterobject *lz)
2986{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002987 PyObject_GC_UnTrack(lz);
2988 Py_XDECREF(lz->func);
2989 Py_XDECREF(lz->it);
2990 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002991}
2992
2993static int
2994ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
2995{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00002996 Py_VISIT(lz->it);
2997 Py_VISIT(lz->func);
2998 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002999}
3000
3001static PyObject *
3002ifilter_next(ifilterobject *lz)
3003{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003004 PyObject *item;
3005 PyObject *it = lz->it;
3006 long ok;
3007 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003008
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003009 iternext = *Py_TYPE(it)->tp_iternext;
3010 for (;;) {
3011 item = iternext(it);
3012 if (item == NULL)
3013 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003014
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003015 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3016 ok = PyObject_IsTrue(item);
3017 } else {
3018 PyObject *good;
3019 good = PyObject_CallFunctionObjArgs(lz->func,
3020 item, NULL);
3021 if (good == NULL) {
3022 Py_DECREF(item);
3023 return NULL;
3024 }
3025 ok = PyObject_IsTrue(good);
3026 Py_DECREF(good);
3027 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003028 if (ok > 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003029 return item;
3030 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003031 if (ok < 0)
3032 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003033 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003034}
3035
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003036PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003037"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003038\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003039Return those items of sequence for which function(item) is true.\n\
3040If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003041
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003042static PyTypeObject ifilter_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003043 PyVarObject_HEAD_INIT(NULL, 0)
3044 "itertools.ifilter", /* tp_name */
3045 sizeof(ifilterobject), /* tp_basicsize */
3046 0, /* tp_itemsize */
3047 /* methods */
3048 (destructor)ifilter_dealloc, /* tp_dealloc */
3049 0, /* tp_print */
3050 0, /* tp_getattr */
3051 0, /* tp_setattr */
3052 0, /* tp_compare */
3053 0, /* tp_repr */
3054 0, /* tp_as_number */
3055 0, /* tp_as_sequence */
3056 0, /* tp_as_mapping */
3057 0, /* tp_hash */
3058 0, /* tp_call */
3059 0, /* tp_str */
3060 PyObject_GenericGetAttr, /* tp_getattro */
3061 0, /* tp_setattro */
3062 0, /* tp_as_buffer */
3063 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3064 Py_TPFLAGS_BASETYPE, /* tp_flags */
3065 ifilter_doc, /* tp_doc */
3066 (traverseproc)ifilter_traverse, /* tp_traverse */
3067 0, /* tp_clear */
3068 0, /* tp_richcompare */
3069 0, /* tp_weaklistoffset */
3070 PyObject_SelfIter, /* tp_iter */
3071 (iternextfunc)ifilter_next, /* tp_iternext */
3072 0, /* tp_methods */
3073 0, /* tp_members */
3074 0, /* tp_getset */
3075 0, /* tp_base */
3076 0, /* tp_dict */
3077 0, /* tp_descr_get */
3078 0, /* tp_descr_set */
3079 0, /* tp_dictoffset */
3080 0, /* tp_init */
3081 0, /* tp_alloc */
3082 ifilter_new, /* tp_new */
3083 PyObject_GC_Del, /* tp_free */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003084};
3085
3086
3087/* ifilterfalse object ************************************************************/
3088
3089typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003090 PyObject_HEAD
3091 PyObject *func;
3092 PyObject *it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003093} ifilterfalseobject;
3094
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003095static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003096
3097static PyObject *
3098ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3099{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003100 PyObject *func, *seq;
3101 PyObject *it;
3102 ifilterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003103
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003104 if (type == &ifilterfalse_type &&
3105 !_PyArg_NoKeywords("ifilterfalse()", kwds))
3106 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003107
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003108 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
3109 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003110
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003111 /* Get iterator. */
3112 it = PyObject_GetIter(seq);
3113 if (it == NULL)
3114 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003115
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003116 /* create ifilterfalseobject structure */
3117 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
3118 if (lz == NULL) {
3119 Py_DECREF(it);
3120 return NULL;
3121 }
3122 Py_INCREF(func);
3123 lz->func = func;
3124 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003125
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003126 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003127}
3128
3129static void
3130ifilterfalse_dealloc(ifilterfalseobject *lz)
3131{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003132 PyObject_GC_UnTrack(lz);
3133 Py_XDECREF(lz->func);
3134 Py_XDECREF(lz->it);
3135 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003136}
3137
3138static int
3139ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
3140{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003141 Py_VISIT(lz->it);
3142 Py_VISIT(lz->func);
3143 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003144}
3145
3146static PyObject *
3147ifilterfalse_next(ifilterfalseobject *lz)
3148{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003149 PyObject *item;
3150 PyObject *it = lz->it;
3151 long ok;
3152 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003153
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003154 iternext = *Py_TYPE(it)->tp_iternext;
3155 for (;;) {
3156 item = iternext(it);
3157 if (item == NULL)
3158 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003159
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003160 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3161 ok = PyObject_IsTrue(item);
3162 } else {
3163 PyObject *good;
3164 good = PyObject_CallFunctionObjArgs(lz->func,
3165 item, NULL);
3166 if (good == NULL) {
3167 Py_DECREF(item);
3168 return NULL;
3169 }
3170 ok = PyObject_IsTrue(good);
3171 Py_DECREF(good);
3172 }
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003173 if (ok == 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003174 return item;
3175 Py_DECREF(item);
Antoine Pitrouc5bef752012-08-15 23:16:51 +02003176 if (ok < 0)
3177 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003178 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003179}
3180
Raymond Hettinger60eca932003-02-09 06:40:58 +00003181PyDoc_STRVAR(ifilterfalse_doc,
3182"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
3183\n\
3184Return those items of sequence for which function(item) is false.\n\
3185If function is None, return the items that are false.");
3186
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003187static PyTypeObject ifilterfalse_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003188 PyVarObject_HEAD_INIT(NULL, 0)
3189 "itertools.ifilterfalse", /* tp_name */
3190 sizeof(ifilterfalseobject), /* tp_basicsize */
3191 0, /* tp_itemsize */
3192 /* methods */
3193 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
3194 0, /* tp_print */
3195 0, /* tp_getattr */
3196 0, /* tp_setattr */
3197 0, /* tp_compare */
3198 0, /* tp_repr */
3199 0, /* tp_as_number */
3200 0, /* tp_as_sequence */
3201 0, /* tp_as_mapping */
3202 0, /* tp_hash */
3203 0, /* tp_call */
3204 0, /* tp_str */
3205 PyObject_GenericGetAttr, /* tp_getattro */
3206 0, /* tp_setattro */
3207 0, /* tp_as_buffer */
3208 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3209 Py_TPFLAGS_BASETYPE, /* tp_flags */
3210 ifilterfalse_doc, /* tp_doc */
3211 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
3212 0, /* tp_clear */
3213 0, /* tp_richcompare */
3214 0, /* tp_weaklistoffset */
3215 PyObject_SelfIter, /* tp_iter */
3216 (iternextfunc)ifilterfalse_next, /* tp_iternext */
3217 0, /* tp_methods */
3218 0, /* tp_members */
3219 0, /* tp_getset */
3220 0, /* tp_base */
3221 0, /* tp_dict */
3222 0, /* tp_descr_get */
3223 0, /* tp_descr_set */
3224 0, /* tp_dictoffset */
3225 0, /* tp_init */
3226 0, /* tp_alloc */
3227 ifilterfalse_new, /* tp_new */
3228 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003229};
3230
3231
3232/* count object ************************************************************/
3233
3234typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003235 PyObject_HEAD
3236 Py_ssize_t cnt;
3237 PyObject *long_cnt;
3238 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003239} countobject;
3240
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003241/* Counting logic and invariants:
3242
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003243fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003244
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003245 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3246 Advances with: cnt += 1
3247 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003248
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003249slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003250
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003251 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3252 All counting is done with python objects (no overflows or underflows).
3253 Advances with: long_cnt += long_step
3254 Step may be zero -- effectively a slow version of repeat(cnt).
3255 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003256*/
3257
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003258static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003259
3260static PyObject *
3261count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3262{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003263 countobject *lz;
3264 int slow_mode = 0;
3265 Py_ssize_t cnt = 0;
3266 PyObject *long_cnt = NULL;
3267 PyObject *long_step = NULL;
3268 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003269
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003270 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3271 kwlist, &long_cnt, &long_step))
3272 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003273
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003274 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3275 (long_step != NULL && !PyNumber_Check(long_step))) {
3276 PyErr_SetString(PyExc_TypeError, "a number is required");
3277 return NULL;
3278 }
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003279
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003280 if (long_cnt != NULL) {
3281 cnt = PyInt_AsSsize_t(long_cnt);
3282 if ((cnt == -1 && PyErr_Occurred()) || !PyInt_Check(long_cnt)) {
3283 PyErr_Clear();
3284 slow_mode = 1;
3285 }
3286 Py_INCREF(long_cnt);
3287 } else {
3288 cnt = 0;
3289 long_cnt = PyInt_FromLong(0);
3290 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003291
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003292 /* If not specified, step defaults to 1 */
3293 if (long_step == NULL) {
3294 long_step = PyInt_FromLong(1);
3295 if (long_step == NULL) {
3296 Py_DECREF(long_cnt);
3297 return NULL;
3298 }
3299 } else
3300 Py_INCREF(long_step);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003301
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003302 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003303
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003304 /* Fast mode only works when the step is 1 */
3305 if (!PyInt_Check(long_step) ||
3306 PyInt_AS_LONG(long_step) != 1) {
3307 slow_mode = 1;
3308 }
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003309
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003310 if (slow_mode)
3311 cnt = PY_SSIZE_T_MAX;
3312 else
3313 Py_CLEAR(long_cnt);
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003314
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003315 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3316 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3317 assert(slow_mode ||
3318 (PyInt_Check(long_step) && PyInt_AS_LONG(long_step) == 1));
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003319
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003320 /* create countobject structure */
3321 lz = (countobject *)type->tp_alloc(type, 0);
3322 if (lz == NULL) {
3323 Py_XDECREF(long_cnt);
3324 return NULL;
3325 }
3326 lz->cnt = cnt;
3327 lz->long_cnt = long_cnt;
3328 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003329
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003330 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003331}
3332
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003333static void
3334count_dealloc(countobject *lz)
3335{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003336 PyObject_GC_UnTrack(lz);
3337 Py_XDECREF(lz->long_cnt);
3338 Py_XDECREF(lz->long_step);
3339 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003340}
3341
Benjamin Peterson062a7c32009-02-16 21:07:52 +00003342static int
Raymond Hettingerb21d8102009-02-16 20:39:12 +00003343count_traverse(countobject *lz, visitproc visit, void *arg)
3344{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003345 Py_VISIT(lz->long_cnt);
3346 Py_VISIT(lz->long_step);
3347 return 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003348}
3349
3350static PyObject *
3351count_nextlong(countobject *lz)
3352{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003353 PyObject *long_cnt;
3354 PyObject *stepped_up;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003355
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003356 long_cnt = lz->long_cnt;
3357 if (long_cnt == NULL) {
3358 /* Switch to slow_mode */
3359 long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
3360 if (long_cnt == NULL)
3361 return NULL;
3362 }
3363 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003364
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003365 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3366 if (stepped_up == NULL)
3367 return NULL;
3368 lz->long_cnt = stepped_up;
3369 return long_cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003370}
3371
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003372static PyObject *
3373count_next(countobject *lz)
3374{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003375 if (lz->cnt == PY_SSIZE_T_MAX)
3376 return count_nextlong(lz);
3377 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003378}
3379
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003380static PyObject *
3381count_repr(countobject *lz)
3382{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003383 PyObject *cnt_repr, *step_repr = NULL;
3384 PyObject *result = NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003385
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003386 if (lz->cnt != PY_SSIZE_T_MAX)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003387 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00003388
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003389 cnt_repr = PyObject_Repr(lz->long_cnt);
3390 if (cnt_repr == NULL)
3391 return NULL;
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003392
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003393 if (PyInt_Check(lz->long_step) && PyInt_AS_LONG(lz->long_step) == 1) {
3394 /* Don't display step when it is an integer equal to 1 */
3395 result = PyString_FromFormat("count(%s)",
3396 PyString_AS_STRING(cnt_repr));
3397 } else {
3398 step_repr = PyObject_Repr(lz->long_step);
3399 if (step_repr != NULL)
3400 result = PyString_FromFormat("count(%s, %s)",
3401 PyString_AS_STRING(cnt_repr),
3402 PyString_AS_STRING(step_repr));
3403 }
3404 Py_DECREF(cnt_repr);
3405 Py_XDECREF(step_repr);
3406 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003407}
3408
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003409static PyObject *
3410count_reduce(countobject *lz)
3411{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003412 if (lz->cnt == PY_SSIZE_T_MAX)
3413 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3414 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003415}
3416
3417PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3418
3419static PyMethodDef count_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003420 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3421 count_reduce_doc},
3422 {NULL, NULL} /* sentinel */
Raymond Hettingere09f45a2009-11-30 19:44:40 +00003423};
3424
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003425PyDoc_STRVAR(count_doc,
Ezio Melottifdc1e0d2010-06-11 02:21:25 +00003426 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003427\n\
Raymond Hettingeraa681c72009-02-21 07:17:22 +00003428Return a count object whose .next() method returns consecutive values.\n\
3429Equivalent to:\n\n\
Raymond Hettinger31c769c2009-02-12 05:39:46 +00003430 def count(firstval=0, step=1):\n\
Eli Benderskye6802db2013-09-03 06:41:58 -07003431 x = firstval\n\
3432 while 1:\n\
3433 yield x\n\
3434 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003435
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003436static PyTypeObject count_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003437 PyVarObject_HEAD_INIT(NULL, 0)
3438 "itertools.count", /* tp_name */
3439 sizeof(countobject), /* tp_basicsize */
3440 0, /* tp_itemsize */
3441 /* methods */
3442 (destructor)count_dealloc, /* tp_dealloc */
3443 0, /* tp_print */
3444 0, /* tp_getattr */
3445 0, /* tp_setattr */
3446 0, /* tp_compare */
3447 (reprfunc)count_repr, /* tp_repr */
3448 0, /* tp_as_number */
3449 0, /* tp_as_sequence */
3450 0, /* tp_as_mapping */
3451 0, /* tp_hash */
3452 0, /* tp_call */
3453 0, /* tp_str */
3454 PyObject_GenericGetAttr, /* tp_getattro */
3455 0, /* tp_setattro */
3456 0, /* tp_as_buffer */
3457 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3458 Py_TPFLAGS_BASETYPE, /* tp_flags */
3459 count_doc, /* tp_doc */
3460 (traverseproc)count_traverse, /* tp_traverse */
3461 0, /* tp_clear */
3462 0, /* tp_richcompare */
3463 0, /* tp_weaklistoffset */
3464 PyObject_SelfIter, /* tp_iter */
3465 (iternextfunc)count_next, /* tp_iternext */
3466 count_methods, /* tp_methods */
3467 0, /* tp_members */
3468 0, /* tp_getset */
3469 0, /* tp_base */
3470 0, /* tp_dict */
3471 0, /* tp_descr_get */
3472 0, /* tp_descr_set */
3473 0, /* tp_dictoffset */
3474 0, /* tp_init */
3475 0, /* tp_alloc */
3476 count_new, /* tp_new */
3477 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003478};
3479
3480
3481/* izip object ************************************************************/
3482
3483#include "Python.h"
3484
3485typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003486 PyObject_HEAD
3487 Py_ssize_t tuplesize;
3488 PyObject *ittuple; /* tuple of iterators */
3489 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003490} izipobject;
3491
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003492static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003493
3494static PyObject *
3495izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3496{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003497 izipobject *lz;
3498 Py_ssize_t i;
3499 PyObject *ittuple; /* tuple of iterators */
3500 PyObject *result;
3501 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003502
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003503 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
3504 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003505
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003506 /* args must be a tuple */
3507 assert(PyTuple_Check(args));
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003508
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003509 /* obtain iterators */
3510 ittuple = PyTuple_New(tuplesize);
3511 if (ittuple == NULL)
3512 return NULL;
3513 for (i=0; i < tuplesize; ++i) {
3514 PyObject *item = PyTuple_GET_ITEM(args, i);
3515 PyObject *it = PyObject_GetIter(item);
3516 if (it == NULL) {
3517 if (PyErr_ExceptionMatches(PyExc_TypeError))
3518 PyErr_Format(PyExc_TypeError,
3519 "izip argument #%zd must support iteration",
3520 i+1);
3521 Py_DECREF(ittuple);
3522 return NULL;
3523 }
3524 PyTuple_SET_ITEM(ittuple, i, it);
3525 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003526
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003527 /* create a result holder */
3528 result = PyTuple_New(tuplesize);
3529 if (result == NULL) {
3530 Py_DECREF(ittuple);
3531 return NULL;
3532 }
3533 for (i=0 ; i < tuplesize ; i++) {
3534 Py_INCREF(Py_None);
3535 PyTuple_SET_ITEM(result, i, Py_None);
3536 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003537
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003538 /* create izipobject structure */
3539 lz = (izipobject *)type->tp_alloc(type, 0);
3540 if (lz == NULL) {
3541 Py_DECREF(ittuple);
3542 Py_DECREF(result);
3543 return NULL;
3544 }
3545 lz->ittuple = ittuple;
3546 lz->tuplesize = tuplesize;
3547 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003548
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003549 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003550}
3551
3552static void
3553izip_dealloc(izipobject *lz)
3554{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003555 PyObject_GC_UnTrack(lz);
3556 Py_XDECREF(lz->ittuple);
3557 Py_XDECREF(lz->result);
3558 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003559}
3560
3561static int
3562izip_traverse(izipobject *lz, visitproc visit, void *arg)
3563{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003564 Py_VISIT(lz->ittuple);
3565 Py_VISIT(lz->result);
3566 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003567}
3568
3569static PyObject *
3570izip_next(izipobject *lz)
3571{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003572 Py_ssize_t i;
3573 Py_ssize_t tuplesize = lz->tuplesize;
3574 PyObject *result = lz->result;
3575 PyObject *it;
3576 PyObject *item;
3577 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003578
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003579 if (tuplesize == 0)
3580 return NULL;
3581 if (Py_REFCNT(result) == 1) {
3582 Py_INCREF(result);
3583 for (i=0 ; i < tuplesize ; i++) {
3584 it = PyTuple_GET_ITEM(lz->ittuple, i);
3585 item = (*Py_TYPE(it)->tp_iternext)(it);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003586 if (item == NULL) {
3587 Py_DECREF(result);
3588 return NULL;
3589 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003590 olditem = PyTuple_GET_ITEM(result, i);
3591 PyTuple_SET_ITEM(result, i, item);
3592 Py_DECREF(olditem);
3593 }
3594 } else {
3595 result = PyTuple_New(tuplesize);
3596 if (result == NULL)
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003597 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003598 for (i=0 ; i < tuplesize ; i++) {
3599 it = PyTuple_GET_ITEM(lz->ittuple, i);
3600 item = (*Py_TYPE(it)->tp_iternext)(it);
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03003601 if (item == NULL) {
3602 Py_DECREF(result);
3603 return NULL;
3604 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003605 PyTuple_SET_ITEM(result, i, item);
3606 }
3607 }
3608 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003609}
3610
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003611PyDoc_STRVAR(izip_doc,
3612"izip(iter1 [,iter2 [...]]) --> izip object\n\
3613\n\
3614Return a izip object whose .next() method returns a tuple where\n\
3615the i-th element comes from the i-th iterable argument. The .next()\n\
3616method continues until the shortest iterable in the argument sequence\n\
3617is exhausted and then it raises StopIteration. Works like the zip()\n\
3618function but consumes less memory by returning an iterator instead of\n\
3619a list.");
3620
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003621static PyTypeObject izip_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003622 PyVarObject_HEAD_INIT(NULL, 0)
3623 "itertools.izip", /* tp_name */
3624 sizeof(izipobject), /* tp_basicsize */
3625 0, /* tp_itemsize */
3626 /* methods */
3627 (destructor)izip_dealloc, /* tp_dealloc */
3628 0, /* tp_print */
3629 0, /* tp_getattr */
3630 0, /* tp_setattr */
3631 0, /* tp_compare */
3632 0, /* tp_repr */
3633 0, /* tp_as_number */
3634 0, /* tp_as_sequence */
3635 0, /* tp_as_mapping */
3636 0, /* tp_hash */
3637 0, /* tp_call */
3638 0, /* tp_str */
3639 PyObject_GenericGetAttr, /* tp_getattro */
3640 0, /* tp_setattro */
3641 0, /* tp_as_buffer */
3642 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3643 Py_TPFLAGS_BASETYPE, /* tp_flags */
3644 izip_doc, /* tp_doc */
3645 (traverseproc)izip_traverse, /* tp_traverse */
3646 0, /* tp_clear */
3647 0, /* tp_richcompare */
3648 0, /* tp_weaklistoffset */
3649 PyObject_SelfIter, /* tp_iter */
3650 (iternextfunc)izip_next, /* tp_iternext */
3651 0, /* tp_methods */
3652 0, /* tp_members */
3653 0, /* tp_getset */
3654 0, /* tp_base */
3655 0, /* tp_dict */
3656 0, /* tp_descr_get */
3657 0, /* tp_descr_set */
3658 0, /* tp_dictoffset */
3659 0, /* tp_init */
3660 0, /* tp_alloc */
3661 izip_new, /* tp_new */
3662 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003663};
3664
3665
3666/* repeat object ************************************************************/
3667
3668typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003669 PyObject_HEAD
3670 PyObject *element;
3671 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003672} repeatobject;
3673
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003674static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003675
3676static PyObject *
3677repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3678{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003679 repeatobject *ro;
3680 PyObject *element;
Raymond Hettinger58ad2452014-06-24 21:53:45 -07003681 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003682 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003683
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003684 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3685 &element, &cnt))
3686 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003687
Raymond Hettinger58ad2452014-06-24 21:53:45 -07003688 if (kwds != NULL)
3689 n_kwds = PyDict_Size(kwds);
3690 /* Does user supply times argument? */
3691 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003692 cnt = 0;
3693
3694 ro = (repeatobject *)type->tp_alloc(type, 0);
3695 if (ro == NULL)
3696 return NULL;
3697 Py_INCREF(element);
3698 ro->element = element;
3699 ro->cnt = cnt;
3700 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003701}
3702
3703static void
3704repeat_dealloc(repeatobject *ro)
3705{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003706 PyObject_GC_UnTrack(ro);
3707 Py_XDECREF(ro->element);
3708 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003709}
3710
3711static int
3712repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3713{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003714 Py_VISIT(ro->element);
3715 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003716}
3717
3718static PyObject *
3719repeat_next(repeatobject *ro)
3720{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003721 if (ro->cnt == 0)
3722 return NULL;
3723 if (ro->cnt > 0)
3724 ro->cnt--;
3725 Py_INCREF(ro->element);
3726 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003727}
3728
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003729static PyObject *
3730repeat_repr(repeatobject *ro)
3731{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003732 PyObject *result, *objrepr;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003733
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003734 objrepr = PyObject_Repr(ro->element);
3735 if (objrepr == NULL)
3736 return NULL;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003737
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003738 if (ro->cnt == -1)
3739 result = PyString_FromFormat("repeat(%s)",
3740 PyString_AS_STRING(objrepr));
3741 else
3742 result = PyString_FromFormat("repeat(%s, %zd)",
3743 PyString_AS_STRING(objrepr), ro->cnt);
3744 Py_DECREF(objrepr);
3745 return result;
3746}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003747
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003748static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003749repeat_len(repeatobject *ro)
3750{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003751 if (ro->cnt == -1) {
3752 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3753 return NULL;
3754 }
3755 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003756}
3757
Armin Rigof5b3e362006-02-11 21:32:43 +00003758PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003759
3760static PyMethodDef repeat_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003761 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3762 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003763};
3764
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003765PyDoc_STRVAR(repeat_doc,
Raymond Hettinger182edae2009-02-19 02:38:25 +00003766"repeat(object [,times]) -> create an iterator which returns the object\n\
3767for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003768endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003769
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003770static PyTypeObject repeat_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003771 PyVarObject_HEAD_INIT(NULL, 0)
3772 "itertools.repeat", /* tp_name */
3773 sizeof(repeatobject), /* tp_basicsize */
3774 0, /* tp_itemsize */
3775 /* methods */
3776 (destructor)repeat_dealloc, /* tp_dealloc */
3777 0, /* tp_print */
3778 0, /* tp_getattr */
3779 0, /* tp_setattr */
3780 0, /* tp_compare */
3781 (reprfunc)repeat_repr, /* tp_repr */
3782 0, /* tp_as_number */
3783 0, /* tp_as_sequence */
3784 0, /* tp_as_mapping */
3785 0, /* tp_hash */
3786 0, /* tp_call */
3787 0, /* tp_str */
3788 PyObject_GenericGetAttr, /* tp_getattro */
3789 0, /* tp_setattro */
3790 0, /* tp_as_buffer */
3791 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3792 Py_TPFLAGS_BASETYPE, /* tp_flags */
3793 repeat_doc, /* tp_doc */
3794 (traverseproc)repeat_traverse, /* tp_traverse */
3795 0, /* tp_clear */
3796 0, /* tp_richcompare */
3797 0, /* tp_weaklistoffset */
3798 PyObject_SelfIter, /* tp_iter */
3799 (iternextfunc)repeat_next, /* tp_iternext */
3800 repeat_methods, /* tp_methods */
3801 0, /* tp_members */
3802 0, /* tp_getset */
3803 0, /* tp_base */
3804 0, /* tp_dict */
3805 0, /* tp_descr_get */
3806 0, /* tp_descr_set */
3807 0, /* tp_dictoffset */
3808 0, /* tp_init */
3809 0, /* tp_alloc */
3810 repeat_new, /* tp_new */
3811 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003812};
3813
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003814/* iziplongest object ************************************************************/
3815
3816#include "Python.h"
3817
3818typedef struct {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003819 PyObject_HEAD
3820 Py_ssize_t tuplesize;
3821 Py_ssize_t numactive;
3822 PyObject *ittuple; /* tuple of iterators */
3823 PyObject *result;
3824 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003825} iziplongestobject;
3826
3827static PyTypeObject iziplongest_type;
3828
3829static PyObject *
3830izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3831{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003832 iziplongestobject *lz;
3833 Py_ssize_t i;
3834 PyObject *ittuple; /* tuple of iterators */
3835 PyObject *result;
3836 PyObject *fillvalue = Py_None;
3837 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003838
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003839 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3840 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3841 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3842 PyErr_SetString(PyExc_TypeError,
3843 "izip_longest() got an unexpected keyword argument");
3844 return NULL;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003845 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003846 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003847
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003848 /* args must be a tuple */
3849 assert(PyTuple_Check(args));
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003850
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003851 /* obtain iterators */
3852 ittuple = PyTuple_New(tuplesize);
3853 if (ittuple == NULL)
3854 return NULL;
3855 for (i=0; i < tuplesize; ++i) {
3856 PyObject *item = PyTuple_GET_ITEM(args, i);
3857 PyObject *it = PyObject_GetIter(item);
3858 if (it == NULL) {
3859 if (PyErr_ExceptionMatches(PyExc_TypeError))
3860 PyErr_Format(PyExc_TypeError,
3861 "izip_longest argument #%zd must support iteration",
3862 i+1);
3863 Py_DECREF(ittuple);
3864 return NULL;
3865 }
3866 PyTuple_SET_ITEM(ittuple, i, it);
3867 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003868
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003869 /* create a result holder */
3870 result = PyTuple_New(tuplesize);
3871 if (result == NULL) {
3872 Py_DECREF(ittuple);
3873 return NULL;
3874 }
3875 for (i=0 ; i < tuplesize ; i++) {
3876 Py_INCREF(Py_None);
3877 PyTuple_SET_ITEM(result, i, Py_None);
3878 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003879
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003880 /* create iziplongestobject structure */
3881 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3882 if (lz == NULL) {
3883 Py_DECREF(ittuple);
3884 Py_DECREF(result);
3885 return NULL;
3886 }
3887 lz->ittuple = ittuple;
3888 lz->tuplesize = tuplesize;
3889 lz->numactive = tuplesize;
3890 lz->result = result;
3891 Py_INCREF(fillvalue);
3892 lz->fillvalue = fillvalue;
3893 return (PyObject *)lz;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003894}
3895
3896static void
3897izip_longest_dealloc(iziplongestobject *lz)
3898{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003899 PyObject_GC_UnTrack(lz);
3900 Py_XDECREF(lz->ittuple);
3901 Py_XDECREF(lz->result);
3902 Py_XDECREF(lz->fillvalue);
3903 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003904}
3905
3906static int
3907izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3908{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003909 Py_VISIT(lz->ittuple);
3910 Py_VISIT(lz->result);
3911 Py_VISIT(lz->fillvalue);
3912 return 0;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003913}
3914
3915static PyObject *
3916izip_longest_next(iziplongestobject *lz)
3917{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003918 Py_ssize_t i;
3919 Py_ssize_t tuplesize = lz->tuplesize;
3920 PyObject *result = lz->result;
3921 PyObject *it;
3922 PyObject *item;
3923 PyObject *olditem;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003924
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003925 if (tuplesize == 0)
3926 return NULL;
3927 if (lz->numactive == 0)
3928 return NULL;
3929 if (Py_REFCNT(result) == 1) {
3930 Py_INCREF(result);
3931 for (i=0 ; i < tuplesize ; i++) {
3932 it = PyTuple_GET_ITEM(lz->ittuple, i);
3933 if (it == NULL) {
3934 Py_INCREF(lz->fillvalue);
3935 item = lz->fillvalue;
3936 } else {
3937 item = PyIter_Next(it);
3938 if (item == NULL) {
3939 lz->numactive -= 1;
3940 if (lz->numactive == 0 || PyErr_Occurred()) {
3941 lz->numactive = 0;
3942 Py_DECREF(result);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003943 return NULL;
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003944 } else {
3945 Py_INCREF(lz->fillvalue);
3946 item = lz->fillvalue;
3947 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3948 Py_DECREF(it);
3949 }
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003950 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003951 }
3952 olditem = PyTuple_GET_ITEM(result, i);
3953 PyTuple_SET_ITEM(result, i, item);
3954 Py_DECREF(olditem);
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +00003955 }
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003956 } else {
3957 result = PyTuple_New(tuplesize);
3958 if (result == NULL)
3959 return NULL;
3960 for (i=0 ; i < tuplesize ; i++) {
3961 it = PyTuple_GET_ITEM(lz->ittuple, i);
3962 if (it == NULL) {
3963 Py_INCREF(lz->fillvalue);
3964 item = lz->fillvalue;
3965 } else {
3966 item = PyIter_Next(it);
3967 if (item == NULL) {
3968 lz->numactive -= 1;
3969 if (lz->numactive == 0 || PyErr_Occurred()) {
3970 lz->numactive = 0;
3971 Py_DECREF(result);
3972 return NULL;
3973 } else {
3974 Py_INCREF(lz->fillvalue);
3975 item = lz->fillvalue;
3976 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3977 Py_DECREF(it);
3978 }
3979 }
3980 }
3981 PyTuple_SET_ITEM(result, i, item);
3982 }
3983 }
3984 return result;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003985}
3986
3987PyDoc_STRVAR(izip_longest_doc,
3988"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3989\n\
3990Return an izip_longest object whose .next() method returns a tuple where\n\
3991the i-th element comes from the i-th iterable argument. The .next()\n\
3992method continues until the longest iterable in the argument sequence\n\
3993is exhausted and then it raises StopIteration. When the shorter iterables\n\
3994are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3995defaults to None or can be specified by a keyword argument.\n\
3996");
3997
3998static PyTypeObject iziplongest_type = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00003999 PyVarObject_HEAD_INIT(NULL, 0)
4000 "itertools.izip_longest", /* tp_name */
4001 sizeof(iziplongestobject), /* tp_basicsize */
4002 0, /* tp_itemsize */
4003 /* methods */
4004 (destructor)izip_longest_dealloc, /* tp_dealloc */
4005 0, /* tp_print */
4006 0, /* tp_getattr */
4007 0, /* tp_setattr */
4008 0, /* tp_compare */
4009 0, /* tp_repr */
4010 0, /* tp_as_number */
4011 0, /* tp_as_sequence */
4012 0, /* tp_as_mapping */
4013 0, /* tp_hash */
4014 0, /* tp_call */
4015 0, /* tp_str */
4016 PyObject_GenericGetAttr, /* tp_getattro */
4017 0, /* tp_setattro */
4018 0, /* tp_as_buffer */
4019 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4020 Py_TPFLAGS_BASETYPE, /* tp_flags */
4021 izip_longest_doc, /* tp_doc */
4022 (traverseproc)izip_longest_traverse, /* tp_traverse */
4023 0, /* tp_clear */
4024 0, /* tp_richcompare */
4025 0, /* tp_weaklistoffset */
4026 PyObject_SelfIter, /* tp_iter */
4027 (iternextfunc)izip_longest_next, /* tp_iternext */
4028 0, /* tp_methods */
4029 0, /* tp_members */
4030 0, /* tp_getset */
4031 0, /* tp_base */
4032 0, /* tp_dict */
4033 0, /* tp_descr_get */
4034 0, /* tp_descr_set */
4035 0, /* tp_dictoffset */
4036 0, /* tp_init */
4037 0, /* tp_alloc */
4038 izip_longest_new, /* tp_new */
4039 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd36862c2007-02-21 05:20:38 +00004040};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004041
4042/* module level code ********************************************************/
4043
4044PyDoc_STRVAR(module_doc,
4045"Functional tools for creating and using iterators.\n\
4046\n\
4047Infinite iterators:\n\
4048count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004049cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004050repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004051\n\
4052Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004053chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettingere76816b2009-01-29 03:43:44 +00004054compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004055dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4056groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00004057ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
4058ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004059islice(seq, [start,] stop [, step]) --> elements from\n\
4060 seq[start:stop:step]\n\
4061imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
4062starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004063tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004064takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger36d816b2009-01-29 03:21:42 +00004065izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4066izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4067\n\
4068Combinatoric generators:\n\
4069product(p, q, ... [repeat=1]) --> cartesian product\n\
4070permutations(p[, r])\n\
Raymond Hettinger9eac1192009-11-19 01:22:04 +00004071combinations(p, r)\n\
4072combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004073");
4074
4075
Raymond Hettingerad983e72003-11-12 14:32:26 +00004076static PyMethodDef module_methods[] = {
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004077 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4078 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004079};
4080
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004081PyMODINIT_FUNC
4082inititertools(void)
4083{
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004084 int i;
4085 PyObject *m;
4086 char *name;
4087 PyTypeObject *typelist[] = {
4088 &combinations_type,
4089 &cwr_type,
4090 &cycle_type,
4091 &dropwhile_type,
4092 &takewhile_type,
4093 &islice_type,
4094 &starmap_type,
4095 &imap_type,
4096 &chain_type,
4097 &compress_type,
4098 &ifilter_type,
4099 &ifilterfalse_type,
4100 &count_type,
4101 &izip_type,
4102 &iziplongest_type,
4103 &permutations_type,
4104 &product_type,
4105 &repeat_type,
4106 &groupby_type,
4107 NULL
4108 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004109
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004110 Py_TYPE(&teedataobject_type) = &PyType_Type;
4111 m = Py_InitModule3("itertools", module_methods, module_doc);
4112 if (m == NULL)
4113 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004114
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004115 for (i=0 ; typelist[i] != NULL ; i++) {
4116 if (PyType_Ready(typelist[i]) < 0)
4117 return;
4118 name = strchr(typelist[i]->tp_name, '.');
4119 assert (name != NULL);
4120 Py_INCREF(typelist[i]);
4121 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4122 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004123
Antoine Pitrouc83ea132010-05-09 14:46:46 +00004124 if (PyType_Ready(&teedataobject_type) < 0)
4125 return;
4126 if (PyType_Ready(&tee_type) < 0)
4127 return;
4128 if (PyType_Ready(&_grouper_type) < 0)
4129 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004130}