blob: bd10ee416adb2cbb2f350c1eac870fd5a24d2be6 [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 Pitrouf95a1b32010-05-09 15:52:27 +00005/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00006 by Raymond D. Hettinger <python@rcn.com>
7 Copyright (c) 2003 Python Software Foundation.
8 All rights reserved.
9*/
10
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000011
12/* groupby object ***********************************************************/
13
14typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000015 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000021} groupbyobject;
22
23static PyTypeObject groupby_type;
24static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26static PyObject *
27groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000029 static char *kwargs[] = {"iterable", "key", NULL};
30 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000032
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000033 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
34 &it, &keyfunc))
35 return NULL;
36
37 gbo = (groupbyobject *)type->tp_alloc(type, 0);
38 if (gbo == NULL)
39 return NULL;
40 gbo->tgtkey = NULL;
41 gbo->currkey = NULL;
42 gbo->currvalue = NULL;
43 gbo->keyfunc = keyfunc;
44 Py_INCREF(keyfunc);
45 gbo->it = PyObject_GetIter(it);
46 if (gbo->it == NULL) {
47 Py_DECREF(gbo);
48 return NULL;
49 }
50 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000051}
52
53static void
54groupby_dealloc(groupbyobject *gbo)
55{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000056 PyObject_GC_UnTrack(gbo);
57 Py_XDECREF(gbo->it);
58 Py_XDECREF(gbo->keyfunc);
59 Py_XDECREF(gbo->tgtkey);
60 Py_XDECREF(gbo->currkey);
61 Py_XDECREF(gbo->currvalue);
62 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000063}
64
65static int
66groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
67{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000068 Py_VISIT(gbo->it);
69 Py_VISIT(gbo->keyfunc);
70 Py_VISIT(gbo->tgtkey);
71 Py_VISIT(gbo->currkey);
72 Py_VISIT(gbo->currvalue);
73 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000074}
75
76static PyObject *
77groupby_next(groupbyobject *gbo)
78{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000081 /* skip to next iteration group */
82 for (;;) {
83 if (gbo->currkey == NULL)
84 /* pass */;
85 else if (gbo->tgtkey == NULL)
86 break;
87 else {
88 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000089
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000090 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
91 gbo->currkey, Py_EQ);
92 if (rcmp == -1)
93 return NULL;
94 else if (rcmp == 0)
95 break;
96 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000097
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000098 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000101
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000102 if (gbo->keyfunc == Py_None) {
103 newkey = newvalue;
104 Py_INCREF(newvalue);
105 } else {
106 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
107 newvalue, NULL);
108 if (newkey == NULL) {
109 Py_DECREF(newvalue);
110 return NULL;
111 }
112 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000113
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000114 tmp = gbo->currkey;
115 gbo->currkey = newkey;
116 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000117
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000118 tmp = gbo->currvalue;
119 gbo->currvalue = newvalue;
120 Py_XDECREF(tmp);
121 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000123 Py_INCREF(gbo->currkey);
124 tmp = gbo->tgtkey;
125 gbo->tgtkey = gbo->currkey;
126 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000127
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000131
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000135}
136
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000137static PyObject *
138groupby_reduce(groupbyobject *lz)
139{
140 /* reduce as a 'new' call with an optional 'setstate' if groupby
141 * has started
142 */
143 PyObject *value;
144 if (lz->tgtkey && lz->currkey && lz->currvalue)
145 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
146 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
147 else
148 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
149 lz->it, lz->keyfunc);
150
151 return value;
152}
153
154PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
155
156static PyObject *
157groupby_setstate(groupbyobject *lz, PyObject *state)
158{
159 PyObject *currkey, *currvalue, *tgtkey;
160 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey))
161 return NULL;
162 Py_CLEAR(lz->currkey);
163 lz->currkey = currkey;
164 Py_INCREF(lz->currkey);
165 Py_CLEAR(lz->currvalue);
166 lz->currvalue = currvalue;
167 Py_INCREF(lz->currvalue);
168 Py_CLEAR(lz->tgtkey);
169 lz->tgtkey = tgtkey;
170 Py_INCREF(lz->tgtkey);
171 Py_RETURN_NONE;
172}
173
174PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
175
176static PyMethodDef groupby_methods[] = {
177 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
178 reduce_doc},
179 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
180 setstate_doc},
181 {NULL, NULL} /* sentinel */
182};
183
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000184PyDoc_STRVAR(groupby_doc,
185"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
186(key, sub-iterator) grouped by each value of key(value).\n");
187
188static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000189 PyVarObject_HEAD_INIT(NULL, 0)
190 "itertools.groupby", /* tp_name */
191 sizeof(groupbyobject), /* tp_basicsize */
192 0, /* tp_itemsize */
193 /* methods */
194 (destructor)groupby_dealloc, /* tp_dealloc */
195 0, /* tp_print */
196 0, /* tp_getattr */
197 0, /* tp_setattr */
198 0, /* tp_reserved */
199 0, /* tp_repr */
200 0, /* tp_as_number */
201 0, /* tp_as_sequence */
202 0, /* tp_as_mapping */
203 0, /* tp_hash */
204 0, /* tp_call */
205 0, /* tp_str */
206 PyObject_GenericGetAttr, /* tp_getattro */
207 0, /* tp_setattro */
208 0, /* tp_as_buffer */
209 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
210 Py_TPFLAGS_BASETYPE, /* tp_flags */
211 groupby_doc, /* tp_doc */
212 (traverseproc)groupby_traverse, /* tp_traverse */
213 0, /* tp_clear */
214 0, /* tp_richcompare */
215 0, /* tp_weaklistoffset */
216 PyObject_SelfIter, /* tp_iter */
217 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000218 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000219 0, /* tp_members */
220 0, /* tp_getset */
221 0, /* tp_base */
222 0, /* tp_dict */
223 0, /* tp_descr_get */
224 0, /* tp_descr_set */
225 0, /* tp_dictoffset */
226 0, /* tp_init */
227 0, /* tp_alloc */
228 groupby_new, /* tp_new */
229 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000230};
231
232
233/* _grouper object (internal) ************************************************/
234
235typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000236 PyObject_HEAD
237 PyObject *parent;
238 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000239} _grouperobject;
240
241static PyTypeObject _grouper_type;
242
243static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000244_grouper_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
245{
246 PyObject *parent, *tgtkey;
247
248 if (!PyArg_ParseTuple(args, "O!O", &groupby_type, &parent, &tgtkey))
249 return NULL;
250
251 return _grouper_create((groupbyobject*) parent, tgtkey);
252}
253
254static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000255_grouper_create(groupbyobject *parent, PyObject *tgtkey)
256{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000257 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000259 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
260 if (igo == NULL)
261 return NULL;
262 igo->parent = (PyObject *)parent;
263 Py_INCREF(parent);
264 igo->tgtkey = tgtkey;
265 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000266
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000267 PyObject_GC_Track(igo);
268 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000269}
270
271static void
272_grouper_dealloc(_grouperobject *igo)
273{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000274 PyObject_GC_UnTrack(igo);
275 Py_DECREF(igo->parent);
276 Py_DECREF(igo->tgtkey);
277 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000278}
279
280static int
281_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
282{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000283 Py_VISIT(igo->parent);
284 Py_VISIT(igo->tgtkey);
285 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000286}
287
288static PyObject *
289_grouper_next(_grouperobject *igo)
290{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000291 groupbyobject *gbo = (groupbyobject *)igo->parent;
292 PyObject *newvalue, *newkey, *r;
293 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000294
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000295 if (gbo->currvalue == NULL) {
296 newvalue = PyIter_Next(gbo->it);
297 if (newvalue == NULL)
298 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000299
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000300 if (gbo->keyfunc == Py_None) {
301 newkey = newvalue;
302 Py_INCREF(newvalue);
303 } else {
304 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
305 newvalue, NULL);
306 if (newkey == NULL) {
307 Py_DECREF(newvalue);
308 return NULL;
309 }
310 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000311
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000312 assert(gbo->currkey == NULL);
313 gbo->currkey = newkey;
314 gbo->currvalue = newvalue;
315 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000316
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000317 assert(gbo->currkey != NULL);
318 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
319 if (rcmp <= 0)
320 /* got any error or current group is end */
321 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000322
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000323 r = gbo->currvalue;
324 gbo->currvalue = NULL;
325 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000326
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000327 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000328}
329
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000330static PyObject *
331_grouper_reduce(_grouperobject *lz)
332{
333 return Py_BuildValue("O(OO)", Py_TYPE(lz),
334 lz->parent, lz->tgtkey);
335}
336
337static PyMethodDef _grouper_methods[] = {
338 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
339 reduce_doc},
340 {NULL, NULL} /* sentinel */
341};
342
343
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000344static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000345 PyVarObject_HEAD_INIT(NULL, 0)
346 "itertools._grouper", /* tp_name */
347 sizeof(_grouperobject), /* tp_basicsize */
348 0, /* tp_itemsize */
349 /* methods */
350 (destructor)_grouper_dealloc, /* tp_dealloc */
351 0, /* tp_print */
352 0, /* tp_getattr */
353 0, /* tp_setattr */
354 0, /* tp_reserved */
355 0, /* tp_repr */
356 0, /* tp_as_number */
357 0, /* tp_as_sequence */
358 0, /* tp_as_mapping */
359 0, /* tp_hash */
360 0, /* tp_call */
361 0, /* tp_str */
362 PyObject_GenericGetAttr, /* tp_getattro */
363 0, /* tp_setattro */
364 0, /* tp_as_buffer */
365 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
366 0, /* tp_doc */
367 (traverseproc)_grouper_traverse,/* tp_traverse */
368 0, /* tp_clear */
369 0, /* tp_richcompare */
370 0, /* tp_weaklistoffset */
371 PyObject_SelfIter, /* tp_iter */
372 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000373 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000374 0, /* tp_members */
375 0, /* tp_getset */
376 0, /* tp_base */
377 0, /* tp_dict */
378 0, /* tp_descr_get */
379 0, /* tp_descr_set */
380 0, /* tp_dictoffset */
381 0, /* tp_init */
382 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000383 _grouper_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000384 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000385};
386
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000387
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000388
Raymond Hettingerad983e72003-11-12 14:32:26 +0000389/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000390
Raymond Hettingerad983e72003-11-12 14:32:26 +0000391/* The teedataobject pre-allocates space for LINKCELLS number of objects.
392 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000393 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000394 the other structure members including PyHEAD overhead. The larger the
395 value, the less memory overhead per object and the less time spent
396 allocating/deallocating new links. The smaller the number, the less
397 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000398*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000399#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000400
401typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000402 PyObject_HEAD
403 PyObject *it;
404 int numread;
405 PyObject *nextlink;
406 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000407} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000408
409typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000410 PyObject_HEAD
411 teedataobject *dataobj;
412 int index;
413 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000414} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415
Raymond Hettingerad983e72003-11-12 14:32:26 +0000416static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417
418static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000419teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000420{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000421 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000422
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000423 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
424 if (tdo == NULL)
425 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000426
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000427 tdo->numread = 0;
428 tdo->nextlink = NULL;
429 Py_INCREF(it);
430 tdo->it = it;
431 PyObject_GC_Track(tdo);
432 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000433}
434
435static PyObject *
436teedataobject_jumplink(teedataobject *tdo)
437{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000438 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000439 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000440 Py_XINCREF(tdo->nextlink);
441 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000442}
443
444static PyObject *
445teedataobject_getitem(teedataobject *tdo, int i)
446{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000447 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000448
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000449 assert(i < LINKCELLS);
450 if (i < tdo->numread)
451 value = tdo->values[i];
452 else {
453 /* this is the lead iterator, so fetch more data */
454 assert(i == tdo->numread);
455 value = PyIter_Next(tdo->it);
456 if (value == NULL)
457 return NULL;
458 tdo->numread++;
459 tdo->values[i] = value;
460 }
461 Py_INCREF(value);
462 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000463}
464
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000465static int
466teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
467{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000468 int i;
469 Py_VISIT(tdo->it);
470 for (i = 0; i < tdo->numread; i++)
471 Py_VISIT(tdo->values[i]);
472 Py_VISIT(tdo->nextlink);
473 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000474}
475
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200476static void
477teedataobject_safe_decref(PyObject *obj)
478{
479 while (obj && Py_TYPE(obj) == &teedataobject_type &&
480 Py_REFCNT(obj) == 1) {
481 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
482 ((teedataobject *)obj)->nextlink = NULL;
483 Py_DECREF(obj);
484 obj = nextlink;
485 }
486 Py_XDECREF(obj);
487}
488
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000489static int
490teedataobject_clear(teedataobject *tdo)
491{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000492 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200493 PyObject *tmp;
494
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000495 Py_CLEAR(tdo->it);
496 for (i=0 ; i<tdo->numread ; i++)
497 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200498 tmp = tdo->nextlink;
499 tdo->nextlink = NULL;
500 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000501 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000502}
503
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000504static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000505teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000506{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000507 PyObject_GC_UnTrack(tdo);
508 teedataobject_clear(tdo);
509 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000510}
511
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000512static PyObject *
513teedataobject_reduce(teedataobject *tdo)
514{
515 int i;
516 /* create a temporary list of already iterated values */
517 PyObject *values = PyList_New(tdo->numread);
518 if (!values)
519 return NULL;
520 for (i=0 ; i<tdo->numread ; i++) {
521 Py_INCREF(tdo->values[i]);
522 PyList_SET_ITEM(values, i, tdo->values[i]);
523 }
524 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
525 values,
526 tdo->nextlink ? tdo->nextlink : Py_None);
527}
528
529static PyTypeObject teedataobject_type;
530
531static PyObject *
532teedataobject_new(PyTypeObject *type, PyObject *args, PyObject *kw)
533{
534 teedataobject *tdo;
535 PyObject *it, *values, *next;
536 Py_ssize_t i, len;
537
538 assert(type == &teedataobject_type);
539 if (!PyArg_ParseTuple(args, "OO!O", &it, &PyList_Type, &values, &next))
540 return NULL;
541
542 tdo = (teedataobject *)teedataobject_newinternal(it);
543 if (!tdo)
544 return NULL;
545
546 len = PyList_GET_SIZE(values);
547 if (len > LINKCELLS)
548 goto err;
549 for (i=0; i<len; i++) {
550 tdo->values[i] = PyList_GET_ITEM(values, i);
551 Py_INCREF(tdo->values[i]);
552 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200553 /* len <= LINKCELLS < INT_MAX */
554 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000555
556 if (len == LINKCELLS) {
557 if (next != Py_None) {
558 if (Py_TYPE(next) != &teedataobject_type)
559 goto err;
560 assert(tdo->nextlink == NULL);
561 Py_INCREF(next);
562 tdo->nextlink = next;
563 }
564 } else {
565 if (next != Py_None)
566 goto err; /* shouldn't have a next if we are not full */
567 }
568 return (PyObject*)tdo;
569
570err:
571 Py_XDECREF(tdo);
572 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
573 return NULL;
574}
575
576static PyMethodDef teedataobject_methods[] = {
577 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
578 reduce_doc},
579 {NULL, NULL} /* sentinel */
580};
581
Raymond Hettingerad983e72003-11-12 14:32:26 +0000582PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000583
Raymond Hettingerad983e72003-11-12 14:32:26 +0000584static PyTypeObject teedataobject_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000585 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000586 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000587 sizeof(teedataobject), /* tp_basicsize */
588 0, /* tp_itemsize */
589 /* methods */
590 (destructor)teedataobject_dealloc, /* tp_dealloc */
591 0, /* tp_print */
592 0, /* tp_getattr */
593 0, /* tp_setattr */
594 0, /* tp_reserved */
595 0, /* tp_repr */
596 0, /* tp_as_number */
597 0, /* tp_as_sequence */
598 0, /* tp_as_mapping */
599 0, /* tp_hash */
600 0, /* tp_call */
601 0, /* tp_str */
602 PyObject_GenericGetAttr, /* tp_getattro */
603 0, /* tp_setattro */
604 0, /* tp_as_buffer */
605 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
606 teedataobject_doc, /* tp_doc */
607 (traverseproc)teedataobject_traverse, /* tp_traverse */
608 (inquiry)teedataobject_clear, /* tp_clear */
609 0, /* tp_richcompare */
610 0, /* tp_weaklistoffset */
611 0, /* tp_iter */
612 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000613 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000614 0, /* tp_members */
615 0, /* tp_getset */
616 0, /* tp_base */
617 0, /* tp_dict */
618 0, /* tp_descr_get */
619 0, /* tp_descr_set */
620 0, /* tp_dictoffset */
621 0, /* tp_init */
622 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000623 teedataobject_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000624 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000625};
626
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000627
628static PyTypeObject tee_type;
629
630static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000631tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000632{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000633 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000634
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000635 if (to->index >= LINKCELLS) {
636 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200637 if (link == NULL)
638 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000639 Py_DECREF(to->dataobj);
640 to->dataobj = (teedataobject *)link;
641 to->index = 0;
642 }
643 value = teedataobject_getitem(to->dataobj, to->index);
644 if (value == NULL)
645 return NULL;
646 to->index++;
647 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000648}
649
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000650static int
651tee_traverse(teeobject *to, visitproc visit, void *arg)
652{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000653 Py_VISIT((PyObject *)to->dataobj);
654 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000655}
656
Raymond Hettingerad983e72003-11-12 14:32:26 +0000657static PyObject *
658tee_copy(teeobject *to)
659{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000660 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000661
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000662 newto = PyObject_GC_New(teeobject, &tee_type);
663 if (newto == NULL)
664 return NULL;
665 Py_INCREF(to->dataobj);
666 newto->dataobj = to->dataobj;
667 newto->index = to->index;
668 newto->weakreflist = NULL;
669 PyObject_GC_Track(newto);
670 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000671}
672
673PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
674
675static PyObject *
676tee_fromiterable(PyObject *iterable)
677{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000678 teeobject *to;
679 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000680
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000681 it = PyObject_GetIter(iterable);
682 if (it == NULL)
683 return NULL;
684 if (PyObject_TypeCheck(it, &tee_type)) {
685 to = (teeobject *)tee_copy((teeobject *)it);
686 goto done;
687 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000688
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000689 to = PyObject_GC_New(teeobject, &tee_type);
690 if (to == NULL)
691 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000692 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000693 if (!to->dataobj) {
694 PyObject_GC_Del(to);
695 to = NULL;
696 goto done;
697 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000698
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000699 to->index = 0;
700 to->weakreflist = NULL;
701 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000702done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000703 Py_XDECREF(it);
704 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000705}
706
707static PyObject *
708tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
709{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000710 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000711
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000712 if (!PyArg_UnpackTuple(args, "_tee", 1, 1, &iterable))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000713 return NULL;
714 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000715}
716
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000717static int
718tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000719{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000720 if (to->weakreflist != NULL)
721 PyObject_ClearWeakRefs((PyObject *) to);
722 Py_CLEAR(to->dataobj);
723 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000724}
725
726static void
727tee_dealloc(teeobject *to)
728{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000729 PyObject_GC_UnTrack(to);
730 tee_clear(to);
731 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000732}
733
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000734static PyObject *
735tee_reduce(teeobject *to)
736{
737 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
738}
739
740static PyObject *
741tee_setstate(teeobject *to, PyObject *state)
742{
743 teedataobject *tdo;
744 int index;
745 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index))
746 return NULL;
747 if (index < 0 || index > LINKCELLS) {
748 PyErr_SetString(PyExc_ValueError, "Index out of range");
749 return NULL;
750 }
751 Py_CLEAR(to->dataobj);
752 to->dataobj = tdo;
753 Py_INCREF(to->dataobj);
754 to->index = index;
755 Py_RETURN_NONE;
756}
757
Raymond Hettingerad983e72003-11-12 14:32:26 +0000758PyDoc_STRVAR(teeobject_doc,
759"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000760
Raymond Hettingerad983e72003-11-12 14:32:26 +0000761static PyMethodDef tee_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000762 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000763 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
764 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000765 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000766};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000767
768static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000769 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000770 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000771 sizeof(teeobject), /* tp_basicsize */
772 0, /* tp_itemsize */
773 /* methods */
774 (destructor)tee_dealloc, /* tp_dealloc */
775 0, /* tp_print */
776 0, /* tp_getattr */
777 0, /* tp_setattr */
778 0, /* tp_reserved */
779 0, /* tp_repr */
780 0, /* tp_as_number */
781 0, /* tp_as_sequence */
782 0, /* tp_as_mapping */
783 0, /* tp_hash */
784 0, /* tp_call */
785 0, /* tp_str */
786 0, /* tp_getattro */
787 0, /* tp_setattro */
788 0, /* tp_as_buffer */
789 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
790 teeobject_doc, /* tp_doc */
791 (traverseproc)tee_traverse, /* tp_traverse */
792 (inquiry)tee_clear, /* tp_clear */
793 0, /* tp_richcompare */
794 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
795 PyObject_SelfIter, /* tp_iter */
796 (iternextfunc)tee_next, /* tp_iternext */
797 tee_methods, /* tp_methods */
798 0, /* tp_members */
799 0, /* tp_getset */
800 0, /* tp_base */
801 0, /* tp_dict */
802 0, /* tp_descr_get */
803 0, /* tp_descr_set */
804 0, /* tp_dictoffset */
805 0, /* tp_init */
806 0, /* tp_alloc */
807 tee_new, /* tp_new */
808 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000809};
810
Raymond Hettingerad983e72003-11-12 14:32:26 +0000811static PyObject *
812tee(PyObject *self, PyObject *args)
813{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000814 Py_ssize_t i, n=2;
815 PyObject *it, *iterable, *copyable, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200816 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000817
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000818 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
819 return NULL;
820 if (n < 0) {
821 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
822 return NULL;
823 }
824 result = PyTuple_New(n);
825 if (result == NULL)
826 return NULL;
827 if (n == 0)
828 return result;
829 it = PyObject_GetIter(iterable);
830 if (it == NULL) {
831 Py_DECREF(result);
832 return NULL;
833 }
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200834 if (!_PyObject_HasAttrId(it, &PyId___copy__)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000835 copyable = tee_fromiterable(it);
836 Py_DECREF(it);
837 if (copyable == NULL) {
838 Py_DECREF(result);
839 return NULL;
840 }
841 } else
842 copyable = it;
843 PyTuple_SET_ITEM(result, 0, copyable);
844 for (i=1 ; i<n ; i++) {
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200845
846 copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000847 if (copyable == NULL) {
848 Py_DECREF(result);
849 return NULL;
850 }
851 PyTuple_SET_ITEM(result, i, copyable);
852 }
853 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000854}
855
856PyDoc_STRVAR(tee_doc,
857"tee(iterable, n=2) --> tuple of n independent iterators.");
858
859
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000860/* cycle object **********************************************************/
861
862typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000863 PyObject_HEAD
864 PyObject *it;
865 PyObject *saved;
866 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000867} cycleobject;
868
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000869static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000870
871static PyObject *
872cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
873{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000874 PyObject *it;
875 PyObject *iterable;
876 PyObject *saved;
877 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000878
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000879 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
880 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000881
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000882 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
883 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000884
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000885 /* Get iterator. */
886 it = PyObject_GetIter(iterable);
887 if (it == NULL)
888 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000889
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000890 saved = PyList_New(0);
891 if (saved == NULL) {
892 Py_DECREF(it);
893 return NULL;
894 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000895
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000896 /* create cycleobject structure */
897 lz = (cycleobject *)type->tp_alloc(type, 0);
898 if (lz == NULL) {
899 Py_DECREF(it);
900 Py_DECREF(saved);
901 return NULL;
902 }
903 lz->it = it;
904 lz->saved = saved;
905 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000906
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000907 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000908}
909
910static void
911cycle_dealloc(cycleobject *lz)
912{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000913 PyObject_GC_UnTrack(lz);
914 Py_XDECREF(lz->saved);
915 Py_XDECREF(lz->it);
916 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000917}
918
919static int
920cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
921{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000922 Py_VISIT(lz->it);
923 Py_VISIT(lz->saved);
924 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000925}
926
927static PyObject *
928cycle_next(cycleobject *lz)
929{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000930 PyObject *item;
931 PyObject *it;
932 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000933
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000934 while (1) {
935 item = PyIter_Next(lz->it);
936 if (item != NULL) {
937 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
938 Py_DECREF(item);
939 return NULL;
940 }
941 return item;
942 }
943 if (PyErr_Occurred()) {
944 if (PyErr_ExceptionMatches(PyExc_StopIteration))
945 PyErr_Clear();
946 else
947 return NULL;
948 }
949 if (PyList_Size(lz->saved) == 0)
950 return NULL;
951 it = PyObject_GetIter(lz->saved);
952 if (it == NULL)
953 return NULL;
954 tmp = lz->it;
955 lz->it = it;
956 lz->firstpass = 1;
957 Py_DECREF(tmp);
958 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000959}
960
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000961static PyObject *
962cycle_reduce(cycleobject *lz)
963{
964 /* Create a new cycle with the iterator tuple, then set
965 * the saved state on it.
966 */
Serhiy Storchakad269b5e2013-01-25 13:38:56 +0200967 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz),
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000968 lz->it, lz->saved, lz->firstpass);
969 }
970
971static PyObject *
972cycle_setstate(cycleobject *lz, PyObject *state)
973{
974 PyObject *saved=NULL;
975 int firstpass;
976 if (!PyArg_ParseTuple(state, "Oi", &saved, &firstpass))
977 return NULL;
978 Py_CLEAR(lz->saved);
979 lz->saved = saved;
980 Py_XINCREF(lz->saved);
981 lz->firstpass = firstpass != 0;
982 Py_RETURN_NONE;
983}
984
985static PyMethodDef cycle_methods[] = {
986 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
987 reduce_doc},
988 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
989 setstate_doc},
990 {NULL, NULL} /* sentinel */
991};
992
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000993PyDoc_STRVAR(cycle_doc,
994"cycle(iterable) --> cycle object\n\
995\n\
996Return elements from the iterable until it is exhausted.\n\
997Then repeat the sequence indefinitely.");
998
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000999static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001000 PyVarObject_HEAD_INIT(NULL, 0)
1001 "itertools.cycle", /* tp_name */
1002 sizeof(cycleobject), /* tp_basicsize */
1003 0, /* tp_itemsize */
1004 /* methods */
1005 (destructor)cycle_dealloc, /* tp_dealloc */
1006 0, /* tp_print */
1007 0, /* tp_getattr */
1008 0, /* tp_setattr */
1009 0, /* tp_reserved */
1010 0, /* tp_repr */
1011 0, /* tp_as_number */
1012 0, /* tp_as_sequence */
1013 0, /* tp_as_mapping */
1014 0, /* tp_hash */
1015 0, /* tp_call */
1016 0, /* tp_str */
1017 PyObject_GenericGetAttr, /* tp_getattro */
1018 0, /* tp_setattro */
1019 0, /* tp_as_buffer */
1020 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1021 Py_TPFLAGS_BASETYPE, /* tp_flags */
1022 cycle_doc, /* tp_doc */
1023 (traverseproc)cycle_traverse, /* tp_traverse */
1024 0, /* tp_clear */
1025 0, /* tp_richcompare */
1026 0, /* tp_weaklistoffset */
1027 PyObject_SelfIter, /* tp_iter */
1028 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001029 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001030 0, /* tp_members */
1031 0, /* tp_getset */
1032 0, /* tp_base */
1033 0, /* tp_dict */
1034 0, /* tp_descr_get */
1035 0, /* tp_descr_set */
1036 0, /* tp_dictoffset */
1037 0, /* tp_init */
1038 0, /* tp_alloc */
1039 cycle_new, /* tp_new */
1040 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001041};
1042
1043
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001044/* dropwhile object **********************************************************/
1045
1046typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001047 PyObject_HEAD
1048 PyObject *func;
1049 PyObject *it;
1050 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001051} dropwhileobject;
1052
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001053static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001054
1055static PyObject *
1056dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1057{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001058 PyObject *func, *seq;
1059 PyObject *it;
1060 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001061
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001062 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
1063 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001064
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001065 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
1066 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001067
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001068 /* Get iterator. */
1069 it = PyObject_GetIter(seq);
1070 if (it == NULL)
1071 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001072
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001073 /* create dropwhileobject structure */
1074 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1075 if (lz == NULL) {
1076 Py_DECREF(it);
1077 return NULL;
1078 }
1079 Py_INCREF(func);
1080 lz->func = func;
1081 lz->it = it;
1082 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001083
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001084 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001085}
1086
1087static void
1088dropwhile_dealloc(dropwhileobject *lz)
1089{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001090 PyObject_GC_UnTrack(lz);
1091 Py_XDECREF(lz->func);
1092 Py_XDECREF(lz->it);
1093 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001094}
1095
1096static int
1097dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1098{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001099 Py_VISIT(lz->it);
1100 Py_VISIT(lz->func);
1101 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001102}
1103
1104static PyObject *
1105dropwhile_next(dropwhileobject *lz)
1106{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001107 PyObject *item, *good;
1108 PyObject *it = lz->it;
1109 long ok;
1110 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001111
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001112 iternext = *Py_TYPE(it)->tp_iternext;
1113 for (;;) {
1114 item = iternext(it);
1115 if (item == NULL)
1116 return NULL;
1117 if (lz->start == 1)
1118 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001119
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001120 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1121 if (good == NULL) {
1122 Py_DECREF(item);
1123 return NULL;
1124 }
1125 ok = PyObject_IsTrue(good);
1126 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001127 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001128 lz->start = 1;
1129 return item;
1130 }
1131 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001132 if (ok < 0)
1133 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001134 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001135}
1136
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001137static PyObject *
1138dropwhile_reduce(dropwhileobject *lz)
1139{
1140 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1141 lz->func, lz->it, lz->start);
1142}
1143
1144static PyObject *
1145dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1146{
1147 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001148 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001149 return NULL;
1150 lz->start = start;
1151 Py_RETURN_NONE;
1152}
1153
1154static PyMethodDef dropwhile_methods[] = {
1155 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1156 reduce_doc},
1157 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1158 setstate_doc},
1159 {NULL, NULL} /* sentinel */
1160};
1161
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001162PyDoc_STRVAR(dropwhile_doc,
1163"dropwhile(predicate, iterable) --> dropwhile object\n\
1164\n\
1165Drop items from the iterable while predicate(item) is true.\n\
1166Afterwards, return every element until the iterable is exhausted.");
1167
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001168static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001169 PyVarObject_HEAD_INIT(NULL, 0)
1170 "itertools.dropwhile", /* tp_name */
1171 sizeof(dropwhileobject), /* tp_basicsize */
1172 0, /* tp_itemsize */
1173 /* methods */
1174 (destructor)dropwhile_dealloc, /* tp_dealloc */
1175 0, /* tp_print */
1176 0, /* tp_getattr */
1177 0, /* tp_setattr */
1178 0, /* tp_reserved */
1179 0, /* tp_repr */
1180 0, /* tp_as_number */
1181 0, /* tp_as_sequence */
1182 0, /* tp_as_mapping */
1183 0, /* tp_hash */
1184 0, /* tp_call */
1185 0, /* tp_str */
1186 PyObject_GenericGetAttr, /* tp_getattro */
1187 0, /* tp_setattro */
1188 0, /* tp_as_buffer */
1189 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1190 Py_TPFLAGS_BASETYPE, /* tp_flags */
1191 dropwhile_doc, /* tp_doc */
1192 (traverseproc)dropwhile_traverse, /* tp_traverse */
1193 0, /* tp_clear */
1194 0, /* tp_richcompare */
1195 0, /* tp_weaklistoffset */
1196 PyObject_SelfIter, /* tp_iter */
1197 (iternextfunc)dropwhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001198 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001199 0, /* tp_members */
1200 0, /* tp_getset */
1201 0, /* tp_base */
1202 0, /* tp_dict */
1203 0, /* tp_descr_get */
1204 0, /* tp_descr_set */
1205 0, /* tp_dictoffset */
1206 0, /* tp_init */
1207 0, /* tp_alloc */
1208 dropwhile_new, /* tp_new */
1209 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001210};
1211
1212
1213/* takewhile object **********************************************************/
1214
1215typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001216 PyObject_HEAD
1217 PyObject *func;
1218 PyObject *it;
1219 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001220} takewhileobject;
1221
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001222static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001223
1224static PyObject *
1225takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1226{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001227 PyObject *func, *seq;
1228 PyObject *it;
1229 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001230
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001231 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1232 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001233
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001234 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1235 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001236
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001237 /* Get iterator. */
1238 it = PyObject_GetIter(seq);
1239 if (it == NULL)
1240 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001241
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001242 /* create takewhileobject structure */
1243 lz = (takewhileobject *)type->tp_alloc(type, 0);
1244 if (lz == NULL) {
1245 Py_DECREF(it);
1246 return NULL;
1247 }
1248 Py_INCREF(func);
1249 lz->func = func;
1250 lz->it = it;
1251 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001252
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001253 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001254}
1255
1256static void
1257takewhile_dealloc(takewhileobject *lz)
1258{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001259 PyObject_GC_UnTrack(lz);
1260 Py_XDECREF(lz->func);
1261 Py_XDECREF(lz->it);
1262 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001263}
1264
1265static int
1266takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1267{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001268 Py_VISIT(lz->it);
1269 Py_VISIT(lz->func);
1270 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001271}
1272
1273static PyObject *
1274takewhile_next(takewhileobject *lz)
1275{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001276 PyObject *item, *good;
1277 PyObject *it = lz->it;
1278 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001279
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001280 if (lz->stop == 1)
1281 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001282
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001283 item = (*Py_TYPE(it)->tp_iternext)(it);
1284 if (item == NULL)
1285 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001286
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001287 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1288 if (good == NULL) {
1289 Py_DECREF(item);
1290 return NULL;
1291 }
1292 ok = PyObject_IsTrue(good);
1293 Py_DECREF(good);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001294 if (ok == 1)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001295 return item;
1296 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001297 if (ok == 0)
1298 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001299 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001300}
1301
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001302static PyObject *
1303takewhile_reduce(takewhileobject *lz)
1304{
1305 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1306 lz->func, lz->it, lz->stop);
1307}
1308
1309static PyObject *
1310takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1311{
1312 int stop = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001313 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001314 return NULL;
1315 lz->stop = stop;
1316 Py_RETURN_NONE;
1317}
1318
1319static PyMethodDef takewhile_reduce_methods[] = {
1320 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1321 reduce_doc},
1322 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1323 setstate_doc},
1324 {NULL, NULL} /* sentinel */
1325};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001326PyDoc_STRVAR(takewhile_doc,
1327"takewhile(predicate, iterable) --> takewhile object\n\
1328\n\
1329Return successive entries from an iterable as long as the \n\
1330predicate evaluates to true for each entry.");
1331
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001332static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001333 PyVarObject_HEAD_INIT(NULL, 0)
1334 "itertools.takewhile", /* tp_name */
1335 sizeof(takewhileobject), /* tp_basicsize */
1336 0, /* tp_itemsize */
1337 /* methods */
1338 (destructor)takewhile_dealloc, /* tp_dealloc */
1339 0, /* tp_print */
1340 0, /* tp_getattr */
1341 0, /* tp_setattr */
1342 0, /* tp_reserved */
1343 0, /* tp_repr */
1344 0, /* tp_as_number */
1345 0, /* tp_as_sequence */
1346 0, /* tp_as_mapping */
1347 0, /* tp_hash */
1348 0, /* tp_call */
1349 0, /* tp_str */
1350 PyObject_GenericGetAttr, /* tp_getattro */
1351 0, /* tp_setattro */
1352 0, /* tp_as_buffer */
1353 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1354 Py_TPFLAGS_BASETYPE, /* tp_flags */
1355 takewhile_doc, /* tp_doc */
1356 (traverseproc)takewhile_traverse, /* tp_traverse */
1357 0, /* tp_clear */
1358 0, /* tp_richcompare */
1359 0, /* tp_weaklistoffset */
1360 PyObject_SelfIter, /* tp_iter */
1361 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001362 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001363 0, /* tp_members */
1364 0, /* tp_getset */
1365 0, /* tp_base */
1366 0, /* tp_dict */
1367 0, /* tp_descr_get */
1368 0, /* tp_descr_set */
1369 0, /* tp_dictoffset */
1370 0, /* tp_init */
1371 0, /* tp_alloc */
1372 takewhile_new, /* tp_new */
1373 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001374};
1375
1376
1377/* islice object ************************************************************/
1378
1379typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001380 PyObject_HEAD
1381 PyObject *it;
1382 Py_ssize_t next;
1383 Py_ssize_t stop;
1384 Py_ssize_t step;
1385 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001386} isliceobject;
1387
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001388static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001389
1390static PyObject *
1391islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1392{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001393 PyObject *seq;
1394 Py_ssize_t start=0, stop=-1, step=1;
1395 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1396 Py_ssize_t numargs;
1397 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001398
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001399 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1400 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001401
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001402 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1403 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001404
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001405 numargs = PyTuple_Size(args);
1406 if (numargs == 2) {
1407 if (a1 != Py_None) {
1408 stop = PyLong_AsSsize_t(a1);
1409 if (stop == -1) {
1410 if (PyErr_Occurred())
1411 PyErr_Clear();
1412 PyErr_SetString(PyExc_ValueError,
1413 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1414 return NULL;
1415 }
1416 }
1417 } else {
1418 if (a1 != Py_None)
1419 start = PyLong_AsSsize_t(a1);
1420 if (start == -1 && PyErr_Occurred())
1421 PyErr_Clear();
1422 if (a2 != Py_None) {
1423 stop = PyLong_AsSsize_t(a2);
1424 if (stop == -1) {
1425 if (PyErr_Occurred())
1426 PyErr_Clear();
1427 PyErr_SetString(PyExc_ValueError,
1428 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1429 return NULL;
1430 }
1431 }
1432 }
1433 if (start<0 || stop<-1) {
1434 PyErr_SetString(PyExc_ValueError,
1435 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1436 return NULL;
1437 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001438
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001439 if (a3 != NULL) {
1440 if (a3 != Py_None)
1441 step = PyLong_AsSsize_t(a3);
1442 if (step == -1 && PyErr_Occurred())
1443 PyErr_Clear();
1444 }
1445 if (step<1) {
1446 PyErr_SetString(PyExc_ValueError,
1447 "Step for islice() must be a positive integer or None.");
1448 return NULL;
1449 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001450
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001451 /* Get iterator. */
1452 it = PyObject_GetIter(seq);
1453 if (it == NULL)
1454 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001455
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001456 /* create isliceobject structure */
1457 lz = (isliceobject *)type->tp_alloc(type, 0);
1458 if (lz == NULL) {
1459 Py_DECREF(it);
1460 return NULL;
1461 }
1462 lz->it = it;
1463 lz->next = start;
1464 lz->stop = stop;
1465 lz->step = step;
1466 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001467
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001468 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469}
1470
1471static void
1472islice_dealloc(isliceobject *lz)
1473{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001474 PyObject_GC_UnTrack(lz);
1475 Py_XDECREF(lz->it);
1476 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001477}
1478
1479static int
1480islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1481{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001482 Py_VISIT(lz->it);
1483 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001484}
1485
1486static PyObject *
1487islice_next(isliceobject *lz)
1488{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001489 PyObject *item;
1490 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001491 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001492 Py_ssize_t oldnext;
1493 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001494
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001495 iternext = *Py_TYPE(it)->tp_iternext;
1496 while (lz->cnt < lz->next) {
1497 item = iternext(it);
1498 if (item == NULL)
1499 return NULL;
1500 Py_DECREF(item);
1501 lz->cnt++;
1502 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001503 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001504 return NULL;
1505 item = iternext(it);
1506 if (item == NULL)
1507 return NULL;
1508 lz->cnt++;
1509 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001510 /* The (size_t) cast below avoids the danger of undefined
1511 behaviour from signed integer overflow. */
1512 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001513 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1514 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001515 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001516}
1517
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001518static PyObject *
1519islice_reduce(isliceobject *lz)
1520{
1521 /* When unpickled, generate a new object with the same bounds,
1522 * then 'setstate' with the next and count
1523 */
1524 PyObject *stop;
1525 if (lz->stop == -1) {
1526 stop = Py_None;
1527 Py_INCREF(stop);
1528 } else {
1529 stop = PyLong_FromSsize_t(lz->stop);
1530 if (stop == NULL)
1531 return NULL;
1532 }
1533 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1534 lz->it, lz->next, stop, lz->step,
1535 lz->cnt);
1536}
1537
1538static PyObject *
1539islice_setstate(isliceobject *lz, PyObject *state)
1540{
1541 Py_ssize_t cnt = PyLong_AsSsize_t(state);
1542 if (cnt == -1 && PyErr_Occurred())
1543 return NULL;
1544 lz->cnt = cnt;
1545 Py_RETURN_NONE;
1546}
1547
1548static PyMethodDef islice_methods[] = {
1549 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1550 reduce_doc},
1551 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1552 setstate_doc},
1553 {NULL, NULL} /* sentinel */
1554};
1555
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001556PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001557"islice(iterable, stop) --> islice object\n\
1558islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001559\n\
1560Return an iterator whose next() method returns selected values from an\n\
1561iterable. If start is specified, will skip all preceding elements;\n\
1562otherwise, start defaults to zero. Step defaults to one. If\n\
1563specified as another value, step determines how many values are \n\
1564skipped between successive calls. Works like a slice() on a list\n\
1565but returns an iterator.");
1566
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001567static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001568 PyVarObject_HEAD_INIT(NULL, 0)
1569 "itertools.islice", /* tp_name */
1570 sizeof(isliceobject), /* tp_basicsize */
1571 0, /* tp_itemsize */
1572 /* methods */
1573 (destructor)islice_dealloc, /* tp_dealloc */
1574 0, /* tp_print */
1575 0, /* tp_getattr */
1576 0, /* tp_setattr */
1577 0, /* tp_reserved */
1578 0, /* tp_repr */
1579 0, /* tp_as_number */
1580 0, /* tp_as_sequence */
1581 0, /* tp_as_mapping */
1582 0, /* tp_hash */
1583 0, /* tp_call */
1584 0, /* tp_str */
1585 PyObject_GenericGetAttr, /* tp_getattro */
1586 0, /* tp_setattro */
1587 0, /* tp_as_buffer */
1588 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1589 Py_TPFLAGS_BASETYPE, /* tp_flags */
1590 islice_doc, /* tp_doc */
1591 (traverseproc)islice_traverse, /* tp_traverse */
1592 0, /* tp_clear */
1593 0, /* tp_richcompare */
1594 0, /* tp_weaklistoffset */
1595 PyObject_SelfIter, /* tp_iter */
1596 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001597 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001598 0, /* tp_members */
1599 0, /* tp_getset */
1600 0, /* tp_base */
1601 0, /* tp_dict */
1602 0, /* tp_descr_get */
1603 0, /* tp_descr_set */
1604 0, /* tp_dictoffset */
1605 0, /* tp_init */
1606 0, /* tp_alloc */
1607 islice_new, /* tp_new */
1608 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001609};
1610
1611
1612/* starmap object ************************************************************/
1613
1614typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001615 PyObject_HEAD
1616 PyObject *func;
1617 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001618} starmapobject;
1619
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001620static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001621
1622static PyObject *
1623starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1624{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001625 PyObject *func, *seq;
1626 PyObject *it;
1627 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001628
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001629 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1630 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001631
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001632 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1633 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001634
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001635 /* Get iterator. */
1636 it = PyObject_GetIter(seq);
1637 if (it == NULL)
1638 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001639
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001640 /* create starmapobject structure */
1641 lz = (starmapobject *)type->tp_alloc(type, 0);
1642 if (lz == NULL) {
1643 Py_DECREF(it);
1644 return NULL;
1645 }
1646 Py_INCREF(func);
1647 lz->func = func;
1648 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001649
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001650 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001651}
1652
1653static void
1654starmap_dealloc(starmapobject *lz)
1655{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001656 PyObject_GC_UnTrack(lz);
1657 Py_XDECREF(lz->func);
1658 Py_XDECREF(lz->it);
1659 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001660}
1661
1662static int
1663starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1664{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001665 Py_VISIT(lz->it);
1666 Py_VISIT(lz->func);
1667 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001668}
1669
1670static PyObject *
1671starmap_next(starmapobject *lz)
1672{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001673 PyObject *args;
1674 PyObject *result;
1675 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001676
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001677 args = (*Py_TYPE(it)->tp_iternext)(it);
1678 if (args == NULL)
1679 return NULL;
1680 if (!PyTuple_CheckExact(args)) {
1681 PyObject *newargs = PySequence_Tuple(args);
1682 Py_DECREF(args);
1683 if (newargs == NULL)
1684 return NULL;
1685 args = newargs;
1686 }
1687 result = PyObject_Call(lz->func, args, NULL);
1688 Py_DECREF(args);
1689 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001690}
1691
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001692static PyObject *
1693starmap_reduce(starmapobject *lz)
1694{
1695 /* Just pickle the iterator */
1696 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1697}
1698
1699static PyMethodDef starmap_methods[] = {
1700 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1701 reduce_doc},
1702 {NULL, NULL} /* sentinel */
1703};
1704
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001705PyDoc_STRVAR(starmap_doc,
1706"starmap(function, sequence) --> starmap object\n\
1707\n\
1708Return an iterator whose values are returned from the function evaluated\n\
1709with a argument tuple taken from the given sequence.");
1710
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001711static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001712 PyVarObject_HEAD_INIT(NULL, 0)
1713 "itertools.starmap", /* tp_name */
1714 sizeof(starmapobject), /* tp_basicsize */
1715 0, /* tp_itemsize */
1716 /* methods */
1717 (destructor)starmap_dealloc, /* tp_dealloc */
1718 0, /* tp_print */
1719 0, /* tp_getattr */
1720 0, /* tp_setattr */
1721 0, /* tp_reserved */
1722 0, /* tp_repr */
1723 0, /* tp_as_number */
1724 0, /* tp_as_sequence */
1725 0, /* tp_as_mapping */
1726 0, /* tp_hash */
1727 0, /* tp_call */
1728 0, /* tp_str */
1729 PyObject_GenericGetAttr, /* tp_getattro */
1730 0, /* tp_setattro */
1731 0, /* tp_as_buffer */
1732 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1733 Py_TPFLAGS_BASETYPE, /* tp_flags */
1734 starmap_doc, /* tp_doc */
1735 (traverseproc)starmap_traverse, /* tp_traverse */
1736 0, /* tp_clear */
1737 0, /* tp_richcompare */
1738 0, /* tp_weaklistoffset */
1739 PyObject_SelfIter, /* tp_iter */
1740 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001741 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001742 0, /* tp_members */
1743 0, /* tp_getset */
1744 0, /* tp_base */
1745 0, /* tp_dict */
1746 0, /* tp_descr_get */
1747 0, /* tp_descr_set */
1748 0, /* tp_dictoffset */
1749 0, /* tp_init */
1750 0, /* tp_alloc */
1751 starmap_new, /* tp_new */
1752 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001753};
1754
1755
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001756/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001757
1758typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001759 PyObject_HEAD
1760 PyObject *source; /* Iterator over input iterables */
1761 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001762} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001763
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001764static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001765
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001766static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001767chain_new_internal(PyTypeObject *type, PyObject *source)
1768{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001769 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001770
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001771 lz = (chainobject *)type->tp_alloc(type, 0);
1772 if (lz == NULL) {
1773 Py_DECREF(source);
1774 return NULL;
1775 }
1776
1777 lz->source = source;
1778 lz->active = NULL;
1779 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001780}
1781
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001782static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001783chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001784{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001785 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001786
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001787 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1788 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001789
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001790 source = PyObject_GetIter(args);
1791 if (source == NULL)
1792 return NULL;
1793
1794 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001795}
1796
1797static PyObject *
1798chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1799{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001800 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001801
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001802 source = PyObject_GetIter(arg);
1803 if (source == NULL)
1804 return NULL;
1805
1806 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001807}
1808
1809static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001810chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001811{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001812 PyObject_GC_UnTrack(lz);
1813 Py_XDECREF(lz->active);
1814 Py_XDECREF(lz->source);
1815 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001816}
1817
Raymond Hettinger2012f172003-02-07 05:32:58 +00001818static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001819chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001820{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001821 Py_VISIT(lz->source);
1822 Py_VISIT(lz->active);
1823 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001824}
1825
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001826static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001827chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001828{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001829 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001830
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001831 if (lz->source == NULL)
1832 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001833
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001834 if (lz->active == NULL) {
1835 PyObject *iterable = PyIter_Next(lz->source);
1836 if (iterable == NULL) {
1837 Py_CLEAR(lz->source);
1838 return NULL; /* no more input sources */
1839 }
1840 lz->active = PyObject_GetIter(iterable);
1841 Py_DECREF(iterable);
1842 if (lz->active == NULL) {
1843 Py_CLEAR(lz->source);
1844 return NULL; /* input not iterable */
1845 }
1846 }
1847 item = PyIter_Next(lz->active);
1848 if (item != NULL)
1849 return item;
1850 if (PyErr_Occurred()) {
1851 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1852 PyErr_Clear();
1853 else
1854 return NULL; /* input raised an exception */
1855 }
1856 Py_CLEAR(lz->active);
1857 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001858}
1859
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001860static PyObject *
1861chain_reduce(chainobject *lz)
1862{
1863 if (lz->source) {
1864 /* we can't pickle function objects (itertools.from_iterable) so
1865 * we must use setstate to replace the iterable. One day we
1866 * will fix pickling of functions
1867 */
1868 if (lz->active) {
1869 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1870 } else {
1871 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1872 }
1873 } else {
1874 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1875 }
1876 return NULL;
1877}
1878
1879static PyObject *
1880chain_setstate(chainobject *lz, PyObject *state)
1881{
1882 PyObject *source, *active=NULL;
1883 if (! PyArg_ParseTuple(state, "O|O", &source, &active))
1884 return NULL;
1885
1886 Py_CLEAR(lz->source);
1887 lz->source = source;
1888 Py_INCREF(lz->source);
1889 Py_CLEAR(lz->active);
1890 lz->active = active;
1891 Py_XINCREF(lz->active);
1892 Py_RETURN_NONE;
1893}
1894
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001895PyDoc_STRVAR(chain_doc,
1896"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001897\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001898Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001899first iterable until it is exhausted, then elements from the next\n\
1900iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001901
Christian Heimesf16baeb2008-02-29 14:57:44 +00001902PyDoc_STRVAR(chain_from_iterable_doc,
1903"chain.from_iterable(iterable) --> chain object\n\
1904\n\
1905Alternate chain() contructor taking a single iterable argument\n\
1906that evaluates lazily.");
1907
1908static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001909 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1910 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001911 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1912 reduce_doc},
1913 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1914 setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001915 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001916};
1917
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001918static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001919 PyVarObject_HEAD_INIT(NULL, 0)
1920 "itertools.chain", /* tp_name */
1921 sizeof(chainobject), /* tp_basicsize */
1922 0, /* tp_itemsize */
1923 /* methods */
1924 (destructor)chain_dealloc, /* tp_dealloc */
1925 0, /* tp_print */
1926 0, /* tp_getattr */
1927 0, /* tp_setattr */
1928 0, /* tp_reserved */
1929 0, /* tp_repr */
1930 0, /* tp_as_number */
1931 0, /* tp_as_sequence */
1932 0, /* tp_as_mapping */
1933 0, /* tp_hash */
1934 0, /* tp_call */
1935 0, /* tp_str */
1936 PyObject_GenericGetAttr, /* tp_getattro */
1937 0, /* tp_setattro */
1938 0, /* tp_as_buffer */
1939 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1940 Py_TPFLAGS_BASETYPE, /* tp_flags */
1941 chain_doc, /* tp_doc */
1942 (traverseproc)chain_traverse, /* tp_traverse */
1943 0, /* tp_clear */
1944 0, /* tp_richcompare */
1945 0, /* tp_weaklistoffset */
1946 PyObject_SelfIter, /* tp_iter */
1947 (iternextfunc)chain_next, /* tp_iternext */
1948 chain_methods, /* tp_methods */
1949 0, /* tp_members */
1950 0, /* tp_getset */
1951 0, /* tp_base */
1952 0, /* tp_dict */
1953 0, /* tp_descr_get */
1954 0, /* tp_descr_set */
1955 0, /* tp_dictoffset */
1956 0, /* tp_init */
1957 0, /* tp_alloc */
1958 chain_new, /* tp_new */
1959 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001960};
1961
1962
Christian Heimesc3f30c42008-02-22 16:37:40 +00001963/* product object ************************************************************/
1964
1965typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001966 PyObject_HEAD
1967 PyObject *pools; /* tuple of pool tuples */
1968 Py_ssize_t *indices; /* one index per pool */
1969 PyObject *result; /* most recently returned result tuple */
1970 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001971} productobject;
1972
1973static PyTypeObject product_type;
1974
1975static PyObject *
1976product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1977{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001978 productobject *lz;
1979 Py_ssize_t nargs, npools, repeat=1;
1980 PyObject *pools = NULL;
1981 Py_ssize_t *indices = NULL;
1982 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001983
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001984 if (kwds != NULL) {
1985 char *kwlist[] = {"repeat", 0};
1986 PyObject *tmpargs = PyTuple_New(0);
1987 if (tmpargs == NULL)
1988 return NULL;
1989 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1990 Py_DECREF(tmpargs);
1991 return NULL;
1992 }
1993 Py_DECREF(tmpargs);
1994 if (repeat < 0) {
1995 PyErr_SetString(PyExc_ValueError,
1996 "repeat argument cannot be negative");
1997 return NULL;
1998 }
1999 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002000
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002001 assert(PyTuple_Check(args));
2002 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
2003 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002005 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
2006 if (indices == NULL) {
2007 PyErr_NoMemory();
2008 goto error;
2009 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002010
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002011 pools = PyTuple_New(npools);
2012 if (pools == NULL)
2013 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002014
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002015 for (i=0; i < nargs ; ++i) {
2016 PyObject *item = PyTuple_GET_ITEM(args, i);
2017 PyObject *pool = PySequence_Tuple(item);
2018 if (pool == NULL)
2019 goto error;
2020 PyTuple_SET_ITEM(pools, i, pool);
2021 indices[i] = 0;
2022 }
2023 for ( ; i < npools; ++i) {
2024 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2025 Py_INCREF(pool);
2026 PyTuple_SET_ITEM(pools, i, pool);
2027 indices[i] = 0;
2028 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002029
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002030 /* create productobject structure */
2031 lz = (productobject *)type->tp_alloc(type, 0);
2032 if (lz == NULL)
2033 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002034
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002035 lz->pools = pools;
2036 lz->indices = indices;
2037 lz->result = NULL;
2038 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002039
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002040 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002041
2042error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002043 if (indices != NULL)
2044 PyMem_Free(indices);
2045 Py_XDECREF(pools);
2046 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002047}
2048
2049static void
2050product_dealloc(productobject *lz)
2051{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002052 PyObject_GC_UnTrack(lz);
2053 Py_XDECREF(lz->pools);
2054 Py_XDECREF(lz->result);
2055 if (lz->indices != NULL)
2056 PyMem_Free(lz->indices);
2057 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002058}
2059
2060static int
2061product_traverse(productobject *lz, visitproc visit, void *arg)
2062{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002063 Py_VISIT(lz->pools);
2064 Py_VISIT(lz->result);
2065 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002066}
2067
2068static PyObject *
2069product_next(productobject *lz)
2070{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002071 PyObject *pool;
2072 PyObject *elem;
2073 PyObject *oldelem;
2074 PyObject *pools = lz->pools;
2075 PyObject *result = lz->result;
2076 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2077 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002078
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002079 if (lz->stopped)
2080 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002081
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002082 if (result == NULL) {
2083 /* On the first pass, return an initial tuple filled with the
2084 first element from each pool. */
2085 result = PyTuple_New(npools);
2086 if (result == NULL)
2087 goto empty;
2088 lz->result = result;
2089 for (i=0; i < npools; i++) {
2090 pool = PyTuple_GET_ITEM(pools, i);
2091 if (PyTuple_GET_SIZE(pool) == 0)
2092 goto empty;
2093 elem = PyTuple_GET_ITEM(pool, 0);
2094 Py_INCREF(elem);
2095 PyTuple_SET_ITEM(result, i, elem);
2096 }
2097 } else {
2098 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002099
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002100 /* Copy the previous result tuple or re-use it if available */
2101 if (Py_REFCNT(result) > 1) {
2102 PyObject *old_result = result;
2103 result = PyTuple_New(npools);
2104 if (result == NULL)
2105 goto empty;
2106 lz->result = result;
2107 for (i=0; i < npools; i++) {
2108 elem = PyTuple_GET_ITEM(old_result, i);
2109 Py_INCREF(elem);
2110 PyTuple_SET_ITEM(result, i, elem);
2111 }
2112 Py_DECREF(old_result);
2113 }
2114 /* Now, we've got the only copy so we can update it in-place */
2115 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002116
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002117 /* Update the pool indices right-to-left. Only advance to the
2118 next pool when the previous one rolls-over */
2119 for (i=npools-1 ; i >= 0 ; i--) {
2120 pool = PyTuple_GET_ITEM(pools, i);
2121 indices[i]++;
2122 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2123 /* Roll-over and advance to next pool */
2124 indices[i] = 0;
2125 elem = PyTuple_GET_ITEM(pool, 0);
2126 Py_INCREF(elem);
2127 oldelem = PyTuple_GET_ITEM(result, i);
2128 PyTuple_SET_ITEM(result, i, elem);
2129 Py_DECREF(oldelem);
2130 } else {
2131 /* No rollover. Just increment and stop here. */
2132 elem = PyTuple_GET_ITEM(pool, indices[i]);
2133 Py_INCREF(elem);
2134 oldelem = PyTuple_GET_ITEM(result, i);
2135 PyTuple_SET_ITEM(result, i, elem);
2136 Py_DECREF(oldelem);
2137 break;
2138 }
2139 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002140
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002141 /* If i is negative, then the indices have all rolled-over
2142 and we're done. */
2143 if (i < 0)
2144 goto empty;
2145 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002146
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002147 Py_INCREF(result);
2148 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002149
2150empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002151 lz->stopped = 1;
2152 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002153}
2154
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002155static PyObject *
2156product_reduce(productobject *lz)
2157{
2158 if (lz->stopped) {
2159 return Py_BuildValue("O(())", Py_TYPE(lz));
2160 } else if (lz->result == NULL) {
2161 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2162 } else {
2163 PyObject *indices;
2164 Py_ssize_t n, i;
2165
2166 /* we must pickle the indices use them for setstate, and
2167 * additionally indicate that the iterator has started
2168 */
2169 n = PyTuple_GET_SIZE(lz->pools);
2170 indices = PyTuple_New(n);
2171 if (indices == NULL)
2172 return NULL;
2173 for (i=0; i<n; i++){
2174 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2175 if (!index) {
2176 Py_DECREF(indices);
2177 return NULL;
2178 }
2179 PyTuple_SET_ITEM(indices, i, index);
2180 }
2181 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2182 }
2183}
2184
2185static PyObject *
2186product_setstate(productobject *lz, PyObject *state)
2187{
2188 PyObject *result;
2189 Py_ssize_t n, i;
2190
2191 n = PyTuple_GET_SIZE(lz->pools);
2192 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2193 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2194 return NULL;
2195 }
2196 for (i=0; i<n; i++)
2197 {
2198 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2199 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2200 if (index < 0 && PyErr_Occurred())
2201 return NULL; /* not an integer */
2202 /* clamp the index */
2203 if (index < 0)
2204 index = 0;
2205 else if (index > n-1)
2206 index = n-1;
2207 lz->indices[i] = index;
2208 }
2209
2210 result = PyTuple_New(n);
2211 if (!result)
2212 return NULL;
2213 for (i=0; i<n; i++) {
2214 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2215 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2216 Py_INCREF(element);
2217 PyTuple_SET_ITEM(result, i, element);
2218 }
2219 Py_CLEAR(lz->result);
2220 lz->result = result;
2221 Py_RETURN_NONE;
2222}
2223
2224static PyMethodDef product_methods[] = {
2225 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2226 reduce_doc},
2227 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2228 setstate_doc},
2229 {NULL, NULL} /* sentinel */
2230};
2231
Christian Heimesc3f30c42008-02-22 16:37:40 +00002232PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002233"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002234\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002235Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002236For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2237The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2238cycle in a manner similar to an odometer (with the rightmost element changing\n\
2239on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002240To compute the product of an iterable with itself, specify the number\n\
2241of repetitions with the optional repeat keyword argument. For example,\n\
2242product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002243product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2244product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2245
2246static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002247 PyVarObject_HEAD_INIT(NULL, 0)
2248 "itertools.product", /* tp_name */
2249 sizeof(productobject), /* tp_basicsize */
2250 0, /* tp_itemsize */
2251 /* methods */
2252 (destructor)product_dealloc, /* tp_dealloc */
2253 0, /* tp_print */
2254 0, /* tp_getattr */
2255 0, /* tp_setattr */
2256 0, /* tp_reserved */
2257 0, /* tp_repr */
2258 0, /* tp_as_number */
2259 0, /* tp_as_sequence */
2260 0, /* tp_as_mapping */
2261 0, /* tp_hash */
2262 0, /* tp_call */
2263 0, /* tp_str */
2264 PyObject_GenericGetAttr, /* tp_getattro */
2265 0, /* tp_setattro */
2266 0, /* tp_as_buffer */
2267 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2268 Py_TPFLAGS_BASETYPE, /* tp_flags */
2269 product_doc, /* tp_doc */
2270 (traverseproc)product_traverse, /* tp_traverse */
2271 0, /* tp_clear */
2272 0, /* tp_richcompare */
2273 0, /* tp_weaklistoffset */
2274 PyObject_SelfIter, /* tp_iter */
2275 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002276 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002277 0, /* tp_members */
2278 0, /* tp_getset */
2279 0, /* tp_base */
2280 0, /* tp_dict */
2281 0, /* tp_descr_get */
2282 0, /* tp_descr_set */
2283 0, /* tp_dictoffset */
2284 0, /* tp_init */
2285 0, /* tp_alloc */
2286 product_new, /* tp_new */
2287 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002288};
2289
2290
Christian Heimes380f7f22008-02-28 11:19:05 +00002291/* combinations object ************************************************************/
2292
2293typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002294 PyObject_HEAD
2295 PyObject *pool; /* input converted to a tuple */
2296 Py_ssize_t *indices; /* one index per result element */
2297 PyObject *result; /* most recently returned result tuple */
2298 Py_ssize_t r; /* size of result tuple */
2299 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002300} combinationsobject;
2301
2302static PyTypeObject combinations_type;
2303
2304static PyObject *
2305combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2306{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002307 combinationsobject *co;
2308 Py_ssize_t n;
2309 Py_ssize_t r;
2310 PyObject *pool = NULL;
2311 PyObject *iterable = NULL;
2312 Py_ssize_t *indices = NULL;
2313 Py_ssize_t i;
2314 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002315
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002316 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2317 &iterable, &r))
2318 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002319
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002320 pool = PySequence_Tuple(iterable);
2321 if (pool == NULL)
2322 goto error;
2323 n = PyTuple_GET_SIZE(pool);
2324 if (r < 0) {
2325 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2326 goto error;
2327 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002328
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002329 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2330 if (indices == NULL) {
2331 PyErr_NoMemory();
2332 goto error;
2333 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002334
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002335 for (i=0 ; i<r ; i++)
2336 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002337
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002338 /* create combinationsobject structure */
2339 co = (combinationsobject *)type->tp_alloc(type, 0);
2340 if (co == NULL)
2341 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002342
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002343 co->pool = pool;
2344 co->indices = indices;
2345 co->result = NULL;
2346 co->r = r;
2347 co->stopped = r > n ? 1 : 0;
2348
2349 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002350
2351error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002352 if (indices != NULL)
2353 PyMem_Free(indices);
2354 Py_XDECREF(pool);
2355 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002356}
2357
2358static void
2359combinations_dealloc(combinationsobject *co)
2360{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002361 PyObject_GC_UnTrack(co);
2362 Py_XDECREF(co->pool);
2363 Py_XDECREF(co->result);
2364 if (co->indices != NULL)
2365 PyMem_Free(co->indices);
2366 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002367}
2368
2369static int
2370combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2371{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002372 Py_VISIT(co->pool);
2373 Py_VISIT(co->result);
2374 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002375}
2376
2377static PyObject *
2378combinations_next(combinationsobject *co)
2379{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002380 PyObject *elem;
2381 PyObject *oldelem;
2382 PyObject *pool = co->pool;
2383 Py_ssize_t *indices = co->indices;
2384 PyObject *result = co->result;
2385 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2386 Py_ssize_t r = co->r;
2387 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002388
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002389 if (co->stopped)
2390 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002391
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002392 if (result == NULL) {
2393 /* On the first pass, initialize result tuple using the indices */
2394 result = PyTuple_New(r);
2395 if (result == NULL)
2396 goto empty;
2397 co->result = result;
2398 for (i=0; i<r ; i++) {
2399 index = indices[i];
2400 elem = PyTuple_GET_ITEM(pool, index);
2401 Py_INCREF(elem);
2402 PyTuple_SET_ITEM(result, i, elem);
2403 }
2404 } else {
2405 /* Copy the previous result tuple or re-use it if available */
2406 if (Py_REFCNT(result) > 1) {
2407 PyObject *old_result = result;
2408 result = PyTuple_New(r);
2409 if (result == NULL)
2410 goto empty;
2411 co->result = result;
2412 for (i=0; i<r ; i++) {
2413 elem = PyTuple_GET_ITEM(old_result, i);
2414 Py_INCREF(elem);
2415 PyTuple_SET_ITEM(result, i, elem);
2416 }
2417 Py_DECREF(old_result);
2418 }
2419 /* Now, we've got the only copy so we can update it in-place
2420 * CPython's empty tuple is a singleton and cached in
2421 * PyTuple's freelist.
2422 */
2423 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002424
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002425 /* Scan indices right-to-left until finding one that is not
2426 at its maximum (i + n - r). */
2427 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2428 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002429
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002430 /* If i is negative, then the indices are all at
2431 their maximum value and we're done. */
2432 if (i < 0)
2433 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002434
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002435 /* Increment the current index which we know is not at its
2436 maximum. Then move back to the right setting each index
2437 to its lowest possible value (one higher than the index
2438 to its left -- this maintains the sort order invariant). */
2439 indices[i]++;
2440 for (j=i+1 ; j<r ; j++)
2441 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002442
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002443 /* Update the result tuple for the new indices
2444 starting with i, the leftmost index that changed */
2445 for ( ; i<r ; i++) {
2446 index = indices[i];
2447 elem = PyTuple_GET_ITEM(pool, index);
2448 Py_INCREF(elem);
2449 oldelem = PyTuple_GET_ITEM(result, i);
2450 PyTuple_SET_ITEM(result, i, elem);
2451 Py_DECREF(oldelem);
2452 }
2453 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002454
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002455 Py_INCREF(result);
2456 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002457
2458empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002459 co->stopped = 1;
2460 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002461}
2462
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002463static PyObject *
2464combinations_reduce(combinationsobject *lz)
2465{
2466 if (lz->result == NULL) {
2467 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2468 } else if (lz->stopped) {
2469 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2470 } else {
2471 PyObject *indices;
2472 Py_ssize_t i;
2473
2474 /* we must pickle the indices and use them for setstate */
2475 indices = PyTuple_New(lz->r);
2476 if (!indices)
2477 return NULL;
2478 for (i=0; i<lz->r; i++)
2479 {
2480 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2481 if (!index) {
2482 Py_DECREF(indices);
2483 return NULL;
2484 }
2485 PyTuple_SET_ITEM(indices, i, index);
2486 }
2487
2488 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2489 }
2490}
2491
2492static PyObject *
2493combinations_setstate(combinationsobject *lz, PyObject *state)
2494{
2495 PyObject *result;
2496 Py_ssize_t i;
2497 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2498
2499 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2500 {
2501 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2502 return NULL;
2503 }
2504
2505 for (i=0; i<lz->r; i++)
2506 {
2507 Py_ssize_t max;
2508 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2509 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2510 if (index == -1 && PyErr_Occurred())
2511 return NULL; /* not an integer */
2512 max = i + n - lz->r;
2513 /* clamp the index (beware of negative max) */
2514 if (index > max)
2515 index = max;
2516 if (index < 0)
2517 index = 0;
2518 lz->indices[i] = index;
2519 }
2520
2521 result = PyTuple_New(lz->r);
2522 if (result == NULL)
2523 return NULL;
2524 for (i=0; i<lz->r; i++) {
2525 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2526 Py_INCREF(element);
2527 PyTuple_SET_ITEM(result, i, element);
2528 }
2529
2530 Py_CLEAR(lz->result);
2531 lz->result = result;
2532 Py_RETURN_NONE;
2533}
2534
2535static PyMethodDef combinations_methods[] = {
2536 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2537 reduce_doc},
2538 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2539 setstate_doc},
2540 {NULL, NULL} /* sentinel */
2541};
2542
Christian Heimes380f7f22008-02-28 11:19:05 +00002543PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002544"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002545\n\
2546Return successive r-length combinations of elements in the iterable.\n\n\
2547combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2548
2549static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002550 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002551 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002552 sizeof(combinationsobject), /* tp_basicsize */
2553 0, /* tp_itemsize */
2554 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002555 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002556 0, /* tp_print */
2557 0, /* tp_getattr */
2558 0, /* tp_setattr */
2559 0, /* tp_reserved */
2560 0, /* tp_repr */
2561 0, /* tp_as_number */
2562 0, /* tp_as_sequence */
2563 0, /* tp_as_mapping */
2564 0, /* tp_hash */
2565 0, /* tp_call */
2566 0, /* tp_str */
2567 PyObject_GenericGetAttr, /* tp_getattro */
2568 0, /* tp_setattro */
2569 0, /* tp_as_buffer */
2570 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2571 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002572 combinations_doc, /* tp_doc */
2573 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002574 0, /* tp_clear */
2575 0, /* tp_richcompare */
2576 0, /* tp_weaklistoffset */
2577 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002578 (iternextfunc)combinations_next, /* tp_iternext */
2579 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002580 0, /* tp_members */
2581 0, /* tp_getset */
2582 0, /* tp_base */
2583 0, /* tp_dict */
2584 0, /* tp_descr_get */
2585 0, /* tp_descr_set */
2586 0, /* tp_dictoffset */
2587 0, /* tp_init */
2588 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002589 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002590 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002591};
2592
2593
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002594/* combinations with replacement object *******************************************/
2595
2596/* Equivalent to:
2597
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002598 def combinations_with_replacement(iterable, r):
2599 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2600 # number items returned: (n+r-1)! / r! / (n-1)!
2601 pool = tuple(iterable)
2602 n = len(pool)
2603 indices = [0] * r
2604 yield tuple(pool[i] for i in indices)
2605 while 1:
2606 for i in reversed(range(r)):
2607 if indices[i] != n - 1:
2608 break
2609 else:
2610 return
2611 indices[i:] = [indices[i] + 1] * (r - i)
2612 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002613
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002614 def combinations_with_replacement2(iterable, r):
2615 'Alternate version that filters from product()'
2616 pool = tuple(iterable)
2617 n = len(pool)
2618 for indices in product(range(n), repeat=r):
2619 if sorted(indices) == list(indices):
2620 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002621*/
2622typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002623 PyObject_HEAD
2624 PyObject *pool; /* input converted to a tuple */
2625 Py_ssize_t *indices; /* one index per result element */
2626 PyObject *result; /* most recently returned result tuple */
2627 Py_ssize_t r; /* size of result tuple */
2628 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002629} cwrobject;
2630
2631static PyTypeObject cwr_type;
2632
2633static PyObject *
2634cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2635{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002636 cwrobject *co;
2637 Py_ssize_t n;
2638 Py_ssize_t r;
2639 PyObject *pool = NULL;
2640 PyObject *iterable = NULL;
2641 Py_ssize_t *indices = NULL;
2642 Py_ssize_t i;
2643 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002644
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002645 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2646 &iterable, &r))
2647 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002648
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002649 pool = PySequence_Tuple(iterable);
2650 if (pool == NULL)
2651 goto error;
2652 n = PyTuple_GET_SIZE(pool);
2653 if (r < 0) {
2654 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2655 goto error;
2656 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002657
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002658 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2659 if (indices == NULL) {
2660 PyErr_NoMemory();
2661 goto error;
2662 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002663
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002664 for (i=0 ; i<r ; i++)
2665 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002666
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002667 /* create cwrobject structure */
2668 co = (cwrobject *)type->tp_alloc(type, 0);
2669 if (co == NULL)
2670 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002671
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002672 co->pool = pool;
2673 co->indices = indices;
2674 co->result = NULL;
2675 co->r = r;
2676 co->stopped = !n && r;
2677
2678 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002679
2680error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002681 if (indices != NULL)
2682 PyMem_Free(indices);
2683 Py_XDECREF(pool);
2684 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002685}
2686
2687static void
2688cwr_dealloc(cwrobject *co)
2689{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002690 PyObject_GC_UnTrack(co);
2691 Py_XDECREF(co->pool);
2692 Py_XDECREF(co->result);
2693 if (co->indices != NULL)
2694 PyMem_Free(co->indices);
2695 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002696}
2697
2698static int
2699cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2700{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002701 Py_VISIT(co->pool);
2702 Py_VISIT(co->result);
2703 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002704}
2705
2706static PyObject *
2707cwr_next(cwrobject *co)
2708{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002709 PyObject *elem;
2710 PyObject *oldelem;
2711 PyObject *pool = co->pool;
2712 Py_ssize_t *indices = co->indices;
2713 PyObject *result = co->result;
2714 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2715 Py_ssize_t r = co->r;
2716 Py_ssize_t i, j, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002717
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002718 if (co->stopped)
2719 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002720
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002721 if (result == NULL) {
2722 /* On the first pass, initialize result tuple using the indices */
2723 result = PyTuple_New(r);
2724 if (result == NULL)
2725 goto empty;
2726 co->result = result;
2727 for (i=0; i<r ; i++) {
2728 index = indices[i];
2729 elem = PyTuple_GET_ITEM(pool, index);
2730 Py_INCREF(elem);
2731 PyTuple_SET_ITEM(result, i, elem);
2732 }
2733 } else {
2734 /* Copy the previous result tuple or re-use it if available */
2735 if (Py_REFCNT(result) > 1) {
2736 PyObject *old_result = result;
2737 result = PyTuple_New(r);
2738 if (result == NULL)
2739 goto empty;
2740 co->result = result;
2741 for (i=0; i<r ; i++) {
2742 elem = PyTuple_GET_ITEM(old_result, i);
2743 Py_INCREF(elem);
2744 PyTuple_SET_ITEM(result, i, elem);
2745 }
2746 Py_DECREF(old_result);
2747 }
2748 /* Now, we've got the only copy so we can update it in-place CPython's
2749 empty tuple is a singleton and cached in PyTuple's freelist. */
2750 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002751
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002752 /* Scan indices right-to-left until finding one that is not
2753 * at its maximum (n-1). */
2754 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2755 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002756
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002757 /* If i is negative, then the indices are all at
2758 their maximum value and we're done. */
2759 if (i < 0)
2760 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002761
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002762 /* Increment the current index which we know is not at its
2763 maximum. Then set all to the right to the same value. */
2764 indices[i]++;
2765 for (j=i+1 ; j<r ; j++)
2766 indices[j] = indices[j-1];
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002767
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002768 /* Update the result tuple for the new indices
2769 starting with i, the leftmost index that changed */
2770 for ( ; i<r ; i++) {
2771 index = indices[i];
2772 elem = PyTuple_GET_ITEM(pool, index);
2773 Py_INCREF(elem);
2774 oldelem = PyTuple_GET_ITEM(result, i);
2775 PyTuple_SET_ITEM(result, i, elem);
2776 Py_DECREF(oldelem);
2777 }
2778 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002779
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002780 Py_INCREF(result);
2781 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002782
2783empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002784 co->stopped = 1;
2785 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002786}
2787
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002788static PyObject *
2789cwr_reduce(cwrobject *lz)
2790{
2791 if (lz->result == NULL) {
2792 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2793 } else if (lz->stopped) {
2794 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2795 } else {
2796 PyObject *indices;
2797 Py_ssize_t i;
2798
2799 /* we must pickle the indices and use them for setstate */
2800 indices = PyTuple_New(lz->r);
2801 if (!indices)
2802 return NULL;
2803 for (i=0; i<lz->r; i++)
2804 {
2805 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2806 if (!index) {
2807 Py_DECREF(indices);
2808 return NULL;
2809 }
2810 PyTuple_SET_ITEM(indices, i, index);
2811 }
2812
2813 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2814 }
2815}
2816
2817static PyObject *
2818cwr_setstate(cwrobject *lz, PyObject *state)
2819{
2820 PyObject *result;
2821 Py_ssize_t n, i;
2822
2823 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2824 {
2825 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2826 return NULL;
2827 }
2828
2829 n = PyTuple_GET_SIZE(lz->pool);
2830 for (i=0; i<lz->r; i++)
2831 {
2832 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2833 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2834 if (index < 0 && PyErr_Occurred())
2835 return NULL; /* not an integer */
2836 /* clamp the index */
2837 if (index < 0)
2838 index = 0;
2839 else if (index > n-1)
2840 index = n-1;
2841 lz->indices[i] = index;
2842 }
2843 result = PyTuple_New(lz->r);
2844 if (result == NULL)
2845 return NULL;
2846 for (i=0; i<lz->r; i++) {
2847 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2848 Py_INCREF(element);
2849 PyTuple_SET_ITEM(result, i, element);
2850 }
2851 Py_CLEAR(lz->result);
2852 lz->result = result;
2853 Py_RETURN_NONE;
2854}
2855
2856static PyMethodDef cwr_methods[] = {
2857 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2858 reduce_doc},
2859 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2860 setstate_doc},
2861 {NULL, NULL} /* sentinel */
2862};
2863
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002864PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002865"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002866\n\
2867Return successive r-length combinations of elements in the iterable\n\
2868allowing individual elements to have successive repeats.\n\
2869combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2870
2871static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002872 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002873 "itertools.combinations_with_replacement", /* tp_name */
2874 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002875 0, /* tp_itemsize */
2876 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002877 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002878 0, /* tp_print */
2879 0, /* tp_getattr */
2880 0, /* tp_setattr */
2881 0, /* tp_reserved */
2882 0, /* tp_repr */
2883 0, /* tp_as_number */
2884 0, /* tp_as_sequence */
2885 0, /* tp_as_mapping */
2886 0, /* tp_hash */
2887 0, /* tp_call */
2888 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002889 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002890 0, /* tp_setattro */
2891 0, /* tp_as_buffer */
2892 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002893 Py_TPFLAGS_BASETYPE, /* tp_flags */
2894 cwr_doc, /* tp_doc */
2895 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002896 0, /* tp_clear */
2897 0, /* tp_richcompare */
2898 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002899 PyObject_SelfIter, /* tp_iter */
2900 (iternextfunc)cwr_next, /* tp_iternext */
2901 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002902 0, /* tp_members */
2903 0, /* tp_getset */
2904 0, /* tp_base */
2905 0, /* tp_dict */
2906 0, /* tp_descr_get */
2907 0, /* tp_descr_set */
2908 0, /* tp_dictoffset */
2909 0, /* tp_init */
2910 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002911 cwr_new, /* tp_new */
2912 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002913};
2914
2915
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002916/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002917
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002918def permutations(iterable, r=None):
2919 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2920 pool = tuple(iterable)
2921 n = len(pool)
2922 r = n if r is None else r
2923 indices = range(n)
2924 cycles = range(n-r+1, n+1)[::-1]
2925 yield tuple(pool[i] for i in indices[:r])
2926 while n:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002927 for i in reversed(range(r)):
2928 cycles[i] -= 1
2929 if cycles[i] == 0:
2930 indices[i:] = indices[i+1:] + indices[i:i+1]
2931 cycles[i] = n - i
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002932 else:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002933 j = cycles[i]
2934 indices[i], indices[-j] = indices[-j], indices[i]
2935 yield tuple(pool[i] for i in indices[:r])
2936 break
2937 else:
2938 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002939*/
2940
2941typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002942 PyObject_HEAD
2943 PyObject *pool; /* input converted to a tuple */
2944 Py_ssize_t *indices; /* one index per element in the pool */
2945 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2946 PyObject *result; /* most recently returned result tuple */
2947 Py_ssize_t r; /* size of result tuple */
2948 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002949} permutationsobject;
2950
2951static PyTypeObject permutations_type;
2952
2953static PyObject *
2954permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2955{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002956 permutationsobject *po;
2957 Py_ssize_t n;
2958 Py_ssize_t r;
2959 PyObject *robj = Py_None;
2960 PyObject *pool = NULL;
2961 PyObject *iterable = NULL;
2962 Py_ssize_t *indices = NULL;
2963 Py_ssize_t *cycles = NULL;
2964 Py_ssize_t i;
2965 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002966
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002967 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2968 &iterable, &robj))
2969 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002970
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002971 pool = PySequence_Tuple(iterable);
2972 if (pool == NULL)
2973 goto error;
2974 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002975
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002976 r = n;
2977 if (robj != Py_None) {
2978 if (!PyLong_Check(robj)) {
2979 PyErr_SetString(PyExc_TypeError, "Expected int as r");
2980 goto error;
2981 }
2982 r = PyLong_AsSsize_t(robj);
2983 if (r == -1 && PyErr_Occurred())
2984 goto error;
2985 }
2986 if (r < 0) {
2987 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2988 goto error;
2989 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002990
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002991 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2992 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2993 if (indices == NULL || cycles == NULL) {
2994 PyErr_NoMemory();
2995 goto error;
2996 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002997
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002998 for (i=0 ; i<n ; i++)
2999 indices[i] = i;
3000 for (i=0 ; i<r ; i++)
3001 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003002
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003003 /* create permutationsobject structure */
3004 po = (permutationsobject *)type->tp_alloc(type, 0);
3005 if (po == NULL)
3006 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003007
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003008 po->pool = pool;
3009 po->indices = indices;
3010 po->cycles = cycles;
3011 po->result = NULL;
3012 po->r = r;
3013 po->stopped = r > n ? 1 : 0;
3014
3015 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003016
3017error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003018 if (indices != NULL)
3019 PyMem_Free(indices);
3020 if (cycles != NULL)
3021 PyMem_Free(cycles);
3022 Py_XDECREF(pool);
3023 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003024}
3025
3026static void
3027permutations_dealloc(permutationsobject *po)
3028{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003029 PyObject_GC_UnTrack(po);
3030 Py_XDECREF(po->pool);
3031 Py_XDECREF(po->result);
3032 PyMem_Free(po->indices);
3033 PyMem_Free(po->cycles);
3034 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003035}
3036
3037static int
3038permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3039{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003040 Py_VISIT(po->pool);
3041 Py_VISIT(po->result);
3042 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003043}
3044
3045static PyObject *
3046permutations_next(permutationsobject *po)
3047{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003048 PyObject *elem;
3049 PyObject *oldelem;
3050 PyObject *pool = po->pool;
3051 Py_ssize_t *indices = po->indices;
3052 Py_ssize_t *cycles = po->cycles;
3053 PyObject *result = po->result;
3054 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3055 Py_ssize_t r = po->r;
3056 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003057
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003058 if (po->stopped)
3059 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003060
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003061 if (result == NULL) {
3062 /* On the first pass, initialize result tuple using the indices */
3063 result = PyTuple_New(r);
3064 if (result == NULL)
3065 goto empty;
3066 po->result = result;
3067 for (i=0; i<r ; i++) {
3068 index = indices[i];
3069 elem = PyTuple_GET_ITEM(pool, index);
3070 Py_INCREF(elem);
3071 PyTuple_SET_ITEM(result, i, elem);
3072 }
3073 } else {
3074 if (n == 0)
3075 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003076
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003077 /* Copy the previous result tuple or re-use it if available */
3078 if (Py_REFCNT(result) > 1) {
3079 PyObject *old_result = result;
3080 result = PyTuple_New(r);
3081 if (result == NULL)
3082 goto empty;
3083 po->result = result;
3084 for (i=0; i<r ; i++) {
3085 elem = PyTuple_GET_ITEM(old_result, i);
3086 Py_INCREF(elem);
3087 PyTuple_SET_ITEM(result, i, elem);
3088 }
3089 Py_DECREF(old_result);
3090 }
3091 /* Now, we've got the only copy so we can update it in-place */
3092 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003093
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003094 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3095 for (i=r-1 ; i>=0 ; i--) {
3096 cycles[i] -= 1;
3097 if (cycles[i] == 0) {
3098 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3099 index = indices[i];
3100 for (j=i ; j<n-1 ; j++)
3101 indices[j] = indices[j+1];
3102 indices[n-1] = index;
3103 cycles[i] = n - i;
3104 } else {
3105 j = cycles[i];
3106 index = indices[i];
3107 indices[i] = indices[n-j];
3108 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003109
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003110 for (k=i; k<r ; k++) {
3111 /* start with i, the leftmost element that changed */
3112 /* yield tuple(pool[k] for k in indices[:r]) */
3113 index = indices[k];
3114 elem = PyTuple_GET_ITEM(pool, index);
3115 Py_INCREF(elem);
3116 oldelem = PyTuple_GET_ITEM(result, k);
3117 PyTuple_SET_ITEM(result, k, elem);
3118 Py_DECREF(oldelem);
3119 }
3120 break;
3121 }
3122 }
3123 /* If i is negative, then the cycles have all
3124 rolled-over and we're done. */
3125 if (i < 0)
3126 goto empty;
3127 }
3128 Py_INCREF(result);
3129 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003130
3131empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003132 po->stopped = 1;
3133 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003134}
3135
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003136static PyObject *
3137permutations_reduce(permutationsobject *po)
3138{
3139 if (po->result == NULL) {
3140 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3141 } else if (po->stopped) {
3142 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3143 } else {
3144 PyObject *indices=NULL, *cycles=NULL;
3145 Py_ssize_t n, i;
3146
3147 /* we must pickle the indices and cycles and use them for setstate */
3148 n = PyTuple_GET_SIZE(po->pool);
3149 indices = PyTuple_New(n);
3150 if (indices == NULL)
3151 goto err;
3152 for (i=0; i<n; i++){
3153 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3154 if (!index)
3155 goto err;
3156 PyTuple_SET_ITEM(indices, i, index);
3157 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003158
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003159 cycles = PyTuple_New(po->r);
3160 if (cycles == NULL)
3161 goto err;
3162 for (i=0; i<po->r; i++)
3163 {
3164 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3165 if (!index)
3166 goto err;
3167 PyTuple_SET_ITEM(cycles, i, index);
3168 }
3169 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3170 po->pool, po->r,
3171 indices, cycles);
3172 err:
3173 Py_XDECREF(indices);
3174 Py_XDECREF(cycles);
3175 return NULL;
3176 }
3177}
3178
3179static PyObject *
3180permutations_setstate(permutationsobject *po, PyObject *state)
3181{
3182 PyObject *indices, *cycles, *result;
3183 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003184
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003185 if (!PyArg_ParseTuple(state, "O!O!",
3186 &PyTuple_Type, &indices,
3187 &PyTuple_Type, &cycles))
3188 return NULL;
3189
3190 n = PyTuple_GET_SIZE(po->pool);
3191 if (PyTuple_GET_SIZE(indices) != n ||
3192 PyTuple_GET_SIZE(cycles) != po->r)
3193 {
3194 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3195 return NULL;
3196 }
3197
3198 for (i=0; i<n; i++)
3199 {
3200 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3201 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3202 if (index < 0 && PyErr_Occurred())
3203 return NULL; /* not an integer */
3204 /* clamp the index */
3205 if (index < 0)
3206 index = 0;
3207 else if (index > n-1)
3208 index = n-1;
3209 po->indices[i] = index;
3210 }
3211
3212 for (i=0; i<po->r; i++)
3213 {
3214 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3215 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3216 if (index < 0 && PyErr_Occurred())
3217 return NULL; /* not an integer */
3218 if (index < 1)
3219 index = 1;
3220 else if (index > n-i)
3221 index = n-i;
3222 po->cycles[i] = index;
3223 }
3224 result = PyTuple_New(po->r);
3225 if (result == NULL)
3226 return NULL;
3227 for (i=0; i<po->r; i++) {
3228 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3229 Py_INCREF(element);
3230 PyTuple_SET_ITEM(result, i, element);
3231 }
3232 Py_CLEAR(po->result);
3233 po->result = result;
3234 Py_RETURN_NONE;
3235}
3236
3237static PyMethodDef permuations_methods[] = {
3238 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3239 reduce_doc},
3240 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3241 setstate_doc},
3242 {NULL, NULL} /* sentinel */
3243};
3244
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003245PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003246"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003247\n\
3248Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003249permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003250
3251static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003252 PyVarObject_HEAD_INIT(NULL, 0)
3253 "itertools.permutations", /* tp_name */
3254 sizeof(permutationsobject), /* tp_basicsize */
3255 0, /* tp_itemsize */
3256 /* methods */
3257 (destructor)permutations_dealloc, /* tp_dealloc */
3258 0, /* tp_print */
3259 0, /* tp_getattr */
3260 0, /* tp_setattr */
3261 0, /* tp_reserved */
3262 0, /* tp_repr */
3263 0, /* tp_as_number */
3264 0, /* tp_as_sequence */
3265 0, /* tp_as_mapping */
3266 0, /* tp_hash */
3267 0, /* tp_call */
3268 0, /* tp_str */
3269 PyObject_GenericGetAttr, /* tp_getattro */
3270 0, /* tp_setattro */
3271 0, /* tp_as_buffer */
3272 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3273 Py_TPFLAGS_BASETYPE, /* tp_flags */
3274 permutations_doc, /* tp_doc */
3275 (traverseproc)permutations_traverse, /* tp_traverse */
3276 0, /* tp_clear */
3277 0, /* tp_richcompare */
3278 0, /* tp_weaklistoffset */
3279 PyObject_SelfIter, /* tp_iter */
3280 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003281 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003282 0, /* tp_members */
3283 0, /* tp_getset */
3284 0, /* tp_base */
3285 0, /* tp_dict */
3286 0, /* tp_descr_get */
3287 0, /* tp_descr_set */
3288 0, /* tp_dictoffset */
3289 0, /* tp_init */
3290 0, /* tp_alloc */
3291 permutations_new, /* tp_new */
3292 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003293};
3294
Raymond Hettinger482ba772010-12-01 22:48:00 +00003295/* accumulate object ************************************************************/
3296
3297typedef struct {
3298 PyObject_HEAD
3299 PyObject *total;
3300 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003301 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003302} accumulateobject;
3303
3304static PyTypeObject accumulate_type;
3305
3306static PyObject *
3307accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3308{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003309 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003310 PyObject *iterable;
3311 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003312 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003313 accumulateobject *lz;
3314
Raymond Hettinger5d446132011-03-27 18:52:10 -07003315 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3316 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003317 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003318
3319 /* Get iterator. */
3320 it = PyObject_GetIter(iterable);
3321 if (it == NULL)
3322 return NULL;
3323
Raymond Hettinger482ba772010-12-01 22:48:00 +00003324 /* create accumulateobject structure */
3325 lz = (accumulateobject *)type->tp_alloc(type, 0);
3326 if (lz == NULL) {
3327 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003328 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003329 }
3330
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003331 if (binop != Py_None) {
3332 Py_XINCREF(binop);
3333 lz->binop = binop;
3334 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003335 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003336 lz->it = it;
3337 return (PyObject *)lz;
3338}
3339
3340static void
3341accumulate_dealloc(accumulateobject *lz)
3342{
3343 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003344 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003345 Py_XDECREF(lz->total);
3346 Py_XDECREF(lz->it);
3347 Py_TYPE(lz)->tp_free(lz);
3348}
3349
3350static int
3351accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3352{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003353 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003354 Py_VISIT(lz->it);
3355 Py_VISIT(lz->total);
3356 return 0;
3357}
3358
3359static PyObject *
3360accumulate_next(accumulateobject *lz)
3361{
3362 PyObject *val, *oldtotal, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003363
Raymond Hettinger482ba772010-12-01 22:48:00 +00003364 val = PyIter_Next(lz->it);
3365 if (val == NULL)
3366 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003367
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003368 if (lz->total == NULL) {
3369 Py_INCREF(val);
3370 lz->total = val;
3371 return lz->total;
3372 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003373
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003374 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003375 newtotal = PyNumber_Add(lz->total, val);
3376 else
3377 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003378 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003379 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003380 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003381
3382 oldtotal = lz->total;
3383 lz->total = newtotal;
3384 Py_DECREF(oldtotal);
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003385
Raymond Hettinger482ba772010-12-01 22:48:00 +00003386 Py_INCREF(newtotal);
3387 return newtotal;
3388}
3389
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003390static PyObject *
3391accumulate_reduce(accumulateobject *lz)
3392{
3393 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3394 lz->it, lz->binop?lz->binop:Py_None,
3395 lz->total?lz->total:Py_None);
3396 }
3397
3398static PyObject *
3399accumulate_setstate(accumulateobject *lz, PyObject *state)
3400{
3401 Py_CLEAR(lz->total);
3402 lz->total = state;
3403 Py_INCREF(lz->total);
3404 Py_RETURN_NONE;
3405}
3406
3407static PyMethodDef accumulate_methods[] = {
3408 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3409 reduce_doc},
3410 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3411 setstate_doc},
3412 {NULL, NULL} /* sentinel */
3413};
3414
Raymond Hettinger482ba772010-12-01 22:48:00 +00003415PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003416"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003417\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003418Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003419
3420static PyTypeObject accumulate_type = {
3421 PyVarObject_HEAD_INIT(NULL, 0)
3422 "itertools.accumulate", /* tp_name */
3423 sizeof(accumulateobject), /* tp_basicsize */
3424 0, /* tp_itemsize */
3425 /* methods */
3426 (destructor)accumulate_dealloc, /* tp_dealloc */
3427 0, /* tp_print */
3428 0, /* tp_getattr */
3429 0, /* tp_setattr */
3430 0, /* tp_reserved */
3431 0, /* tp_repr */
3432 0, /* tp_as_number */
3433 0, /* tp_as_sequence */
3434 0, /* tp_as_mapping */
3435 0, /* tp_hash */
3436 0, /* tp_call */
3437 0, /* tp_str */
3438 PyObject_GenericGetAttr, /* tp_getattro */
3439 0, /* tp_setattro */
3440 0, /* tp_as_buffer */
3441 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3442 Py_TPFLAGS_BASETYPE, /* tp_flags */
3443 accumulate_doc, /* tp_doc */
3444 (traverseproc)accumulate_traverse, /* tp_traverse */
3445 0, /* tp_clear */
3446 0, /* tp_richcompare */
3447 0, /* tp_weaklistoffset */
3448 PyObject_SelfIter, /* tp_iter */
3449 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003450 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003451 0, /* tp_members */
3452 0, /* tp_getset */
3453 0, /* tp_base */
3454 0, /* tp_dict */
3455 0, /* tp_descr_get */
3456 0, /* tp_descr_set */
3457 0, /* tp_dictoffset */
3458 0, /* tp_init */
3459 0, /* tp_alloc */
3460 accumulate_new, /* tp_new */
3461 PyObject_GC_Del, /* tp_free */
3462};
3463
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003464
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003465/* compress object ************************************************************/
3466
3467/* Equivalent to:
3468
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003469 def compress(data, selectors):
3470 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3471 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003472*/
3473
3474typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003475 PyObject_HEAD
3476 PyObject *data;
3477 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003478} compressobject;
3479
3480static PyTypeObject compress_type;
3481
3482static PyObject *
3483compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3484{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003485 PyObject *seq1, *seq2;
3486 PyObject *data=NULL, *selectors=NULL;
3487 compressobject *lz;
3488 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003489
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003490 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3491 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003492
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003493 data = PyObject_GetIter(seq1);
3494 if (data == NULL)
3495 goto fail;
3496 selectors = PyObject_GetIter(seq2);
3497 if (selectors == NULL)
3498 goto fail;
3499
3500 /* create compressobject structure */
3501 lz = (compressobject *)type->tp_alloc(type, 0);
3502 if (lz == NULL)
3503 goto fail;
3504 lz->data = data;
3505 lz->selectors = selectors;
3506 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003507
3508fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003509 Py_XDECREF(data);
3510 Py_XDECREF(selectors);
3511 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003512}
3513
3514static void
3515compress_dealloc(compressobject *lz)
3516{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003517 PyObject_GC_UnTrack(lz);
3518 Py_XDECREF(lz->data);
3519 Py_XDECREF(lz->selectors);
3520 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003521}
3522
3523static int
3524compress_traverse(compressobject *lz, visitproc visit, void *arg)
3525{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003526 Py_VISIT(lz->data);
3527 Py_VISIT(lz->selectors);
3528 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003529}
3530
3531static PyObject *
3532compress_next(compressobject *lz)
3533{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003534 PyObject *data = lz->data, *selectors = lz->selectors;
3535 PyObject *datum, *selector;
3536 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3537 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3538 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003539
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003540 while (1) {
3541 /* Steps: get datum, get selector, evaluate selector.
3542 Order is important (to match the pure python version
3543 in terms of which input gets a chance to raise an
3544 exception first).
3545 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003546
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003547 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003548 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003549 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003550
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003551 selector = selectornext(selectors);
3552 if (selector == NULL) {
3553 Py_DECREF(datum);
3554 return NULL;
3555 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003556
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003557 ok = PyObject_IsTrue(selector);
3558 Py_DECREF(selector);
3559 if (ok == 1)
3560 return datum;
3561 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003562 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003563 return NULL;
3564 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003565}
3566
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003567static PyObject *
3568compress_reduce(compressobject *lz)
3569{
3570 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3571 lz->data, lz->selectors);
3572 }
3573
3574static PyMethodDef compress_methods[] = {
3575 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3576 reduce_doc},
3577 {NULL, NULL} /* sentinel */
3578};
3579
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003580PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003581"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003582\n\
3583Return data elements corresponding to true selector elements.\n\
3584Forms a shorter iterator from selected data elements using the\n\
3585selectors to choose the data elements.");
3586
3587static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003588 PyVarObject_HEAD_INIT(NULL, 0)
3589 "itertools.compress", /* tp_name */
3590 sizeof(compressobject), /* tp_basicsize */
3591 0, /* tp_itemsize */
3592 /* methods */
3593 (destructor)compress_dealloc, /* tp_dealloc */
3594 0, /* tp_print */
3595 0, /* tp_getattr */
3596 0, /* tp_setattr */
3597 0, /* tp_reserved */
3598 0, /* tp_repr */
3599 0, /* tp_as_number */
3600 0, /* tp_as_sequence */
3601 0, /* tp_as_mapping */
3602 0, /* tp_hash */
3603 0, /* tp_call */
3604 0, /* tp_str */
3605 PyObject_GenericGetAttr, /* tp_getattro */
3606 0, /* tp_setattro */
3607 0, /* tp_as_buffer */
3608 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3609 Py_TPFLAGS_BASETYPE, /* tp_flags */
3610 compress_doc, /* tp_doc */
3611 (traverseproc)compress_traverse, /* tp_traverse */
3612 0, /* tp_clear */
3613 0, /* tp_richcompare */
3614 0, /* tp_weaklistoffset */
3615 PyObject_SelfIter, /* tp_iter */
3616 (iternextfunc)compress_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003617 compress_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003618 0, /* tp_members */
3619 0, /* tp_getset */
3620 0, /* tp_base */
3621 0, /* tp_dict */
3622 0, /* tp_descr_get */
3623 0, /* tp_descr_set */
3624 0, /* tp_dictoffset */
3625 0, /* tp_init */
3626 0, /* tp_alloc */
3627 compress_new, /* tp_new */
3628 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003629};
3630
3631
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003632/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003633
3634typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003635 PyObject_HEAD
3636 PyObject *func;
3637 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003638} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003639
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003640static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003641
3642static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003643filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003644{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003645 PyObject *func, *seq;
3646 PyObject *it;
3647 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003648
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003649 if (type == &filterfalse_type &&
3650 !_PyArg_NoKeywords("filterfalse()", kwds))
3651 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003652
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003653 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3654 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003655
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003656 /* Get iterator. */
3657 it = PyObject_GetIter(seq);
3658 if (it == NULL)
3659 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003660
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003661 /* create filterfalseobject structure */
3662 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3663 if (lz == NULL) {
3664 Py_DECREF(it);
3665 return NULL;
3666 }
3667 Py_INCREF(func);
3668 lz->func = func;
3669 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003670
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003671 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003672}
3673
3674static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003675filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003676{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003677 PyObject_GC_UnTrack(lz);
3678 Py_XDECREF(lz->func);
3679 Py_XDECREF(lz->it);
3680 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003681}
3682
3683static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003684filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003685{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003686 Py_VISIT(lz->it);
3687 Py_VISIT(lz->func);
3688 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003689}
3690
3691static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003692filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003693{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003694 PyObject *item;
3695 PyObject *it = lz->it;
3696 long ok;
3697 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003698
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003699 iternext = *Py_TYPE(it)->tp_iternext;
3700 for (;;) {
3701 item = iternext(it);
3702 if (item == NULL)
3703 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003704
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003705 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3706 ok = PyObject_IsTrue(item);
3707 } else {
3708 PyObject *good;
3709 good = PyObject_CallFunctionObjArgs(lz->func,
3710 item, NULL);
3711 if (good == NULL) {
3712 Py_DECREF(item);
3713 return NULL;
3714 }
3715 ok = PyObject_IsTrue(good);
3716 Py_DECREF(good);
3717 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003718 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003719 return item;
3720 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003721 if (ok < 0)
3722 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003723 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003724}
3725
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003726static PyObject *
3727filterfalse_reduce(filterfalseobject *lz)
3728{
3729 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3730 lz->func, lz->it);
3731 }
3732
3733static PyMethodDef filterfalse_methods[] = {
3734 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3735 reduce_doc},
3736 {NULL, NULL} /* sentinel */
3737};
3738
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003739PyDoc_STRVAR(filterfalse_doc,
3740"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003741\n\
3742Return those items of sequence for which function(item) is false.\n\
3743If function is None, return the items that are false.");
3744
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003745static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003746 PyVarObject_HEAD_INIT(NULL, 0)
3747 "itertools.filterfalse", /* tp_name */
3748 sizeof(filterfalseobject), /* tp_basicsize */
3749 0, /* tp_itemsize */
3750 /* methods */
3751 (destructor)filterfalse_dealloc, /* tp_dealloc */
3752 0, /* tp_print */
3753 0, /* tp_getattr */
3754 0, /* tp_setattr */
3755 0, /* tp_reserved */
3756 0, /* tp_repr */
3757 0, /* tp_as_number */
3758 0, /* tp_as_sequence */
3759 0, /* tp_as_mapping */
3760 0, /* tp_hash */
3761 0, /* tp_call */
3762 0, /* tp_str */
3763 PyObject_GenericGetAttr, /* tp_getattro */
3764 0, /* tp_setattro */
3765 0, /* tp_as_buffer */
3766 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3767 Py_TPFLAGS_BASETYPE, /* tp_flags */
3768 filterfalse_doc, /* tp_doc */
3769 (traverseproc)filterfalse_traverse, /* tp_traverse */
3770 0, /* tp_clear */
3771 0, /* tp_richcompare */
3772 0, /* tp_weaklistoffset */
3773 PyObject_SelfIter, /* tp_iter */
3774 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003775 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003776 0, /* tp_members */
3777 0, /* tp_getset */
3778 0, /* tp_base */
3779 0, /* tp_dict */
3780 0, /* tp_descr_get */
3781 0, /* tp_descr_set */
3782 0, /* tp_dictoffset */
3783 0, /* tp_init */
3784 0, /* tp_alloc */
3785 filterfalse_new, /* tp_new */
3786 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003787};
3788
3789
3790/* count object ************************************************************/
3791
3792typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003793 PyObject_HEAD
3794 Py_ssize_t cnt;
3795 PyObject *long_cnt;
3796 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003797} countobject;
3798
Raymond Hettinger30729212009-02-12 06:28:27 +00003799/* Counting logic and invariants:
3800
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003801fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003802
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003803 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3804 Advances with: cnt += 1
3805 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003806
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003807slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003808
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003809 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3810 All counting is done with python objects (no overflows or underflows).
3811 Advances with: long_cnt += long_step
3812 Step may be zero -- effectively a slow version of repeat(cnt).
3813 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003814*/
3815
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003816static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003817
3818static PyObject *
3819count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3820{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003821 countobject *lz;
3822 int slow_mode = 0;
3823 Py_ssize_t cnt = 0;
3824 PyObject *long_cnt = NULL;
3825 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003826 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003827 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003828
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003829 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3830 kwlist, &long_cnt, &long_step))
3831 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003832
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003833 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3834 (long_step != NULL && !PyNumber_Check(long_step))) {
3835 PyErr_SetString(PyExc_TypeError, "a number is required");
3836 return NULL;
3837 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003838
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003839 if (long_cnt != NULL) {
3840 cnt = PyLong_AsSsize_t(long_cnt);
3841 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3842 PyErr_Clear();
3843 slow_mode = 1;
3844 }
3845 Py_INCREF(long_cnt);
3846 } else {
3847 cnt = 0;
3848 long_cnt = PyLong_FromLong(0);
3849 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003850
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003851 /* If not specified, step defaults to 1 */
3852 if (long_step == NULL) {
3853 long_step = PyLong_FromLong(1);
3854 if (long_step == NULL) {
3855 Py_DECREF(long_cnt);
3856 return NULL;
3857 }
3858 } else
3859 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003860
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003861 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003862
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003863 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003864 step = PyLong_AsLong(long_step);
3865 if (step != 1) {
3866 slow_mode = 1;
3867 if (step == -1 && PyErr_Occurred())
3868 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003869 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003870
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003871 if (slow_mode)
3872 cnt = PY_SSIZE_T_MAX;
3873 else
3874 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003875
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003876 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3877 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3878 assert(slow_mode ||
3879 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003880
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003881 /* create countobject structure */
3882 lz = (countobject *)type->tp_alloc(type, 0);
3883 if (lz == NULL) {
3884 Py_XDECREF(long_cnt);
3885 return NULL;
3886 }
3887 lz->cnt = cnt;
3888 lz->long_cnt = long_cnt;
3889 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003890
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003891 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003892}
3893
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003894static void
3895count_dealloc(countobject *lz)
3896{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003897 PyObject_GC_UnTrack(lz);
3898 Py_XDECREF(lz->long_cnt);
3899 Py_XDECREF(lz->long_step);
3900 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003901}
3902
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003903static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003904count_traverse(countobject *lz, visitproc visit, void *arg)
3905{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003906 Py_VISIT(lz->long_cnt);
3907 Py_VISIT(lz->long_step);
3908 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003909}
3910
3911static PyObject *
3912count_nextlong(countobject *lz)
3913{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003914 PyObject *long_cnt;
3915 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003916
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003917 long_cnt = lz->long_cnt;
3918 if (long_cnt == NULL) {
3919 /* Switch to slow_mode */
3920 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3921 if (long_cnt == NULL)
3922 return NULL;
3923 }
3924 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00003925
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003926 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3927 if (stepped_up == NULL)
3928 return NULL;
3929 lz->long_cnt = stepped_up;
3930 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003931}
3932
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003933static PyObject *
3934count_next(countobject *lz)
3935{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003936 if (lz->cnt == PY_SSIZE_T_MAX)
3937 return count_nextlong(lz);
3938 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003939}
3940
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003941static PyObject *
3942count_repr(countobject *lz)
3943{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003944 if (lz->cnt != PY_SSIZE_T_MAX)
3945 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003946
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003947 if (PyLong_Check(lz->long_step)) {
3948 long step = PyLong_AsLong(lz->long_step);
3949 if (step == -1 && PyErr_Occurred()) {
3950 PyErr_Clear();
3951 }
3952 if (step == 1) {
3953 /* Don't display step when it is an integer equal to 1 */
3954 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
3955 }
3956 }
3957 return PyUnicode_FromFormat("count(%R, %R)",
3958 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003959}
3960
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003961static PyObject *
3962count_reduce(countobject *lz)
3963{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003964 if (lz->cnt == PY_SSIZE_T_MAX)
3965 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3966 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003967}
3968
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003969static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003970 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003971 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003972 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003973};
3974
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003975PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00003976 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003977\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003978Return a count object whose .__next__() method returns consecutive values.\n\
3979Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00003980 def count(firstval=0, step=1):\n\
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003981 x = firstval\n\
3982 while 1:\n\
3983 yield x\n\
3984 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003985
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003986static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003987 PyVarObject_HEAD_INIT(NULL, 0)
3988 "itertools.count", /* tp_name */
3989 sizeof(countobject), /* tp_basicsize */
3990 0, /* tp_itemsize */
3991 /* methods */
3992 (destructor)count_dealloc, /* tp_dealloc */
3993 0, /* tp_print */
3994 0, /* tp_getattr */
3995 0, /* tp_setattr */
3996 0, /* tp_reserved */
3997 (reprfunc)count_repr, /* tp_repr */
3998 0, /* tp_as_number */
3999 0, /* tp_as_sequence */
4000 0, /* tp_as_mapping */
4001 0, /* tp_hash */
4002 0, /* tp_call */
4003 0, /* tp_str */
4004 PyObject_GenericGetAttr, /* tp_getattro */
4005 0, /* tp_setattro */
4006 0, /* tp_as_buffer */
4007 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4008 Py_TPFLAGS_BASETYPE, /* tp_flags */
4009 count_doc, /* tp_doc */
4010 (traverseproc)count_traverse, /* tp_traverse */
4011 0, /* tp_clear */
4012 0, /* tp_richcompare */
4013 0, /* tp_weaklistoffset */
4014 PyObject_SelfIter, /* tp_iter */
4015 (iternextfunc)count_next, /* tp_iternext */
4016 count_methods, /* tp_methods */
4017 0, /* tp_members */
4018 0, /* tp_getset */
4019 0, /* tp_base */
4020 0, /* tp_dict */
4021 0, /* tp_descr_get */
4022 0, /* tp_descr_set */
4023 0, /* tp_dictoffset */
4024 0, /* tp_init */
4025 0, /* tp_alloc */
4026 count_new, /* tp_new */
4027 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004028};
4029
4030
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004031/* repeat object ************************************************************/
4032
4033typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004034 PyObject_HEAD
4035 PyObject *element;
4036 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004037} repeatobject;
4038
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004039static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004040
4041static PyObject *
4042repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4043{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004044 repeatobject *ro;
4045 PyObject *element;
4046 Py_ssize_t cnt = -1;
4047 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004048
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004049 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4050 &element, &cnt))
4051 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004052
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004053 if (PyTuple_Size(args) == 2 && cnt < 0)
4054 cnt = 0;
4055
4056 ro = (repeatobject *)type->tp_alloc(type, 0);
4057 if (ro == NULL)
4058 return NULL;
4059 Py_INCREF(element);
4060 ro->element = element;
4061 ro->cnt = cnt;
4062 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004063}
4064
4065static void
4066repeat_dealloc(repeatobject *ro)
4067{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004068 PyObject_GC_UnTrack(ro);
4069 Py_XDECREF(ro->element);
4070 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004071}
4072
4073static int
4074repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4075{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004076 Py_VISIT(ro->element);
4077 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004078}
4079
4080static PyObject *
4081repeat_next(repeatobject *ro)
4082{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004083 if (ro->cnt == 0)
4084 return NULL;
4085 if (ro->cnt > 0)
4086 ro->cnt--;
4087 Py_INCREF(ro->element);
4088 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004089}
4090
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004091static PyObject *
4092repeat_repr(repeatobject *ro)
4093{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004094 if (ro->cnt == -1)
4095 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4096 else
4097 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4098}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004099
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004100static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004101repeat_len(repeatobject *ro)
4102{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004103 if (ro->cnt == -1) {
4104 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4105 return NULL;
4106 }
4107 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004108}
4109
Armin Rigof5b3e362006-02-11 21:32:43 +00004110PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004111
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004112static PyObject *
4113repeat_reduce(repeatobject *ro)
4114{
4115 /* unpickle this so that a new repeat iterator is constructed with an
4116 * object, then call __setstate__ on it to set cnt
4117 */
4118 if (ro->cnt >= 0)
4119 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4120 else
4121 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4122}
4123
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004124static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004125 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004126 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004127 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004128};
4129
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004130PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004131"repeat(object [,times]) -> create an iterator which returns the object\n\
4132for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004133endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004134
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004135static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004136 PyVarObject_HEAD_INIT(NULL, 0)
4137 "itertools.repeat", /* tp_name */
4138 sizeof(repeatobject), /* tp_basicsize */
4139 0, /* tp_itemsize */
4140 /* methods */
4141 (destructor)repeat_dealloc, /* tp_dealloc */
4142 0, /* tp_print */
4143 0, /* tp_getattr */
4144 0, /* tp_setattr */
4145 0, /* tp_reserved */
4146 (reprfunc)repeat_repr, /* tp_repr */
4147 0, /* tp_as_number */
4148 0, /* tp_as_sequence */
4149 0, /* tp_as_mapping */
4150 0, /* tp_hash */
4151 0, /* tp_call */
4152 0, /* tp_str */
4153 PyObject_GenericGetAttr, /* tp_getattro */
4154 0, /* tp_setattro */
4155 0, /* tp_as_buffer */
4156 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4157 Py_TPFLAGS_BASETYPE, /* tp_flags */
4158 repeat_doc, /* tp_doc */
4159 (traverseproc)repeat_traverse, /* tp_traverse */
4160 0, /* tp_clear */
4161 0, /* tp_richcompare */
4162 0, /* tp_weaklistoffset */
4163 PyObject_SelfIter, /* tp_iter */
4164 (iternextfunc)repeat_next, /* tp_iternext */
4165 repeat_methods, /* tp_methods */
4166 0, /* tp_members */
4167 0, /* tp_getset */
4168 0, /* tp_base */
4169 0, /* tp_dict */
4170 0, /* tp_descr_get */
4171 0, /* tp_descr_set */
4172 0, /* tp_dictoffset */
4173 0, /* tp_init */
4174 0, /* tp_alloc */
4175 repeat_new, /* tp_new */
4176 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004177};
4178
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004179/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004180
4181#include "Python.h"
4182
4183typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004184 PyObject_HEAD
4185 Py_ssize_t tuplesize;
4186 Py_ssize_t numactive;
4187 PyObject *ittuple; /* tuple of iterators */
4188 PyObject *result;
4189 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004190} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004191
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004192static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004193
4194static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004195zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004196{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004197 ziplongestobject *lz;
4198 Py_ssize_t i;
4199 PyObject *ittuple; /* tuple of iterators */
4200 PyObject *result;
4201 PyObject *fillvalue = Py_None;
4202 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004203
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004204 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
4205 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4206 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
4207 PyErr_SetString(PyExc_TypeError,
4208 "zip_longest() got an unexpected keyword argument");
4209 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004210 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004211 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004212
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004213 /* args must be a tuple */
4214 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00004215
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004216 /* obtain iterators */
4217 ittuple = PyTuple_New(tuplesize);
4218 if (ittuple == NULL)
4219 return NULL;
4220 for (i=0; i < tuplesize; ++i) {
4221 PyObject *item = PyTuple_GET_ITEM(args, i);
4222 PyObject *it = PyObject_GetIter(item);
4223 if (it == NULL) {
4224 if (PyErr_ExceptionMatches(PyExc_TypeError))
4225 PyErr_Format(PyExc_TypeError,
4226 "zip_longest argument #%zd must support iteration",
4227 i+1);
4228 Py_DECREF(ittuple);
4229 return NULL;
4230 }
4231 PyTuple_SET_ITEM(ittuple, i, it);
4232 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004233
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004234 /* create a result holder */
4235 result = PyTuple_New(tuplesize);
4236 if (result == NULL) {
4237 Py_DECREF(ittuple);
4238 return NULL;
4239 }
4240 for (i=0 ; i < tuplesize ; i++) {
4241 Py_INCREF(Py_None);
4242 PyTuple_SET_ITEM(result, i, Py_None);
4243 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004244
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004245 /* create ziplongestobject structure */
4246 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4247 if (lz == NULL) {
4248 Py_DECREF(ittuple);
4249 Py_DECREF(result);
4250 return NULL;
4251 }
4252 lz->ittuple = ittuple;
4253 lz->tuplesize = tuplesize;
4254 lz->numactive = tuplesize;
4255 lz->result = result;
4256 Py_INCREF(fillvalue);
4257 lz->fillvalue = fillvalue;
4258 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004259}
4260
4261static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004262zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004263{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004264 PyObject_GC_UnTrack(lz);
4265 Py_XDECREF(lz->ittuple);
4266 Py_XDECREF(lz->result);
4267 Py_XDECREF(lz->fillvalue);
4268 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004269}
4270
4271static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004272zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004273{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004274 Py_VISIT(lz->ittuple);
4275 Py_VISIT(lz->result);
4276 Py_VISIT(lz->fillvalue);
4277 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004278}
4279
4280static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004281zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004282{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004283 Py_ssize_t i;
4284 Py_ssize_t tuplesize = lz->tuplesize;
4285 PyObject *result = lz->result;
4286 PyObject *it;
4287 PyObject *item;
4288 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004289
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004290 if (tuplesize == 0)
4291 return NULL;
4292 if (lz->numactive == 0)
4293 return NULL;
4294 if (Py_REFCNT(result) == 1) {
4295 Py_INCREF(result);
4296 for (i=0 ; i < tuplesize ; i++) {
4297 it = PyTuple_GET_ITEM(lz->ittuple, i);
4298 if (it == NULL) {
4299 Py_INCREF(lz->fillvalue);
4300 item = lz->fillvalue;
4301 } else {
4302 item = PyIter_Next(it);
4303 if (item == NULL) {
4304 lz->numactive -= 1;
4305 if (lz->numactive == 0 || PyErr_Occurred()) {
4306 lz->numactive = 0;
4307 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004308 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004309 } else {
4310 Py_INCREF(lz->fillvalue);
4311 item = lz->fillvalue;
4312 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4313 Py_DECREF(it);
4314 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004315 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004316 }
4317 olditem = PyTuple_GET_ITEM(result, i);
4318 PyTuple_SET_ITEM(result, i, item);
4319 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004320 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004321 } else {
4322 result = PyTuple_New(tuplesize);
4323 if (result == NULL)
4324 return NULL;
4325 for (i=0 ; i < tuplesize ; i++) {
4326 it = PyTuple_GET_ITEM(lz->ittuple, i);
4327 if (it == NULL) {
4328 Py_INCREF(lz->fillvalue);
4329 item = lz->fillvalue;
4330 } else {
4331 item = PyIter_Next(it);
4332 if (item == NULL) {
4333 lz->numactive -= 1;
4334 if (lz->numactive == 0 || PyErr_Occurred()) {
4335 lz->numactive = 0;
4336 Py_DECREF(result);
4337 return NULL;
4338 } else {
4339 Py_INCREF(lz->fillvalue);
4340 item = lz->fillvalue;
4341 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4342 Py_DECREF(it);
4343 }
4344 }
4345 }
4346 PyTuple_SET_ITEM(result, i, item);
4347 }
4348 }
4349 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004350}
4351
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004352static PyObject *
4353zip_longest_reduce(ziplongestobject *lz)
4354{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004355
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004356 /* Create a new tuple with empty sequences where appropriate to pickle.
4357 * Then use setstate to set the fillvalue
4358 */
4359 int i;
4360 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
4361 if (args == NULL)
4362 return NULL;
4363 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4364 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4365 if (elem == NULL) {
4366 elem = PyTuple_New(0);
4367 if (elem == NULL) {
4368 Py_DECREF(args);
4369 return NULL;
4370 }
4371 } else
4372 Py_INCREF(elem);
4373 PyTuple_SET_ITEM(args, i, elem);
4374 }
4375 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4376}
4377
4378static PyObject *
4379zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4380{
4381 Py_CLEAR(lz->fillvalue);
4382 lz->fillvalue = state;
4383 Py_INCREF(lz->fillvalue);
4384 Py_RETURN_NONE;
4385}
4386
4387static PyMethodDef zip_longest_methods[] = {
4388 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4389 reduce_doc},
4390 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4391 setstate_doc},
4392 {NULL, NULL} /* sentinel */
4393};
4394
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004395PyDoc_STRVAR(zip_longest_doc,
4396"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004397\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004398Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004399the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004400method continues until the longest iterable in the argument sequence\n\
4401is exhausted and then it raises StopIteration. When the shorter iterables\n\
4402are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4403defaults to None or can be specified by a keyword argument.\n\
4404");
4405
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004406static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004407 PyVarObject_HEAD_INIT(NULL, 0)
4408 "itertools.zip_longest", /* tp_name */
4409 sizeof(ziplongestobject), /* tp_basicsize */
4410 0, /* tp_itemsize */
4411 /* methods */
4412 (destructor)zip_longest_dealloc, /* tp_dealloc */
4413 0, /* tp_print */
4414 0, /* tp_getattr */
4415 0, /* tp_setattr */
4416 0, /* tp_reserved */
4417 0, /* tp_repr */
4418 0, /* tp_as_number */
4419 0, /* tp_as_sequence */
4420 0, /* tp_as_mapping */
4421 0, /* tp_hash */
4422 0, /* tp_call */
4423 0, /* tp_str */
4424 PyObject_GenericGetAttr, /* tp_getattro */
4425 0, /* tp_setattro */
4426 0, /* tp_as_buffer */
4427 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4428 Py_TPFLAGS_BASETYPE, /* tp_flags */
4429 zip_longest_doc, /* tp_doc */
4430 (traverseproc)zip_longest_traverse, /* tp_traverse */
4431 0, /* tp_clear */
4432 0, /* tp_richcompare */
4433 0, /* tp_weaklistoffset */
4434 PyObject_SelfIter, /* tp_iter */
4435 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004436 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004437 0, /* tp_members */
4438 0, /* tp_getset */
4439 0, /* tp_base */
4440 0, /* tp_dict */
4441 0, /* tp_descr_get */
4442 0, /* tp_descr_set */
4443 0, /* tp_dictoffset */
4444 0, /* tp_init */
4445 0, /* tp_alloc */
4446 zip_longest_new, /* tp_new */
4447 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004448};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004449
4450/* module level code ********************************************************/
4451
4452PyDoc_STRVAR(module_doc,
4453"Functional tools for creating and using iterators.\n\
4454\n\
4455Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004456count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004457cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004458repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004459\n\
4460Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004461accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004462chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
4463compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4464dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4465groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004466filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004467islice(seq, [start,] stop [, step]) --> elements from\n\
4468 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004469starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004470tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004471takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004472zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004473\n\
4474Combinatoric generators:\n\
4475product(p, q, ... [repeat=1]) --> cartesian product\n\
4476permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004477combinations(p, r)\n\
4478combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004479");
4480
4481
Raymond Hettingerad983e72003-11-12 14:32:26 +00004482static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004483 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4484 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004485};
4486
Martin v. Löwis1a214512008-06-11 05:26:20 +00004487
4488static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004489 PyModuleDef_HEAD_INIT,
4490 "itertools",
4491 module_doc,
4492 -1,
4493 module_methods,
4494 NULL,
4495 NULL,
4496 NULL,
4497 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004498};
4499
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004500PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004501PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004502{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004503 int i;
4504 PyObject *m;
4505 char *name;
4506 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004507 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004508 &combinations_type,
4509 &cwr_type,
4510 &cycle_type,
4511 &dropwhile_type,
4512 &takewhile_type,
4513 &islice_type,
4514 &starmap_type,
4515 &chain_type,
4516 &compress_type,
4517 &filterfalse_type,
4518 &count_type,
4519 &ziplongest_type,
4520 &permutations_type,
4521 &product_type,
4522 &repeat_type,
4523 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004524 &_grouper_type,
4525 &tee_type,
4526 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004527 NULL
4528 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004529
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004530 Py_TYPE(&teedataobject_type) = &PyType_Type;
4531 m = PyModule_Create(&itertoolsmodule);
4532 if (m == NULL)
4533 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004534
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004535 for (i=0 ; typelist[i] != NULL ; i++) {
4536 if (PyType_Ready(typelist[i]) < 0)
4537 return NULL;
4538 name = strchr(typelist[i]->tp_name, '.');
4539 assert (name != NULL);
4540 Py_INCREF(typelist[i]);
4541 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4542 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004543
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004544 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004545}