blob: 55d88a8c9f5abdc702f5336b662e5f1747dc4486 [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>
Raymond Hettinger96ef8112003-02-01 00:10:11 +00007*/
8
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00009
10/* groupby object ***********************************************************/
11
12typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000013 PyObject_HEAD
14 PyObject *it;
15 PyObject *keyfunc;
16 PyObject *tgtkey;
17 PyObject *currkey;
18 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000019} groupbyobject;
20
21static PyTypeObject groupby_type;
22static PyObject *_grouper_create(groupbyobject *, PyObject *);
23
24static PyObject *
25groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
26{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000027 static char *kwargs[] = {"iterable", "key", NULL};
28 groupbyobject *gbo;
29 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000030
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000031 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
32 &it, &keyfunc))
33 return NULL;
34
35 gbo = (groupbyobject *)type->tp_alloc(type, 0);
36 if (gbo == NULL)
37 return NULL;
38 gbo->tgtkey = NULL;
39 gbo->currkey = NULL;
40 gbo->currvalue = NULL;
41 gbo->keyfunc = keyfunc;
42 Py_INCREF(keyfunc);
43 gbo->it = PyObject_GetIter(it);
44 if (gbo->it == NULL) {
45 Py_DECREF(gbo);
46 return NULL;
47 }
48 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000049}
50
51static void
52groupby_dealloc(groupbyobject *gbo)
53{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000054 PyObject_GC_UnTrack(gbo);
55 Py_XDECREF(gbo->it);
56 Py_XDECREF(gbo->keyfunc);
57 Py_XDECREF(gbo->tgtkey);
58 Py_XDECREF(gbo->currkey);
59 Py_XDECREF(gbo->currvalue);
60 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000061}
62
63static int
64groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
65{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000066 Py_VISIT(gbo->it);
67 Py_VISIT(gbo->keyfunc);
68 Py_VISIT(gbo->tgtkey);
69 Py_VISIT(gbo->currkey);
70 Py_VISIT(gbo->currvalue);
71 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000072}
73
74static PyObject *
75groupby_next(groupbyobject *gbo)
76{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000077 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000078
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000079 /* skip to next iteration group */
80 for (;;) {
81 if (gbo->currkey == NULL)
82 /* pass */;
83 else if (gbo->tgtkey == NULL)
84 break;
85 else {
86 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000087
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000088 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
89 gbo->currkey, Py_EQ);
90 if (rcmp == -1)
91 return NULL;
92 else if (rcmp == 0)
93 break;
94 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000095
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000096 newvalue = PyIter_Next(gbo->it);
97 if (newvalue == NULL)
98 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000099
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000100 if (gbo->keyfunc == Py_None) {
101 newkey = newvalue;
102 Py_INCREF(newvalue);
103 } else {
104 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
105 newvalue, NULL);
106 if (newkey == NULL) {
107 Py_DECREF(newvalue);
108 return NULL;
109 }
110 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000111
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000112 tmp = gbo->currkey;
113 gbo->currkey = newkey;
114 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000115
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000116 tmp = gbo->currvalue;
117 gbo->currvalue = newvalue;
118 Py_XDECREF(tmp);
119 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000120
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000121 Py_INCREF(gbo->currkey);
122 tmp = gbo->tgtkey;
123 gbo->tgtkey = gbo->currkey;
124 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000125
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000126 grouper = _grouper_create(gbo, gbo->tgtkey);
127 if (grouper == NULL)
128 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000129
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000130 r = PyTuple_Pack(2, gbo->currkey, grouper);
131 Py_DECREF(grouper);
132 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000133}
134
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000135static PyObject *
136groupby_reduce(groupbyobject *lz)
137{
138 /* reduce as a 'new' call with an optional 'setstate' if groupby
139 * has started
140 */
141 PyObject *value;
142 if (lz->tgtkey && lz->currkey && lz->currvalue)
143 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
144 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
145 else
146 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
147 lz->it, lz->keyfunc);
148
149 return value;
150}
151
152PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
153
154static PyObject *
155groupby_setstate(groupbyobject *lz, PyObject *state)
156{
157 PyObject *currkey, *currvalue, *tgtkey;
158 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey))
159 return NULL;
160 Py_CLEAR(lz->currkey);
161 lz->currkey = currkey;
162 Py_INCREF(lz->currkey);
163 Py_CLEAR(lz->currvalue);
164 lz->currvalue = currvalue;
165 Py_INCREF(lz->currvalue);
166 Py_CLEAR(lz->tgtkey);
167 lz->tgtkey = tgtkey;
168 Py_INCREF(lz->tgtkey);
169 Py_RETURN_NONE;
170}
171
172PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
173
174static PyMethodDef groupby_methods[] = {
175 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
176 reduce_doc},
177 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
178 setstate_doc},
179 {NULL, NULL} /* sentinel */
180};
181
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000182PyDoc_STRVAR(groupby_doc,
183"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
184(key, sub-iterator) grouped by each value of key(value).\n");
185
186static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000187 PyVarObject_HEAD_INIT(NULL, 0)
188 "itertools.groupby", /* tp_name */
189 sizeof(groupbyobject), /* tp_basicsize */
190 0, /* tp_itemsize */
191 /* methods */
192 (destructor)groupby_dealloc, /* tp_dealloc */
193 0, /* tp_print */
194 0, /* tp_getattr */
195 0, /* tp_setattr */
196 0, /* tp_reserved */
197 0, /* tp_repr */
198 0, /* tp_as_number */
199 0, /* tp_as_sequence */
200 0, /* tp_as_mapping */
201 0, /* tp_hash */
202 0, /* tp_call */
203 0, /* tp_str */
204 PyObject_GenericGetAttr, /* tp_getattro */
205 0, /* tp_setattro */
206 0, /* tp_as_buffer */
207 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
208 Py_TPFLAGS_BASETYPE, /* tp_flags */
209 groupby_doc, /* tp_doc */
210 (traverseproc)groupby_traverse, /* tp_traverse */
211 0, /* tp_clear */
212 0, /* tp_richcompare */
213 0, /* tp_weaklistoffset */
214 PyObject_SelfIter, /* tp_iter */
215 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000216 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000217 0, /* tp_members */
218 0, /* tp_getset */
219 0, /* tp_base */
220 0, /* tp_dict */
221 0, /* tp_descr_get */
222 0, /* tp_descr_set */
223 0, /* tp_dictoffset */
224 0, /* tp_init */
225 0, /* tp_alloc */
226 groupby_new, /* tp_new */
227 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000228};
229
230
231/* _grouper object (internal) ************************************************/
232
233typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000234 PyObject_HEAD
235 PyObject *parent;
236 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000237} _grouperobject;
238
239static PyTypeObject _grouper_type;
240
241static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000242_grouper_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
243{
244 PyObject *parent, *tgtkey;
245
246 if (!PyArg_ParseTuple(args, "O!O", &groupby_type, &parent, &tgtkey))
247 return NULL;
248
249 return _grouper_create((groupbyobject*) parent, tgtkey);
250}
251
252static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000253_grouper_create(groupbyobject *parent, PyObject *tgtkey)
254{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000255 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000256
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000257 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
258 if (igo == NULL)
259 return NULL;
260 igo->parent = (PyObject *)parent;
261 Py_INCREF(parent);
262 igo->tgtkey = tgtkey;
263 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000264
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000265 PyObject_GC_Track(igo);
266 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000267}
268
269static void
270_grouper_dealloc(_grouperobject *igo)
271{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000272 PyObject_GC_UnTrack(igo);
273 Py_DECREF(igo->parent);
274 Py_DECREF(igo->tgtkey);
275 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000276}
277
278static int
279_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
280{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000281 Py_VISIT(igo->parent);
282 Py_VISIT(igo->tgtkey);
283 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000284}
285
286static PyObject *
287_grouper_next(_grouperobject *igo)
288{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000289 groupbyobject *gbo = (groupbyobject *)igo->parent;
290 PyObject *newvalue, *newkey, *r;
291 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000292
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000293 if (gbo->currvalue == NULL) {
294 newvalue = PyIter_Next(gbo->it);
295 if (newvalue == NULL)
296 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000297
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000298 if (gbo->keyfunc == Py_None) {
299 newkey = newvalue;
300 Py_INCREF(newvalue);
301 } else {
302 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
303 newvalue, NULL);
304 if (newkey == NULL) {
305 Py_DECREF(newvalue);
306 return NULL;
307 }
308 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000309
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000310 assert(gbo->currkey == NULL);
311 gbo->currkey = newkey;
312 gbo->currvalue = newvalue;
313 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000314
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000315 assert(gbo->currkey != NULL);
316 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
317 if (rcmp <= 0)
318 /* got any error or current group is end */
319 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000320
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000321 r = gbo->currvalue;
322 gbo->currvalue = NULL;
323 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000324
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000325 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000326}
327
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000328static PyObject *
329_grouper_reduce(_grouperobject *lz)
330{
331 return Py_BuildValue("O(OO)", Py_TYPE(lz),
332 lz->parent, lz->tgtkey);
333}
334
335static PyMethodDef _grouper_methods[] = {
336 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
337 reduce_doc},
338 {NULL, NULL} /* sentinel */
339};
340
341
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000342static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000343 PyVarObject_HEAD_INIT(NULL, 0)
344 "itertools._grouper", /* tp_name */
345 sizeof(_grouperobject), /* tp_basicsize */
346 0, /* tp_itemsize */
347 /* methods */
348 (destructor)_grouper_dealloc, /* tp_dealloc */
349 0, /* tp_print */
350 0, /* tp_getattr */
351 0, /* tp_setattr */
352 0, /* tp_reserved */
353 0, /* tp_repr */
354 0, /* tp_as_number */
355 0, /* tp_as_sequence */
356 0, /* tp_as_mapping */
357 0, /* tp_hash */
358 0, /* tp_call */
359 0, /* tp_str */
360 PyObject_GenericGetAttr, /* tp_getattro */
361 0, /* tp_setattro */
362 0, /* tp_as_buffer */
363 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
364 0, /* tp_doc */
365 (traverseproc)_grouper_traverse,/* tp_traverse */
366 0, /* tp_clear */
367 0, /* tp_richcompare */
368 0, /* tp_weaklistoffset */
369 PyObject_SelfIter, /* tp_iter */
370 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000371 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000372 0, /* tp_members */
373 0, /* tp_getset */
374 0, /* tp_base */
375 0, /* tp_dict */
376 0, /* tp_descr_get */
377 0, /* tp_descr_set */
378 0, /* tp_dictoffset */
379 0, /* tp_init */
380 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000381 _grouper_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000382 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000383};
384
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000385
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000386
Raymond Hettingerad983e72003-11-12 14:32:26 +0000387/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000388
Raymond Hettingerad983e72003-11-12 14:32:26 +0000389/* The teedataobject pre-allocates space for LINKCELLS number of objects.
390 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000391 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000392 the other structure members including PyHEAD overhead. The larger the
393 value, the less memory overhead per object and the less time spent
394 allocating/deallocating new links. The smaller the number, the less
395 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000396*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000397#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000398
399typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000400 PyObject_HEAD
401 PyObject *it;
Antoine Pitroub4a46cb2013-09-20 22:19:22 +0200402 int numread; /* 0 <= numread <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000403 PyObject *nextlink;
404 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000405} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000406
407typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000408 PyObject_HEAD
409 teedataobject *dataobj;
Antoine Pitroub4a46cb2013-09-20 22:19:22 +0200410 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000411 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000412} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000413
Raymond Hettingerad983e72003-11-12 14:32:26 +0000414static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415
416static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000417teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000418{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000419 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000420
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000421 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
422 if (tdo == NULL)
423 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000424
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000425 tdo->numread = 0;
426 tdo->nextlink = NULL;
427 Py_INCREF(it);
428 tdo->it = it;
429 PyObject_GC_Track(tdo);
430 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000431}
432
433static PyObject *
434teedataobject_jumplink(teedataobject *tdo)
435{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000436 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000437 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000438 Py_XINCREF(tdo->nextlink);
439 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000440}
441
442static PyObject *
443teedataobject_getitem(teedataobject *tdo, int i)
444{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000445 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000446
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000447 assert(i < LINKCELLS);
448 if (i < tdo->numread)
449 value = tdo->values[i];
450 else {
451 /* this is the lead iterator, so fetch more data */
452 assert(i == tdo->numread);
453 value = PyIter_Next(tdo->it);
454 if (value == NULL)
455 return NULL;
456 tdo->numread++;
457 tdo->values[i] = value;
458 }
459 Py_INCREF(value);
460 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000461}
462
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000463static int
464teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
465{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000466 int i;
467 Py_VISIT(tdo->it);
468 for (i = 0; i < tdo->numread; i++)
469 Py_VISIT(tdo->values[i]);
470 Py_VISIT(tdo->nextlink);
471 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000472}
473
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200474static void
475teedataobject_safe_decref(PyObject *obj)
476{
477 while (obj && Py_TYPE(obj) == &teedataobject_type &&
478 Py_REFCNT(obj) == 1) {
479 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
480 ((teedataobject *)obj)->nextlink = NULL;
481 Py_DECREF(obj);
482 obj = nextlink;
483 }
484 Py_XDECREF(obj);
485}
486
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000487static int
488teedataobject_clear(teedataobject *tdo)
489{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000490 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200491 PyObject *tmp;
492
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000493 Py_CLEAR(tdo->it);
494 for (i=0 ; i<tdo->numread ; i++)
495 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200496 tmp = tdo->nextlink;
497 tdo->nextlink = NULL;
498 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000499 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000500}
501
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000502static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000503teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000504{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000505 PyObject_GC_UnTrack(tdo);
506 teedataobject_clear(tdo);
507 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000508}
509
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000510static PyObject *
511teedataobject_reduce(teedataobject *tdo)
512{
513 int i;
514 /* create a temporary list of already iterated values */
515 PyObject *values = PyList_New(tdo->numread);
516 if (!values)
517 return NULL;
518 for (i=0 ; i<tdo->numread ; i++) {
519 Py_INCREF(tdo->values[i]);
520 PyList_SET_ITEM(values, i, tdo->values[i]);
521 }
522 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
523 values,
524 tdo->nextlink ? tdo->nextlink : Py_None);
525}
526
527static PyTypeObject teedataobject_type;
528
529static PyObject *
530teedataobject_new(PyTypeObject *type, PyObject *args, PyObject *kw)
531{
532 teedataobject *tdo;
533 PyObject *it, *values, *next;
534 Py_ssize_t i, len;
535
536 assert(type == &teedataobject_type);
537 if (!PyArg_ParseTuple(args, "OO!O", &it, &PyList_Type, &values, &next))
538 return NULL;
539
540 tdo = (teedataobject *)teedataobject_newinternal(it);
541 if (!tdo)
542 return NULL;
543
544 len = PyList_GET_SIZE(values);
545 if (len > LINKCELLS)
546 goto err;
547 for (i=0; i<len; i++) {
548 tdo->values[i] = PyList_GET_ITEM(values, i);
549 Py_INCREF(tdo->values[i]);
550 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200551 /* len <= LINKCELLS < INT_MAX */
552 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000553
554 if (len == LINKCELLS) {
555 if (next != Py_None) {
556 if (Py_TYPE(next) != &teedataobject_type)
557 goto err;
558 assert(tdo->nextlink == NULL);
559 Py_INCREF(next);
560 tdo->nextlink = next;
561 }
562 } else {
563 if (next != Py_None)
564 goto err; /* shouldn't have a next if we are not full */
565 }
566 return (PyObject*)tdo;
567
568err:
569 Py_XDECREF(tdo);
570 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
571 return NULL;
572}
573
574static PyMethodDef teedataobject_methods[] = {
575 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
576 reduce_doc},
577 {NULL, NULL} /* sentinel */
578};
579
Raymond Hettingerad983e72003-11-12 14:32:26 +0000580PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000581
Raymond Hettingerad983e72003-11-12 14:32:26 +0000582static PyTypeObject teedataobject_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000583 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000584 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000585 sizeof(teedataobject), /* tp_basicsize */
586 0, /* tp_itemsize */
587 /* methods */
588 (destructor)teedataobject_dealloc, /* tp_dealloc */
589 0, /* tp_print */
590 0, /* tp_getattr */
591 0, /* tp_setattr */
592 0, /* tp_reserved */
593 0, /* tp_repr */
594 0, /* tp_as_number */
595 0, /* tp_as_sequence */
596 0, /* tp_as_mapping */
597 0, /* tp_hash */
598 0, /* tp_call */
599 0, /* tp_str */
600 PyObject_GenericGetAttr, /* tp_getattro */
601 0, /* tp_setattro */
602 0, /* tp_as_buffer */
603 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
604 teedataobject_doc, /* tp_doc */
605 (traverseproc)teedataobject_traverse, /* tp_traverse */
606 (inquiry)teedataobject_clear, /* tp_clear */
607 0, /* tp_richcompare */
608 0, /* tp_weaklistoffset */
609 0, /* tp_iter */
610 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000611 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000612 0, /* tp_members */
613 0, /* tp_getset */
614 0, /* tp_base */
615 0, /* tp_dict */
616 0, /* tp_descr_get */
617 0, /* tp_descr_set */
618 0, /* tp_dictoffset */
619 0, /* tp_init */
620 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000621 teedataobject_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000622 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000623};
624
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000625
626static PyTypeObject tee_type;
627
628static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000629tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000630{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000631 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000632
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000633 if (to->index >= LINKCELLS) {
634 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200635 if (link == NULL)
636 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000637 Py_DECREF(to->dataobj);
638 to->dataobj = (teedataobject *)link;
639 to->index = 0;
640 }
641 value = teedataobject_getitem(to->dataobj, to->index);
642 if (value == NULL)
643 return NULL;
644 to->index++;
645 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000646}
647
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000648static int
649tee_traverse(teeobject *to, visitproc visit, void *arg)
650{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000651 Py_VISIT((PyObject *)to->dataobj);
652 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000653}
654
Raymond Hettingerad983e72003-11-12 14:32:26 +0000655static PyObject *
656tee_copy(teeobject *to)
657{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000658 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000659
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000660 newto = PyObject_GC_New(teeobject, &tee_type);
661 if (newto == NULL)
662 return NULL;
663 Py_INCREF(to->dataobj);
664 newto->dataobj = to->dataobj;
665 newto->index = to->index;
666 newto->weakreflist = NULL;
667 PyObject_GC_Track(newto);
668 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000669}
670
671PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
672
673static PyObject *
674tee_fromiterable(PyObject *iterable)
675{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000676 teeobject *to;
677 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000678
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000679 it = PyObject_GetIter(iterable);
680 if (it == NULL)
681 return NULL;
682 if (PyObject_TypeCheck(it, &tee_type)) {
683 to = (teeobject *)tee_copy((teeobject *)it);
684 goto done;
685 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000686
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000687 to = PyObject_GC_New(teeobject, &tee_type);
688 if (to == NULL)
689 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000690 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000691 if (!to->dataobj) {
692 PyObject_GC_Del(to);
693 to = NULL;
694 goto done;
695 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000696
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000697 to->index = 0;
698 to->weakreflist = NULL;
699 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000700done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000701 Py_XDECREF(it);
702 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000703}
704
705static PyObject *
706tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
707{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000708 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000709
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000710 if (!PyArg_UnpackTuple(args, "_tee", 1, 1, &iterable))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000711 return NULL;
712 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000713}
714
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000715static int
716tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000717{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000718 if (to->weakreflist != NULL)
719 PyObject_ClearWeakRefs((PyObject *) to);
720 Py_CLEAR(to->dataobj);
721 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000722}
723
724static void
725tee_dealloc(teeobject *to)
726{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000727 PyObject_GC_UnTrack(to);
728 tee_clear(to);
729 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000730}
731
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000732static PyObject *
733tee_reduce(teeobject *to)
734{
735 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
736}
737
738static PyObject *
739tee_setstate(teeobject *to, PyObject *state)
740{
741 teedataobject *tdo;
742 int index;
743 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index))
744 return NULL;
745 if (index < 0 || index > LINKCELLS) {
746 PyErr_SetString(PyExc_ValueError, "Index out of range");
747 return NULL;
748 }
749 Py_CLEAR(to->dataobj);
750 to->dataobj = tdo;
751 Py_INCREF(to->dataobj);
752 to->index = index;
753 Py_RETURN_NONE;
754}
755
Raymond Hettingerad983e72003-11-12 14:32:26 +0000756PyDoc_STRVAR(teeobject_doc,
757"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000758
Raymond Hettingerad983e72003-11-12 14:32:26 +0000759static PyMethodDef tee_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000760 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000761 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
762 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000763 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000764};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000765
766static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000767 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000768 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000769 sizeof(teeobject), /* tp_basicsize */
770 0, /* tp_itemsize */
771 /* methods */
772 (destructor)tee_dealloc, /* tp_dealloc */
773 0, /* tp_print */
774 0, /* tp_getattr */
775 0, /* tp_setattr */
776 0, /* tp_reserved */
777 0, /* tp_repr */
778 0, /* tp_as_number */
779 0, /* tp_as_sequence */
780 0, /* tp_as_mapping */
781 0, /* tp_hash */
782 0, /* tp_call */
783 0, /* tp_str */
784 0, /* tp_getattro */
785 0, /* tp_setattro */
786 0, /* tp_as_buffer */
787 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
788 teeobject_doc, /* tp_doc */
789 (traverseproc)tee_traverse, /* tp_traverse */
790 (inquiry)tee_clear, /* tp_clear */
791 0, /* tp_richcompare */
792 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
793 PyObject_SelfIter, /* tp_iter */
794 (iternextfunc)tee_next, /* tp_iternext */
795 tee_methods, /* tp_methods */
796 0, /* tp_members */
797 0, /* tp_getset */
798 0, /* tp_base */
799 0, /* tp_dict */
800 0, /* tp_descr_get */
801 0, /* tp_descr_set */
802 0, /* tp_dictoffset */
803 0, /* tp_init */
804 0, /* tp_alloc */
805 tee_new, /* tp_new */
806 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000807};
808
Raymond Hettingerad983e72003-11-12 14:32:26 +0000809static PyObject *
810tee(PyObject *self, PyObject *args)
811{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000812 Py_ssize_t i, n=2;
813 PyObject *it, *iterable, *copyable, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200814 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000815
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000816 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
817 return NULL;
818 if (n < 0) {
819 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
820 return NULL;
821 }
822 result = PyTuple_New(n);
823 if (result == NULL)
824 return NULL;
825 if (n == 0)
826 return result;
827 it = PyObject_GetIter(iterable);
828 if (it == NULL) {
829 Py_DECREF(result);
830 return NULL;
831 }
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200832 if (!_PyObject_HasAttrId(it, &PyId___copy__)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000833 copyable = tee_fromiterable(it);
834 Py_DECREF(it);
835 if (copyable == NULL) {
836 Py_DECREF(result);
837 return NULL;
838 }
839 } else
840 copyable = it;
841 PyTuple_SET_ITEM(result, 0, copyable);
842 for (i=1 ; i<n ; i++) {
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200843
844 copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000845 if (copyable == NULL) {
846 Py_DECREF(result);
847 return NULL;
848 }
849 PyTuple_SET_ITEM(result, i, copyable);
850 }
851 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000852}
853
854PyDoc_STRVAR(tee_doc,
855"tee(iterable, n=2) --> tuple of n independent iterators.");
856
857
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000858/* cycle object **********************************************************/
859
860typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000861 PyObject_HEAD
862 PyObject *it;
863 PyObject *saved;
864 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000865} cycleobject;
866
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000867static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000868
869static PyObject *
870cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
871{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000872 PyObject *it;
873 PyObject *iterable;
874 PyObject *saved;
875 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000876
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000877 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
878 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000879
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000880 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
881 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000882
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000883 /* Get iterator. */
884 it = PyObject_GetIter(iterable);
885 if (it == NULL)
886 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000887
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000888 saved = PyList_New(0);
889 if (saved == NULL) {
890 Py_DECREF(it);
891 return NULL;
892 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000893
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000894 /* create cycleobject structure */
895 lz = (cycleobject *)type->tp_alloc(type, 0);
896 if (lz == NULL) {
897 Py_DECREF(it);
898 Py_DECREF(saved);
899 return NULL;
900 }
901 lz->it = it;
902 lz->saved = saved;
903 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000904
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000905 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000906}
907
908static void
909cycle_dealloc(cycleobject *lz)
910{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000911 PyObject_GC_UnTrack(lz);
912 Py_XDECREF(lz->saved);
913 Py_XDECREF(lz->it);
914 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000915}
916
917static int
918cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
919{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000920 Py_VISIT(lz->it);
921 Py_VISIT(lz->saved);
922 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000923}
924
925static PyObject *
926cycle_next(cycleobject *lz)
927{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000928 PyObject *item;
929 PyObject *it;
930 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000931
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000932 while (1) {
933 item = PyIter_Next(lz->it);
934 if (item != NULL) {
935 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
936 Py_DECREF(item);
937 return NULL;
938 }
939 return item;
940 }
941 if (PyErr_Occurred()) {
942 if (PyErr_ExceptionMatches(PyExc_StopIteration))
943 PyErr_Clear();
944 else
945 return NULL;
946 }
947 if (PyList_Size(lz->saved) == 0)
948 return NULL;
949 it = PyObject_GetIter(lz->saved);
950 if (it == NULL)
951 return NULL;
952 tmp = lz->it;
953 lz->it = it;
954 lz->firstpass = 1;
955 Py_DECREF(tmp);
956 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000957}
958
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000959static PyObject *
960cycle_reduce(cycleobject *lz)
961{
962 /* Create a new cycle with the iterator tuple, then set
963 * the saved state on it.
964 */
Serhiy Storchakad269b5e2013-01-25 13:38:56 +0200965 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz),
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000966 lz->it, lz->saved, lz->firstpass);
967 }
968
969static PyObject *
970cycle_setstate(cycleobject *lz, PyObject *state)
971{
972 PyObject *saved=NULL;
973 int firstpass;
974 if (!PyArg_ParseTuple(state, "Oi", &saved, &firstpass))
975 return NULL;
976 Py_CLEAR(lz->saved);
977 lz->saved = saved;
978 Py_XINCREF(lz->saved);
979 lz->firstpass = firstpass != 0;
980 Py_RETURN_NONE;
981}
982
983static PyMethodDef cycle_methods[] = {
984 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
985 reduce_doc},
986 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
987 setstate_doc},
988 {NULL, NULL} /* sentinel */
989};
990
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000991PyDoc_STRVAR(cycle_doc,
992"cycle(iterable) --> cycle object\n\
993\n\
994Return elements from the iterable until it is exhausted.\n\
995Then repeat the sequence indefinitely.");
996
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000997static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000998 PyVarObject_HEAD_INIT(NULL, 0)
999 "itertools.cycle", /* tp_name */
1000 sizeof(cycleobject), /* tp_basicsize */
1001 0, /* tp_itemsize */
1002 /* methods */
1003 (destructor)cycle_dealloc, /* tp_dealloc */
1004 0, /* tp_print */
1005 0, /* tp_getattr */
1006 0, /* tp_setattr */
1007 0, /* tp_reserved */
1008 0, /* tp_repr */
1009 0, /* tp_as_number */
1010 0, /* tp_as_sequence */
1011 0, /* tp_as_mapping */
1012 0, /* tp_hash */
1013 0, /* tp_call */
1014 0, /* tp_str */
1015 PyObject_GenericGetAttr, /* tp_getattro */
1016 0, /* tp_setattro */
1017 0, /* tp_as_buffer */
1018 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1019 Py_TPFLAGS_BASETYPE, /* tp_flags */
1020 cycle_doc, /* tp_doc */
1021 (traverseproc)cycle_traverse, /* tp_traverse */
1022 0, /* tp_clear */
1023 0, /* tp_richcompare */
1024 0, /* tp_weaklistoffset */
1025 PyObject_SelfIter, /* tp_iter */
1026 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001027 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001028 0, /* tp_members */
1029 0, /* tp_getset */
1030 0, /* tp_base */
1031 0, /* tp_dict */
1032 0, /* tp_descr_get */
1033 0, /* tp_descr_set */
1034 0, /* tp_dictoffset */
1035 0, /* tp_init */
1036 0, /* tp_alloc */
1037 cycle_new, /* tp_new */
1038 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001039};
1040
1041
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001042/* dropwhile object **********************************************************/
1043
1044typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001045 PyObject_HEAD
1046 PyObject *func;
1047 PyObject *it;
1048 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001049} dropwhileobject;
1050
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001051static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001052
1053static PyObject *
1054dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1055{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001056 PyObject *func, *seq;
1057 PyObject *it;
1058 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001059
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001060 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
1061 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001062
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001063 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
1064 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001065
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001066 /* Get iterator. */
1067 it = PyObject_GetIter(seq);
1068 if (it == NULL)
1069 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001070
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001071 /* create dropwhileobject structure */
1072 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1073 if (lz == NULL) {
1074 Py_DECREF(it);
1075 return NULL;
1076 }
1077 Py_INCREF(func);
1078 lz->func = func;
1079 lz->it = it;
1080 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001081
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001082 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001083}
1084
1085static void
1086dropwhile_dealloc(dropwhileobject *lz)
1087{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001088 PyObject_GC_UnTrack(lz);
1089 Py_XDECREF(lz->func);
1090 Py_XDECREF(lz->it);
1091 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001092}
1093
1094static int
1095dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1096{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001097 Py_VISIT(lz->it);
1098 Py_VISIT(lz->func);
1099 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001100}
1101
1102static PyObject *
1103dropwhile_next(dropwhileobject *lz)
1104{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001105 PyObject *item, *good;
1106 PyObject *it = lz->it;
1107 long ok;
1108 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001109
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001110 iternext = *Py_TYPE(it)->tp_iternext;
1111 for (;;) {
1112 item = iternext(it);
1113 if (item == NULL)
1114 return NULL;
1115 if (lz->start == 1)
1116 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001117
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001118 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1119 if (good == NULL) {
1120 Py_DECREF(item);
1121 return NULL;
1122 }
1123 ok = PyObject_IsTrue(good);
1124 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001125 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001126 lz->start = 1;
1127 return item;
1128 }
1129 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001130 if (ok < 0)
1131 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001132 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001133}
1134
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001135static PyObject *
1136dropwhile_reduce(dropwhileobject *lz)
1137{
1138 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1139 lz->func, lz->it, lz->start);
1140}
1141
1142static PyObject *
1143dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1144{
1145 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001146 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001147 return NULL;
1148 lz->start = start;
1149 Py_RETURN_NONE;
1150}
1151
1152static PyMethodDef dropwhile_methods[] = {
1153 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1154 reduce_doc},
1155 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1156 setstate_doc},
1157 {NULL, NULL} /* sentinel */
1158};
1159
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001160PyDoc_STRVAR(dropwhile_doc,
1161"dropwhile(predicate, iterable) --> dropwhile object\n\
1162\n\
1163Drop items from the iterable while predicate(item) is true.\n\
1164Afterwards, return every element until the iterable is exhausted.");
1165
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001166static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001167 PyVarObject_HEAD_INIT(NULL, 0)
1168 "itertools.dropwhile", /* tp_name */
1169 sizeof(dropwhileobject), /* tp_basicsize */
1170 0, /* tp_itemsize */
1171 /* methods */
1172 (destructor)dropwhile_dealloc, /* tp_dealloc */
1173 0, /* tp_print */
1174 0, /* tp_getattr */
1175 0, /* tp_setattr */
1176 0, /* tp_reserved */
1177 0, /* tp_repr */
1178 0, /* tp_as_number */
1179 0, /* tp_as_sequence */
1180 0, /* tp_as_mapping */
1181 0, /* tp_hash */
1182 0, /* tp_call */
1183 0, /* tp_str */
1184 PyObject_GenericGetAttr, /* tp_getattro */
1185 0, /* tp_setattro */
1186 0, /* tp_as_buffer */
1187 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1188 Py_TPFLAGS_BASETYPE, /* tp_flags */
1189 dropwhile_doc, /* tp_doc */
1190 (traverseproc)dropwhile_traverse, /* tp_traverse */
1191 0, /* tp_clear */
1192 0, /* tp_richcompare */
1193 0, /* tp_weaklistoffset */
1194 PyObject_SelfIter, /* tp_iter */
1195 (iternextfunc)dropwhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001196 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001197 0, /* tp_members */
1198 0, /* tp_getset */
1199 0, /* tp_base */
1200 0, /* tp_dict */
1201 0, /* tp_descr_get */
1202 0, /* tp_descr_set */
1203 0, /* tp_dictoffset */
1204 0, /* tp_init */
1205 0, /* tp_alloc */
1206 dropwhile_new, /* tp_new */
1207 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001208};
1209
1210
1211/* takewhile object **********************************************************/
1212
1213typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001214 PyObject_HEAD
1215 PyObject *func;
1216 PyObject *it;
1217 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001218} takewhileobject;
1219
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001220static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001221
1222static PyObject *
1223takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1224{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001225 PyObject *func, *seq;
1226 PyObject *it;
1227 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001228
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001229 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1230 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001231
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001232 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1233 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001234
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001235 /* Get iterator. */
1236 it = PyObject_GetIter(seq);
1237 if (it == NULL)
1238 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001239
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001240 /* create takewhileobject structure */
1241 lz = (takewhileobject *)type->tp_alloc(type, 0);
1242 if (lz == NULL) {
1243 Py_DECREF(it);
1244 return NULL;
1245 }
1246 Py_INCREF(func);
1247 lz->func = func;
1248 lz->it = it;
1249 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001250
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001251 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001252}
1253
1254static void
1255takewhile_dealloc(takewhileobject *lz)
1256{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001257 PyObject_GC_UnTrack(lz);
1258 Py_XDECREF(lz->func);
1259 Py_XDECREF(lz->it);
1260 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001261}
1262
1263static int
1264takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1265{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001266 Py_VISIT(lz->it);
1267 Py_VISIT(lz->func);
1268 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001269}
1270
1271static PyObject *
1272takewhile_next(takewhileobject *lz)
1273{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001274 PyObject *item, *good;
1275 PyObject *it = lz->it;
1276 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001277
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001278 if (lz->stop == 1)
1279 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001280
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001281 item = (*Py_TYPE(it)->tp_iternext)(it);
1282 if (item == NULL)
1283 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001284
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001285 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1286 if (good == NULL) {
1287 Py_DECREF(item);
1288 return NULL;
1289 }
1290 ok = PyObject_IsTrue(good);
1291 Py_DECREF(good);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001292 if (ok == 1)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001293 return item;
1294 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001295 if (ok == 0)
1296 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001297 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001298}
1299
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001300static PyObject *
1301takewhile_reduce(takewhileobject *lz)
1302{
1303 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1304 lz->func, lz->it, lz->stop);
1305}
1306
1307static PyObject *
1308takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1309{
1310 int stop = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001311 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001312 return NULL;
1313 lz->stop = stop;
1314 Py_RETURN_NONE;
1315}
1316
1317static PyMethodDef takewhile_reduce_methods[] = {
1318 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1319 reduce_doc},
1320 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1321 setstate_doc},
1322 {NULL, NULL} /* sentinel */
1323};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001324PyDoc_STRVAR(takewhile_doc,
1325"takewhile(predicate, iterable) --> takewhile object\n\
1326\n\
1327Return successive entries from an iterable as long as the \n\
1328predicate evaluates to true for each entry.");
1329
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001330static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001331 PyVarObject_HEAD_INIT(NULL, 0)
1332 "itertools.takewhile", /* tp_name */
1333 sizeof(takewhileobject), /* tp_basicsize */
1334 0, /* tp_itemsize */
1335 /* methods */
1336 (destructor)takewhile_dealloc, /* tp_dealloc */
1337 0, /* tp_print */
1338 0, /* tp_getattr */
1339 0, /* tp_setattr */
1340 0, /* tp_reserved */
1341 0, /* tp_repr */
1342 0, /* tp_as_number */
1343 0, /* tp_as_sequence */
1344 0, /* tp_as_mapping */
1345 0, /* tp_hash */
1346 0, /* tp_call */
1347 0, /* tp_str */
1348 PyObject_GenericGetAttr, /* tp_getattro */
1349 0, /* tp_setattro */
1350 0, /* tp_as_buffer */
1351 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1352 Py_TPFLAGS_BASETYPE, /* tp_flags */
1353 takewhile_doc, /* tp_doc */
1354 (traverseproc)takewhile_traverse, /* tp_traverse */
1355 0, /* tp_clear */
1356 0, /* tp_richcompare */
1357 0, /* tp_weaklistoffset */
1358 PyObject_SelfIter, /* tp_iter */
1359 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001360 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001361 0, /* tp_members */
1362 0, /* tp_getset */
1363 0, /* tp_base */
1364 0, /* tp_dict */
1365 0, /* tp_descr_get */
1366 0, /* tp_descr_set */
1367 0, /* tp_dictoffset */
1368 0, /* tp_init */
1369 0, /* tp_alloc */
1370 takewhile_new, /* tp_new */
1371 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001372};
1373
1374
1375/* islice object ************************************************************/
1376
1377typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001378 PyObject_HEAD
1379 PyObject *it;
1380 Py_ssize_t next;
1381 Py_ssize_t stop;
1382 Py_ssize_t step;
1383 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001384} isliceobject;
1385
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001386static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001387
1388static PyObject *
1389islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1390{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001391 PyObject *seq;
1392 Py_ssize_t start=0, stop=-1, step=1;
1393 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1394 Py_ssize_t numargs;
1395 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001396
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001397 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1398 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001399
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001400 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1401 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001402
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001403 numargs = PyTuple_Size(args);
1404 if (numargs == 2) {
1405 if (a1 != Py_None) {
1406 stop = PyLong_AsSsize_t(a1);
1407 if (stop == -1) {
1408 if (PyErr_Occurred())
1409 PyErr_Clear();
1410 PyErr_SetString(PyExc_ValueError,
1411 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1412 return NULL;
1413 }
1414 }
1415 } else {
1416 if (a1 != Py_None)
1417 start = PyLong_AsSsize_t(a1);
1418 if (start == -1 && PyErr_Occurred())
1419 PyErr_Clear();
1420 if (a2 != Py_None) {
1421 stop = PyLong_AsSsize_t(a2);
1422 if (stop == -1) {
1423 if (PyErr_Occurred())
1424 PyErr_Clear();
1425 PyErr_SetString(PyExc_ValueError,
1426 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1427 return NULL;
1428 }
1429 }
1430 }
1431 if (start<0 || stop<-1) {
1432 PyErr_SetString(PyExc_ValueError,
1433 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1434 return NULL;
1435 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001436
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001437 if (a3 != NULL) {
1438 if (a3 != Py_None)
1439 step = PyLong_AsSsize_t(a3);
1440 if (step == -1 && PyErr_Occurred())
1441 PyErr_Clear();
1442 }
1443 if (step<1) {
1444 PyErr_SetString(PyExc_ValueError,
1445 "Step for islice() must be a positive integer or None.");
1446 return NULL;
1447 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001448
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001449 /* Get iterator. */
1450 it = PyObject_GetIter(seq);
1451 if (it == NULL)
1452 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001453
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001454 /* create isliceobject structure */
1455 lz = (isliceobject *)type->tp_alloc(type, 0);
1456 if (lz == NULL) {
1457 Py_DECREF(it);
1458 return NULL;
1459 }
1460 lz->it = it;
1461 lz->next = start;
1462 lz->stop = stop;
1463 lz->step = step;
1464 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001465
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001466 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001467}
1468
1469static void
1470islice_dealloc(isliceobject *lz)
1471{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001472 PyObject_GC_UnTrack(lz);
1473 Py_XDECREF(lz->it);
1474 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001475}
1476
1477static int
1478islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1479{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001480 Py_VISIT(lz->it);
1481 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001482}
1483
1484static PyObject *
1485islice_next(isliceobject *lz)
1486{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001487 PyObject *item;
1488 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001489 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001490 Py_ssize_t oldnext;
1491 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001492
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001493 if (it == NULL)
1494 return NULL;
1495
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001496 iternext = *Py_TYPE(it)->tp_iternext;
1497 while (lz->cnt < lz->next) {
1498 item = iternext(it);
1499 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001500 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001501 Py_DECREF(item);
1502 lz->cnt++;
1503 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001504 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001505 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001506 item = iternext(it);
1507 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001508 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001509 lz->cnt++;
1510 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001511 /* The (size_t) cast below avoids the danger of undefined
1512 behaviour from signed integer overflow. */
1513 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001514 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1515 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001516 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001517
1518empty:
1519 Py_CLEAR(lz->it);
1520 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001521}
1522
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001523static PyObject *
1524islice_reduce(isliceobject *lz)
1525{
1526 /* When unpickled, generate a new object with the same bounds,
1527 * then 'setstate' with the next and count
1528 */
1529 PyObject *stop;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001530 if (lz->it == NULL) {
1531 PyObject *empty_list;
1532 PyObject *empty_it;
1533 empty_list = PyList_New(0);
1534 if (empty_list == NULL)
1535 return NULL;
1536 empty_it = PyObject_GetIter(empty_list);
1537 Py_DECREF(empty_list);
1538 if (empty_it == NULL)
1539 return NULL;
1540 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1541 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001542 if (lz->stop == -1) {
1543 stop = Py_None;
1544 Py_INCREF(stop);
1545 } else {
1546 stop = PyLong_FromSsize_t(lz->stop);
1547 if (stop == NULL)
1548 return NULL;
1549 }
1550 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1551 lz->it, lz->next, stop, lz->step,
1552 lz->cnt);
1553}
1554
1555static PyObject *
1556islice_setstate(isliceobject *lz, PyObject *state)
1557{
1558 Py_ssize_t cnt = PyLong_AsSsize_t(state);
1559 if (cnt == -1 && PyErr_Occurred())
1560 return NULL;
1561 lz->cnt = cnt;
1562 Py_RETURN_NONE;
1563}
1564
1565static PyMethodDef islice_methods[] = {
1566 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1567 reduce_doc},
1568 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1569 setstate_doc},
1570 {NULL, NULL} /* sentinel */
1571};
1572
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001573PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001574"islice(iterable, stop) --> islice object\n\
1575islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001576\n\
1577Return an iterator whose next() method returns selected values from an\n\
1578iterable. If start is specified, will skip all preceding elements;\n\
1579otherwise, start defaults to zero. Step defaults to one. If\n\
1580specified as another value, step determines how many values are \n\
1581skipped between successive calls. Works like a slice() on a list\n\
1582but returns an iterator.");
1583
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001584static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001585 PyVarObject_HEAD_INIT(NULL, 0)
1586 "itertools.islice", /* tp_name */
1587 sizeof(isliceobject), /* tp_basicsize */
1588 0, /* tp_itemsize */
1589 /* methods */
1590 (destructor)islice_dealloc, /* tp_dealloc */
1591 0, /* tp_print */
1592 0, /* tp_getattr */
1593 0, /* tp_setattr */
1594 0, /* tp_reserved */
1595 0, /* tp_repr */
1596 0, /* tp_as_number */
1597 0, /* tp_as_sequence */
1598 0, /* tp_as_mapping */
1599 0, /* tp_hash */
1600 0, /* tp_call */
1601 0, /* tp_str */
1602 PyObject_GenericGetAttr, /* tp_getattro */
1603 0, /* tp_setattro */
1604 0, /* tp_as_buffer */
1605 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1606 Py_TPFLAGS_BASETYPE, /* tp_flags */
1607 islice_doc, /* tp_doc */
1608 (traverseproc)islice_traverse, /* tp_traverse */
1609 0, /* tp_clear */
1610 0, /* tp_richcompare */
1611 0, /* tp_weaklistoffset */
1612 PyObject_SelfIter, /* tp_iter */
1613 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001614 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001615 0, /* tp_members */
1616 0, /* tp_getset */
1617 0, /* tp_base */
1618 0, /* tp_dict */
1619 0, /* tp_descr_get */
1620 0, /* tp_descr_set */
1621 0, /* tp_dictoffset */
1622 0, /* tp_init */
1623 0, /* tp_alloc */
1624 islice_new, /* tp_new */
1625 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001626};
1627
1628
1629/* starmap object ************************************************************/
1630
1631typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001632 PyObject_HEAD
1633 PyObject *func;
1634 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001635} starmapobject;
1636
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001637static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001638
1639static PyObject *
1640starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1641{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001642 PyObject *func, *seq;
1643 PyObject *it;
1644 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001645
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001646 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1647 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001648
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001649 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1650 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001651
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001652 /* Get iterator. */
1653 it = PyObject_GetIter(seq);
1654 if (it == NULL)
1655 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001656
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001657 /* create starmapobject structure */
1658 lz = (starmapobject *)type->tp_alloc(type, 0);
1659 if (lz == NULL) {
1660 Py_DECREF(it);
1661 return NULL;
1662 }
1663 Py_INCREF(func);
1664 lz->func = func;
1665 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001666
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001667 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001668}
1669
1670static void
1671starmap_dealloc(starmapobject *lz)
1672{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001673 PyObject_GC_UnTrack(lz);
1674 Py_XDECREF(lz->func);
1675 Py_XDECREF(lz->it);
1676 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001677}
1678
1679static int
1680starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1681{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001682 Py_VISIT(lz->it);
1683 Py_VISIT(lz->func);
1684 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001685}
1686
1687static PyObject *
1688starmap_next(starmapobject *lz)
1689{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001690 PyObject *args;
1691 PyObject *result;
1692 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001693
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001694 args = (*Py_TYPE(it)->tp_iternext)(it);
1695 if (args == NULL)
1696 return NULL;
1697 if (!PyTuple_CheckExact(args)) {
1698 PyObject *newargs = PySequence_Tuple(args);
1699 Py_DECREF(args);
1700 if (newargs == NULL)
1701 return NULL;
1702 args = newargs;
1703 }
1704 result = PyObject_Call(lz->func, args, NULL);
1705 Py_DECREF(args);
1706 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001707}
1708
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001709static PyObject *
1710starmap_reduce(starmapobject *lz)
1711{
1712 /* Just pickle the iterator */
1713 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1714}
1715
1716static PyMethodDef starmap_methods[] = {
1717 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1718 reduce_doc},
1719 {NULL, NULL} /* sentinel */
1720};
1721
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001722PyDoc_STRVAR(starmap_doc,
1723"starmap(function, sequence) --> starmap object\n\
1724\n\
1725Return an iterator whose values are returned from the function evaluated\n\
Serhiy Storchakad65c9492015-11-02 14:10:23 +02001726with an argument tuple taken from the given sequence.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001727
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001728static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001729 PyVarObject_HEAD_INIT(NULL, 0)
1730 "itertools.starmap", /* tp_name */
1731 sizeof(starmapobject), /* tp_basicsize */
1732 0, /* tp_itemsize */
1733 /* methods */
1734 (destructor)starmap_dealloc, /* tp_dealloc */
1735 0, /* tp_print */
1736 0, /* tp_getattr */
1737 0, /* tp_setattr */
1738 0, /* tp_reserved */
1739 0, /* tp_repr */
1740 0, /* tp_as_number */
1741 0, /* tp_as_sequence */
1742 0, /* tp_as_mapping */
1743 0, /* tp_hash */
1744 0, /* tp_call */
1745 0, /* tp_str */
1746 PyObject_GenericGetAttr, /* tp_getattro */
1747 0, /* tp_setattro */
1748 0, /* tp_as_buffer */
1749 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1750 Py_TPFLAGS_BASETYPE, /* tp_flags */
1751 starmap_doc, /* tp_doc */
1752 (traverseproc)starmap_traverse, /* tp_traverse */
1753 0, /* tp_clear */
1754 0, /* tp_richcompare */
1755 0, /* tp_weaklistoffset */
1756 PyObject_SelfIter, /* tp_iter */
1757 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001758 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001759 0, /* tp_members */
1760 0, /* tp_getset */
1761 0, /* tp_base */
1762 0, /* tp_dict */
1763 0, /* tp_descr_get */
1764 0, /* tp_descr_set */
1765 0, /* tp_dictoffset */
1766 0, /* tp_init */
1767 0, /* tp_alloc */
1768 starmap_new, /* tp_new */
1769 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001770};
1771
1772
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001773/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001774
1775typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001776 PyObject_HEAD
1777 PyObject *source; /* Iterator over input iterables */
1778 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001779} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001780
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001781static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001782
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001783static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001784chain_new_internal(PyTypeObject *type, PyObject *source)
1785{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001786 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001787
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001788 lz = (chainobject *)type->tp_alloc(type, 0);
1789 if (lz == NULL) {
1790 Py_DECREF(source);
1791 return NULL;
1792 }
1793
1794 lz->source = source;
1795 lz->active = NULL;
1796 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001797}
1798
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001799static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001800chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001801{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001802 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001803
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001804 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1805 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001806
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001807 source = PyObject_GetIter(args);
1808 if (source == NULL)
1809 return NULL;
1810
1811 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001812}
1813
1814static PyObject *
1815chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1816{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001817 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001818
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001819 source = PyObject_GetIter(arg);
1820 if (source == NULL)
1821 return NULL;
1822
1823 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001824}
1825
1826static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001827chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001828{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001829 PyObject_GC_UnTrack(lz);
1830 Py_XDECREF(lz->active);
1831 Py_XDECREF(lz->source);
1832 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001833}
1834
Raymond Hettinger2012f172003-02-07 05:32:58 +00001835static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001836chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001837{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001838 Py_VISIT(lz->source);
1839 Py_VISIT(lz->active);
1840 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001841}
1842
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001843static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001844chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001845{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001846 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001847
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001848 if (lz->source == NULL)
1849 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001850
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001851 if (lz->active == NULL) {
1852 PyObject *iterable = PyIter_Next(lz->source);
1853 if (iterable == NULL) {
1854 Py_CLEAR(lz->source);
1855 return NULL; /* no more input sources */
1856 }
1857 lz->active = PyObject_GetIter(iterable);
1858 Py_DECREF(iterable);
1859 if (lz->active == NULL) {
1860 Py_CLEAR(lz->source);
1861 return NULL; /* input not iterable */
1862 }
1863 }
1864 item = PyIter_Next(lz->active);
1865 if (item != NULL)
1866 return item;
1867 if (PyErr_Occurred()) {
1868 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1869 PyErr_Clear();
1870 else
1871 return NULL; /* input raised an exception */
1872 }
1873 Py_CLEAR(lz->active);
1874 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001875}
1876
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001877static PyObject *
1878chain_reduce(chainobject *lz)
1879{
1880 if (lz->source) {
1881 /* we can't pickle function objects (itertools.from_iterable) so
1882 * we must use setstate to replace the iterable. One day we
1883 * will fix pickling of functions
1884 */
1885 if (lz->active) {
1886 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1887 } else {
1888 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1889 }
1890 } else {
1891 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1892 }
1893 return NULL;
1894}
1895
1896static PyObject *
1897chain_setstate(chainobject *lz, PyObject *state)
1898{
1899 PyObject *source, *active=NULL;
1900 if (! PyArg_ParseTuple(state, "O|O", &source, &active))
1901 return NULL;
1902
1903 Py_CLEAR(lz->source);
1904 lz->source = source;
1905 Py_INCREF(lz->source);
1906 Py_CLEAR(lz->active);
1907 lz->active = active;
1908 Py_XINCREF(lz->active);
1909 Py_RETURN_NONE;
1910}
1911
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001912PyDoc_STRVAR(chain_doc,
1913"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001914\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001915Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001916first iterable until it is exhausted, then elements from the next\n\
1917iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001918
Christian Heimesf16baeb2008-02-29 14:57:44 +00001919PyDoc_STRVAR(chain_from_iterable_doc,
1920"chain.from_iterable(iterable) --> chain object\n\
1921\n\
1922Alternate chain() contructor taking a single iterable argument\n\
1923that evaluates lazily.");
1924
1925static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001926 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1927 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001928 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1929 reduce_doc},
1930 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1931 setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001932 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001933};
1934
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001935static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001936 PyVarObject_HEAD_INIT(NULL, 0)
1937 "itertools.chain", /* tp_name */
1938 sizeof(chainobject), /* tp_basicsize */
1939 0, /* tp_itemsize */
1940 /* methods */
1941 (destructor)chain_dealloc, /* tp_dealloc */
1942 0, /* tp_print */
1943 0, /* tp_getattr */
1944 0, /* tp_setattr */
1945 0, /* tp_reserved */
1946 0, /* tp_repr */
1947 0, /* tp_as_number */
1948 0, /* tp_as_sequence */
1949 0, /* tp_as_mapping */
1950 0, /* tp_hash */
1951 0, /* tp_call */
1952 0, /* tp_str */
1953 PyObject_GenericGetAttr, /* tp_getattro */
1954 0, /* tp_setattro */
1955 0, /* tp_as_buffer */
1956 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1957 Py_TPFLAGS_BASETYPE, /* tp_flags */
1958 chain_doc, /* tp_doc */
1959 (traverseproc)chain_traverse, /* tp_traverse */
1960 0, /* tp_clear */
1961 0, /* tp_richcompare */
1962 0, /* tp_weaklistoffset */
1963 PyObject_SelfIter, /* tp_iter */
1964 (iternextfunc)chain_next, /* tp_iternext */
1965 chain_methods, /* tp_methods */
1966 0, /* tp_members */
1967 0, /* tp_getset */
1968 0, /* tp_base */
1969 0, /* tp_dict */
1970 0, /* tp_descr_get */
1971 0, /* tp_descr_set */
1972 0, /* tp_dictoffset */
1973 0, /* tp_init */
1974 0, /* tp_alloc */
1975 chain_new, /* tp_new */
1976 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001977};
1978
1979
Christian Heimesc3f30c42008-02-22 16:37:40 +00001980/* product object ************************************************************/
1981
1982typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001983 PyObject_HEAD
1984 PyObject *pools; /* tuple of pool tuples */
1985 Py_ssize_t *indices; /* one index per pool */
1986 PyObject *result; /* most recently returned result tuple */
1987 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001988} productobject;
1989
1990static PyTypeObject product_type;
1991
1992static PyObject *
1993product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1994{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001995 productobject *lz;
1996 Py_ssize_t nargs, npools, repeat=1;
1997 PyObject *pools = NULL;
1998 Py_ssize_t *indices = NULL;
1999 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002000
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002001 if (kwds != NULL) {
2002 char *kwlist[] = {"repeat", 0};
2003 PyObject *tmpargs = PyTuple_New(0);
2004 if (tmpargs == NULL)
2005 return NULL;
2006 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
2007 Py_DECREF(tmpargs);
2008 return NULL;
2009 }
2010 Py_DECREF(tmpargs);
2011 if (repeat < 0) {
2012 PyErr_SetString(PyExc_ValueError,
2013 "repeat argument cannot be negative");
2014 return NULL;
2015 }
2016 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002017
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002018 assert(PyTuple_CheckExact(args));
2019 if (repeat == 0) {
2020 nargs = 0;
2021 } else {
2022 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002023 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002024 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2025 return NULL;
2026 }
2027 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002028 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002029
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002030 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002031 if (indices == NULL) {
2032 PyErr_NoMemory();
2033 goto error;
2034 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002035
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002036 pools = PyTuple_New(npools);
2037 if (pools == NULL)
2038 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002039
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002040 for (i=0; i < nargs ; ++i) {
2041 PyObject *item = PyTuple_GET_ITEM(args, i);
2042 PyObject *pool = PySequence_Tuple(item);
2043 if (pool == NULL)
2044 goto error;
2045 PyTuple_SET_ITEM(pools, i, pool);
2046 indices[i] = 0;
2047 }
2048 for ( ; i < npools; ++i) {
2049 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2050 Py_INCREF(pool);
2051 PyTuple_SET_ITEM(pools, i, pool);
2052 indices[i] = 0;
2053 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002054
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002055 /* create productobject structure */
2056 lz = (productobject *)type->tp_alloc(type, 0);
2057 if (lz == NULL)
2058 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002059
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002060 lz->pools = pools;
2061 lz->indices = indices;
2062 lz->result = NULL;
2063 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002064
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002065 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002066
2067error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002068 if (indices != NULL)
2069 PyMem_Free(indices);
2070 Py_XDECREF(pools);
2071 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002072}
2073
2074static void
2075product_dealloc(productobject *lz)
2076{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002077 PyObject_GC_UnTrack(lz);
2078 Py_XDECREF(lz->pools);
2079 Py_XDECREF(lz->result);
2080 if (lz->indices != NULL)
2081 PyMem_Free(lz->indices);
2082 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002083}
2084
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002085static PyObject *
2086product_sizeof(productobject *lz, void *unused)
2087{
2088 Py_ssize_t res;
2089
2090 res = sizeof(productobject);
2091 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2092 return PyLong_FromSsize_t(res);
2093}
2094
2095PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2096
Christian Heimesc3f30c42008-02-22 16:37:40 +00002097static int
2098product_traverse(productobject *lz, visitproc visit, void *arg)
2099{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002100 Py_VISIT(lz->pools);
2101 Py_VISIT(lz->result);
2102 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002103}
2104
2105static PyObject *
2106product_next(productobject *lz)
2107{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002108 PyObject *pool;
2109 PyObject *elem;
2110 PyObject *oldelem;
2111 PyObject *pools = lz->pools;
2112 PyObject *result = lz->result;
2113 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2114 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002115
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002116 if (lz->stopped)
2117 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002118
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002119 if (result == NULL) {
2120 /* On the first pass, return an initial tuple filled with the
2121 first element from each pool. */
2122 result = PyTuple_New(npools);
2123 if (result == NULL)
2124 goto empty;
2125 lz->result = result;
2126 for (i=0; i < npools; i++) {
2127 pool = PyTuple_GET_ITEM(pools, i);
2128 if (PyTuple_GET_SIZE(pool) == 0)
2129 goto empty;
2130 elem = PyTuple_GET_ITEM(pool, 0);
2131 Py_INCREF(elem);
2132 PyTuple_SET_ITEM(result, i, elem);
2133 }
2134 } else {
2135 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002136
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002137 /* Copy the previous result tuple or re-use it if available */
2138 if (Py_REFCNT(result) > 1) {
2139 PyObject *old_result = result;
2140 result = PyTuple_New(npools);
2141 if (result == NULL)
2142 goto empty;
2143 lz->result = result;
2144 for (i=0; i < npools; i++) {
2145 elem = PyTuple_GET_ITEM(old_result, i);
2146 Py_INCREF(elem);
2147 PyTuple_SET_ITEM(result, i, elem);
2148 }
2149 Py_DECREF(old_result);
2150 }
2151 /* Now, we've got the only copy so we can update it in-place */
2152 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002153
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002154 /* Update the pool indices right-to-left. Only advance to the
2155 next pool when the previous one rolls-over */
2156 for (i=npools-1 ; i >= 0 ; i--) {
2157 pool = PyTuple_GET_ITEM(pools, i);
2158 indices[i]++;
2159 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2160 /* Roll-over and advance to next pool */
2161 indices[i] = 0;
2162 elem = PyTuple_GET_ITEM(pool, 0);
2163 Py_INCREF(elem);
2164 oldelem = PyTuple_GET_ITEM(result, i);
2165 PyTuple_SET_ITEM(result, i, elem);
2166 Py_DECREF(oldelem);
2167 } else {
2168 /* No rollover. Just increment and stop here. */
2169 elem = PyTuple_GET_ITEM(pool, indices[i]);
2170 Py_INCREF(elem);
2171 oldelem = PyTuple_GET_ITEM(result, i);
2172 PyTuple_SET_ITEM(result, i, elem);
2173 Py_DECREF(oldelem);
2174 break;
2175 }
2176 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002177
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002178 /* If i is negative, then the indices have all rolled-over
2179 and we're done. */
2180 if (i < 0)
2181 goto empty;
2182 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002183
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002184 Py_INCREF(result);
2185 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002186
2187empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002188 lz->stopped = 1;
2189 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002190}
2191
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002192static PyObject *
2193product_reduce(productobject *lz)
2194{
2195 if (lz->stopped) {
2196 return Py_BuildValue("O(())", Py_TYPE(lz));
2197 } else if (lz->result == NULL) {
2198 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2199 } else {
2200 PyObject *indices;
2201 Py_ssize_t n, i;
2202
2203 /* we must pickle the indices use them for setstate, and
2204 * additionally indicate that the iterator has started
2205 */
2206 n = PyTuple_GET_SIZE(lz->pools);
2207 indices = PyTuple_New(n);
2208 if (indices == NULL)
2209 return NULL;
2210 for (i=0; i<n; i++){
2211 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2212 if (!index) {
2213 Py_DECREF(indices);
2214 return NULL;
2215 }
2216 PyTuple_SET_ITEM(indices, i, index);
2217 }
2218 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2219 }
2220}
2221
2222static PyObject *
2223product_setstate(productobject *lz, PyObject *state)
2224{
2225 PyObject *result;
2226 Py_ssize_t n, i;
2227
2228 n = PyTuple_GET_SIZE(lz->pools);
2229 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2230 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2231 return NULL;
2232 }
2233 for (i=0; i<n; i++)
2234 {
2235 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2236 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002237 PyObject* pool;
2238 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002239 if (index < 0 && PyErr_Occurred())
2240 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002241 pool = PyTuple_GET_ITEM(lz->pools, i);
2242 poolsize = PyTuple_GET_SIZE(pool);
2243 if (poolsize == 0) {
2244 lz->stopped = 1;
2245 Py_RETURN_NONE;
2246 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002247 /* clamp the index */
2248 if (index < 0)
2249 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002250 else if (index > poolsize-1)
2251 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002252 lz->indices[i] = index;
2253 }
2254
2255 result = PyTuple_New(n);
2256 if (!result)
2257 return NULL;
2258 for (i=0; i<n; i++) {
2259 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2260 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2261 Py_INCREF(element);
2262 PyTuple_SET_ITEM(result, i, element);
2263 }
2264 Py_CLEAR(lz->result);
2265 lz->result = result;
2266 Py_RETURN_NONE;
2267}
2268
2269static PyMethodDef product_methods[] = {
2270 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2271 reduce_doc},
2272 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2273 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002274 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2275 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002276 {NULL, NULL} /* sentinel */
2277};
2278
Christian Heimesc3f30c42008-02-22 16:37:40 +00002279PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002280"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002281\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002282Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002283For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2284The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2285cycle in a manner similar to an odometer (with the rightmost element changing\n\
2286on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002287To compute the product of an iterable with itself, specify the number\n\
2288of repetitions with the optional repeat keyword argument. For example,\n\
2289product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002290product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2291product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2292
2293static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002294 PyVarObject_HEAD_INIT(NULL, 0)
2295 "itertools.product", /* tp_name */
2296 sizeof(productobject), /* tp_basicsize */
2297 0, /* tp_itemsize */
2298 /* methods */
2299 (destructor)product_dealloc, /* tp_dealloc */
2300 0, /* tp_print */
2301 0, /* tp_getattr */
2302 0, /* tp_setattr */
2303 0, /* tp_reserved */
2304 0, /* tp_repr */
2305 0, /* tp_as_number */
2306 0, /* tp_as_sequence */
2307 0, /* tp_as_mapping */
2308 0, /* tp_hash */
2309 0, /* tp_call */
2310 0, /* tp_str */
2311 PyObject_GenericGetAttr, /* tp_getattro */
2312 0, /* tp_setattro */
2313 0, /* tp_as_buffer */
2314 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2315 Py_TPFLAGS_BASETYPE, /* tp_flags */
2316 product_doc, /* tp_doc */
2317 (traverseproc)product_traverse, /* tp_traverse */
2318 0, /* tp_clear */
2319 0, /* tp_richcompare */
2320 0, /* tp_weaklistoffset */
2321 PyObject_SelfIter, /* tp_iter */
2322 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002323 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002324 0, /* tp_members */
2325 0, /* tp_getset */
2326 0, /* tp_base */
2327 0, /* tp_dict */
2328 0, /* tp_descr_get */
2329 0, /* tp_descr_set */
2330 0, /* tp_dictoffset */
2331 0, /* tp_init */
2332 0, /* tp_alloc */
2333 product_new, /* tp_new */
2334 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002335};
2336
2337
Christian Heimes380f7f22008-02-28 11:19:05 +00002338/* combinations object ************************************************************/
2339
2340typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002341 PyObject_HEAD
2342 PyObject *pool; /* input converted to a tuple */
2343 Py_ssize_t *indices; /* one index per result element */
2344 PyObject *result; /* most recently returned result tuple */
2345 Py_ssize_t r; /* size of result tuple */
2346 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002347} combinationsobject;
2348
2349static PyTypeObject combinations_type;
2350
2351static PyObject *
2352combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2353{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002354 combinationsobject *co;
2355 Py_ssize_t n;
2356 Py_ssize_t r;
2357 PyObject *pool = NULL;
2358 PyObject *iterable = NULL;
2359 Py_ssize_t *indices = NULL;
2360 Py_ssize_t i;
2361 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002362
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002363 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2364 &iterable, &r))
2365 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002366
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002367 pool = PySequence_Tuple(iterable);
2368 if (pool == NULL)
2369 goto error;
2370 n = PyTuple_GET_SIZE(pool);
2371 if (r < 0) {
2372 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2373 goto error;
2374 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002375
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002376 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002377 if (indices == NULL) {
2378 PyErr_NoMemory();
2379 goto error;
2380 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002381
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002382 for (i=0 ; i<r ; i++)
2383 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002384
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002385 /* create combinationsobject structure */
2386 co = (combinationsobject *)type->tp_alloc(type, 0);
2387 if (co == NULL)
2388 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002389
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002390 co->pool = pool;
2391 co->indices = indices;
2392 co->result = NULL;
2393 co->r = r;
2394 co->stopped = r > n ? 1 : 0;
2395
2396 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002397
2398error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002399 if (indices != NULL)
2400 PyMem_Free(indices);
2401 Py_XDECREF(pool);
2402 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002403}
2404
2405static void
2406combinations_dealloc(combinationsobject *co)
2407{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002408 PyObject_GC_UnTrack(co);
2409 Py_XDECREF(co->pool);
2410 Py_XDECREF(co->result);
2411 if (co->indices != NULL)
2412 PyMem_Free(co->indices);
2413 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002414}
2415
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002416static PyObject *
2417combinations_sizeof(combinationsobject *co, void *unused)
2418{
2419 Py_ssize_t res;
2420
2421 res = sizeof(combinationsobject);
2422 res += co->r * sizeof(Py_ssize_t);
2423 return PyLong_FromSsize_t(res);
2424}
2425
Christian Heimes380f7f22008-02-28 11:19:05 +00002426static int
2427combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2428{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002429 Py_VISIT(co->pool);
2430 Py_VISIT(co->result);
2431 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002432}
2433
2434static PyObject *
2435combinations_next(combinationsobject *co)
2436{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002437 PyObject *elem;
2438 PyObject *oldelem;
2439 PyObject *pool = co->pool;
2440 Py_ssize_t *indices = co->indices;
2441 PyObject *result = co->result;
2442 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2443 Py_ssize_t r = co->r;
2444 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002445
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002446 if (co->stopped)
2447 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002448
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002449 if (result == NULL) {
2450 /* On the first pass, initialize result tuple using the indices */
2451 result = PyTuple_New(r);
2452 if (result == NULL)
2453 goto empty;
2454 co->result = result;
2455 for (i=0; i<r ; i++) {
2456 index = indices[i];
2457 elem = PyTuple_GET_ITEM(pool, index);
2458 Py_INCREF(elem);
2459 PyTuple_SET_ITEM(result, i, elem);
2460 }
2461 } else {
2462 /* Copy the previous result tuple or re-use it if available */
2463 if (Py_REFCNT(result) > 1) {
2464 PyObject *old_result = result;
2465 result = PyTuple_New(r);
2466 if (result == NULL)
2467 goto empty;
2468 co->result = result;
2469 for (i=0; i<r ; i++) {
2470 elem = PyTuple_GET_ITEM(old_result, i);
2471 Py_INCREF(elem);
2472 PyTuple_SET_ITEM(result, i, elem);
2473 }
2474 Py_DECREF(old_result);
2475 }
2476 /* Now, we've got the only copy so we can update it in-place
2477 * CPython's empty tuple is a singleton and cached in
2478 * PyTuple's freelist.
2479 */
2480 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002481
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002482 /* Scan indices right-to-left until finding one that is not
2483 at its maximum (i + n - r). */
2484 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2485 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002486
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002487 /* If i is negative, then the indices are all at
2488 their maximum value and we're done. */
2489 if (i < 0)
2490 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002491
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002492 /* Increment the current index which we know is not at its
2493 maximum. Then move back to the right setting each index
2494 to its lowest possible value (one higher than the index
2495 to its left -- this maintains the sort order invariant). */
2496 indices[i]++;
2497 for (j=i+1 ; j<r ; j++)
2498 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002499
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002500 /* Update the result tuple for the new indices
2501 starting with i, the leftmost index that changed */
2502 for ( ; i<r ; i++) {
2503 index = indices[i];
2504 elem = PyTuple_GET_ITEM(pool, index);
2505 Py_INCREF(elem);
2506 oldelem = PyTuple_GET_ITEM(result, i);
2507 PyTuple_SET_ITEM(result, i, elem);
2508 Py_DECREF(oldelem);
2509 }
2510 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002511
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002512 Py_INCREF(result);
2513 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002514
2515empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002516 co->stopped = 1;
2517 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002518}
2519
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002520static PyObject *
2521combinations_reduce(combinationsobject *lz)
2522{
2523 if (lz->result == NULL) {
2524 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2525 } else if (lz->stopped) {
2526 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2527 } else {
2528 PyObject *indices;
2529 Py_ssize_t i;
2530
2531 /* we must pickle the indices and use them for setstate */
2532 indices = PyTuple_New(lz->r);
2533 if (!indices)
2534 return NULL;
2535 for (i=0; i<lz->r; i++)
2536 {
2537 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2538 if (!index) {
2539 Py_DECREF(indices);
2540 return NULL;
2541 }
2542 PyTuple_SET_ITEM(indices, i, index);
2543 }
2544
2545 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2546 }
2547}
2548
2549static PyObject *
2550combinations_setstate(combinationsobject *lz, PyObject *state)
2551{
2552 PyObject *result;
2553 Py_ssize_t i;
2554 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2555
2556 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2557 {
2558 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2559 return NULL;
2560 }
2561
2562 for (i=0; i<lz->r; i++)
2563 {
2564 Py_ssize_t max;
2565 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2566 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2567 if (index == -1 && PyErr_Occurred())
2568 return NULL; /* not an integer */
2569 max = i + n - lz->r;
2570 /* clamp the index (beware of negative max) */
2571 if (index > max)
2572 index = max;
2573 if (index < 0)
2574 index = 0;
2575 lz->indices[i] = index;
2576 }
2577
2578 result = PyTuple_New(lz->r);
2579 if (result == NULL)
2580 return NULL;
2581 for (i=0; i<lz->r; i++) {
2582 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2583 Py_INCREF(element);
2584 PyTuple_SET_ITEM(result, i, element);
2585 }
2586
2587 Py_CLEAR(lz->result);
2588 lz->result = result;
2589 Py_RETURN_NONE;
2590}
2591
2592static PyMethodDef combinations_methods[] = {
2593 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2594 reduce_doc},
2595 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2596 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002597 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2598 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002599 {NULL, NULL} /* sentinel */
2600};
2601
Christian Heimes380f7f22008-02-28 11:19:05 +00002602PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002603"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002604\n\
2605Return successive r-length combinations of elements in the iterable.\n\n\
2606combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2607
2608static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002609 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002610 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002611 sizeof(combinationsobject), /* tp_basicsize */
2612 0, /* tp_itemsize */
2613 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002614 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002615 0, /* tp_print */
2616 0, /* tp_getattr */
2617 0, /* tp_setattr */
2618 0, /* tp_reserved */
2619 0, /* tp_repr */
2620 0, /* tp_as_number */
2621 0, /* tp_as_sequence */
2622 0, /* tp_as_mapping */
2623 0, /* tp_hash */
2624 0, /* tp_call */
2625 0, /* tp_str */
2626 PyObject_GenericGetAttr, /* tp_getattro */
2627 0, /* tp_setattro */
2628 0, /* tp_as_buffer */
2629 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2630 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002631 combinations_doc, /* tp_doc */
2632 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002633 0, /* tp_clear */
2634 0, /* tp_richcompare */
2635 0, /* tp_weaklistoffset */
2636 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002637 (iternextfunc)combinations_next, /* tp_iternext */
2638 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002639 0, /* tp_members */
2640 0, /* tp_getset */
2641 0, /* tp_base */
2642 0, /* tp_dict */
2643 0, /* tp_descr_get */
2644 0, /* tp_descr_set */
2645 0, /* tp_dictoffset */
2646 0, /* tp_init */
2647 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002648 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002649 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002650};
2651
2652
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002653/* combinations with replacement object *******************************************/
2654
2655/* Equivalent to:
2656
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002657 def combinations_with_replacement(iterable, r):
2658 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2659 # number items returned: (n+r-1)! / r! / (n-1)!
2660 pool = tuple(iterable)
2661 n = len(pool)
2662 indices = [0] * r
2663 yield tuple(pool[i] for i in indices)
2664 while 1:
2665 for i in reversed(range(r)):
2666 if indices[i] != n - 1:
2667 break
2668 else:
2669 return
2670 indices[i:] = [indices[i] + 1] * (r - i)
2671 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002672
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002673 def combinations_with_replacement2(iterable, r):
2674 'Alternate version that filters from product()'
2675 pool = tuple(iterable)
2676 n = len(pool)
2677 for indices in product(range(n), repeat=r):
2678 if sorted(indices) == list(indices):
2679 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002680*/
2681typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002682 PyObject_HEAD
2683 PyObject *pool; /* input converted to a tuple */
2684 Py_ssize_t *indices; /* one index per result element */
2685 PyObject *result; /* most recently returned result tuple */
2686 Py_ssize_t r; /* size of result tuple */
2687 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002688} cwrobject;
2689
2690static PyTypeObject cwr_type;
2691
2692static PyObject *
2693cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2694{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002695 cwrobject *co;
2696 Py_ssize_t n;
2697 Py_ssize_t r;
2698 PyObject *pool = NULL;
2699 PyObject *iterable = NULL;
2700 Py_ssize_t *indices = NULL;
2701 Py_ssize_t i;
2702 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002703
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002704 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2705 &iterable, &r))
2706 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002707
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002708 pool = PySequence_Tuple(iterable);
2709 if (pool == NULL)
2710 goto error;
2711 n = PyTuple_GET_SIZE(pool);
2712 if (r < 0) {
2713 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2714 goto error;
2715 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002716
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002717 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002718 if (indices == NULL) {
2719 PyErr_NoMemory();
2720 goto error;
2721 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002722
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002723 for (i=0 ; i<r ; i++)
2724 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002725
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002726 /* create cwrobject structure */
2727 co = (cwrobject *)type->tp_alloc(type, 0);
2728 if (co == NULL)
2729 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002730
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002731 co->pool = pool;
2732 co->indices = indices;
2733 co->result = NULL;
2734 co->r = r;
2735 co->stopped = !n && r;
2736
2737 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002738
2739error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002740 if (indices != NULL)
2741 PyMem_Free(indices);
2742 Py_XDECREF(pool);
2743 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002744}
2745
2746static void
2747cwr_dealloc(cwrobject *co)
2748{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002749 PyObject_GC_UnTrack(co);
2750 Py_XDECREF(co->pool);
2751 Py_XDECREF(co->result);
2752 if (co->indices != NULL)
2753 PyMem_Free(co->indices);
2754 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002755}
2756
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002757static PyObject *
2758cwr_sizeof(cwrobject *co, void *unused)
2759{
2760 Py_ssize_t res;
2761
2762 res = sizeof(cwrobject);
2763 res += co->r * sizeof(Py_ssize_t);
2764 return PyLong_FromSsize_t(res);
2765}
2766
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002767static int
2768cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2769{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002770 Py_VISIT(co->pool);
2771 Py_VISIT(co->result);
2772 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002773}
2774
2775static PyObject *
2776cwr_next(cwrobject *co)
2777{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002778 PyObject *elem;
2779 PyObject *oldelem;
2780 PyObject *pool = co->pool;
2781 Py_ssize_t *indices = co->indices;
2782 PyObject *result = co->result;
2783 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2784 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002785 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002786
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002787 if (co->stopped)
2788 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002789
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002790 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002791 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002792 result = PyTuple_New(r);
2793 if (result == NULL)
2794 goto empty;
2795 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002796 if (n > 0) {
2797 elem = PyTuple_GET_ITEM(pool, 0);
2798 for (i=0; i<r ; i++) {
2799 assert(indices[i] == 0);
2800 Py_INCREF(elem);
2801 PyTuple_SET_ITEM(result, i, elem);
2802 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002803 }
2804 } else {
2805 /* Copy the previous result tuple or re-use it if available */
2806 if (Py_REFCNT(result) > 1) {
2807 PyObject *old_result = result;
2808 result = PyTuple_New(r);
2809 if (result == NULL)
2810 goto empty;
2811 co->result = result;
2812 for (i=0; i<r ; i++) {
2813 elem = PyTuple_GET_ITEM(old_result, i);
2814 Py_INCREF(elem);
2815 PyTuple_SET_ITEM(result, i, elem);
2816 }
2817 Py_DECREF(old_result);
2818 }
2819 /* Now, we've got the only copy so we can update it in-place CPython's
2820 empty tuple is a singleton and cached in PyTuple's freelist. */
2821 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002822
Tim Peters9edb1682013-09-03 11:49:31 -05002823 /* Scan indices right-to-left until finding one that is not
2824 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002825 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2826 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002827
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002828 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002829 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002830 if (i < 0)
2831 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002832
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002833 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002834 maximum. Then set all to the right to the same value. */
2835 index = indices[i] + 1;
2836 assert(index < n);
2837 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002838 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002839 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002840 Py_INCREF(elem);
2841 oldelem = PyTuple_GET_ITEM(result, i);
2842 PyTuple_SET_ITEM(result, i, elem);
2843 Py_DECREF(oldelem);
2844 }
2845 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002846
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002847 Py_INCREF(result);
2848 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002849
2850empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002851 co->stopped = 1;
2852 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002853}
2854
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002855static PyObject *
2856cwr_reduce(cwrobject *lz)
2857{
2858 if (lz->result == NULL) {
2859 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2860 } else if (lz->stopped) {
2861 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2862 } else {
2863 PyObject *indices;
2864 Py_ssize_t i;
2865
2866 /* we must pickle the indices and use them for setstate */
2867 indices = PyTuple_New(lz->r);
2868 if (!indices)
2869 return NULL;
2870 for (i=0; i<lz->r; i++)
2871 {
2872 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2873 if (!index) {
2874 Py_DECREF(indices);
2875 return NULL;
2876 }
2877 PyTuple_SET_ITEM(indices, i, index);
2878 }
2879
2880 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2881 }
2882}
2883
2884static PyObject *
2885cwr_setstate(cwrobject *lz, PyObject *state)
2886{
2887 PyObject *result;
2888 Py_ssize_t n, i;
2889
2890 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2891 {
2892 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2893 return NULL;
2894 }
2895
2896 n = PyTuple_GET_SIZE(lz->pool);
2897 for (i=0; i<lz->r; i++)
2898 {
2899 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2900 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2901 if (index < 0 && PyErr_Occurred())
2902 return NULL; /* not an integer */
2903 /* clamp the index */
2904 if (index < 0)
2905 index = 0;
2906 else if (index > n-1)
2907 index = n-1;
2908 lz->indices[i] = index;
2909 }
2910 result = PyTuple_New(lz->r);
2911 if (result == NULL)
2912 return NULL;
2913 for (i=0; i<lz->r; i++) {
2914 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2915 Py_INCREF(element);
2916 PyTuple_SET_ITEM(result, i, element);
2917 }
2918 Py_CLEAR(lz->result);
2919 lz->result = result;
2920 Py_RETURN_NONE;
2921}
2922
2923static PyMethodDef cwr_methods[] = {
2924 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2925 reduce_doc},
2926 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2927 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002928 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
2929 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002930 {NULL, NULL} /* sentinel */
2931};
2932
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002933PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002934"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002935\n\
2936Return successive r-length combinations of elements in the iterable\n\
2937allowing individual elements to have successive repeats.\n\
2938combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2939
2940static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002941 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002942 "itertools.combinations_with_replacement", /* tp_name */
2943 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002944 0, /* tp_itemsize */
2945 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002946 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002947 0, /* tp_print */
2948 0, /* tp_getattr */
2949 0, /* tp_setattr */
2950 0, /* tp_reserved */
2951 0, /* tp_repr */
2952 0, /* tp_as_number */
2953 0, /* tp_as_sequence */
2954 0, /* tp_as_mapping */
2955 0, /* tp_hash */
2956 0, /* tp_call */
2957 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002958 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002959 0, /* tp_setattro */
2960 0, /* tp_as_buffer */
2961 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002962 Py_TPFLAGS_BASETYPE, /* tp_flags */
2963 cwr_doc, /* tp_doc */
2964 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002965 0, /* tp_clear */
2966 0, /* tp_richcompare */
2967 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002968 PyObject_SelfIter, /* tp_iter */
2969 (iternextfunc)cwr_next, /* tp_iternext */
2970 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002971 0, /* tp_members */
2972 0, /* tp_getset */
2973 0, /* tp_base */
2974 0, /* tp_dict */
2975 0, /* tp_descr_get */
2976 0, /* tp_descr_set */
2977 0, /* tp_dictoffset */
2978 0, /* tp_init */
2979 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002980 cwr_new, /* tp_new */
2981 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002982};
2983
2984
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002985/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002986
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002987def permutations(iterable, r=None):
2988 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2989 pool = tuple(iterable)
2990 n = len(pool)
2991 r = n if r is None else r
2992 indices = range(n)
2993 cycles = range(n-r+1, n+1)[::-1]
2994 yield tuple(pool[i] for i in indices[:r])
2995 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03002996 for i in reversed(range(r)):
2997 cycles[i] -= 1
2998 if cycles[i] == 0:
2999 indices[i:] = indices[i+1:] + indices[i:i+1]
3000 cycles[i] = n - i
3001 else:
3002 j = cycles[i]
3003 indices[i], indices[-j] = indices[-j], indices[i]
3004 yield tuple(pool[i] for i in indices[:r])
3005 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003006 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003007 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003008*/
3009
3010typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003011 PyObject_HEAD
3012 PyObject *pool; /* input converted to a tuple */
3013 Py_ssize_t *indices; /* one index per element in the pool */
3014 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3015 PyObject *result; /* most recently returned result tuple */
3016 Py_ssize_t r; /* size of result tuple */
3017 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003018} permutationsobject;
3019
3020static PyTypeObject permutations_type;
3021
3022static PyObject *
3023permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3024{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003025 permutationsobject *po;
3026 Py_ssize_t n;
3027 Py_ssize_t r;
3028 PyObject *robj = Py_None;
3029 PyObject *pool = NULL;
3030 PyObject *iterable = NULL;
3031 Py_ssize_t *indices = NULL;
3032 Py_ssize_t *cycles = NULL;
3033 Py_ssize_t i;
3034 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003035
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003036 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
3037 &iterable, &robj))
3038 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003039
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003040 pool = PySequence_Tuple(iterable);
3041 if (pool == NULL)
3042 goto error;
3043 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003044
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003045 r = n;
3046 if (robj != Py_None) {
3047 if (!PyLong_Check(robj)) {
3048 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3049 goto error;
3050 }
3051 r = PyLong_AsSsize_t(robj);
3052 if (r == -1 && PyErr_Occurred())
3053 goto error;
3054 }
3055 if (r < 0) {
3056 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3057 goto error;
3058 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003059
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003060 indices = PyMem_New(Py_ssize_t, n);
3061 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003062 if (indices == NULL || cycles == NULL) {
3063 PyErr_NoMemory();
3064 goto error;
3065 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003066
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003067 for (i=0 ; i<n ; i++)
3068 indices[i] = i;
3069 for (i=0 ; i<r ; i++)
3070 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003071
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003072 /* create permutationsobject structure */
3073 po = (permutationsobject *)type->tp_alloc(type, 0);
3074 if (po == NULL)
3075 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003076
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003077 po->pool = pool;
3078 po->indices = indices;
3079 po->cycles = cycles;
3080 po->result = NULL;
3081 po->r = r;
3082 po->stopped = r > n ? 1 : 0;
3083
3084 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003085
3086error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003087 if (indices != NULL)
3088 PyMem_Free(indices);
3089 if (cycles != NULL)
3090 PyMem_Free(cycles);
3091 Py_XDECREF(pool);
3092 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003093}
3094
3095static void
3096permutations_dealloc(permutationsobject *po)
3097{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003098 PyObject_GC_UnTrack(po);
3099 Py_XDECREF(po->pool);
3100 Py_XDECREF(po->result);
3101 PyMem_Free(po->indices);
3102 PyMem_Free(po->cycles);
3103 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003104}
3105
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003106static PyObject *
3107permutations_sizeof(permutationsobject *po, void *unused)
3108{
3109 Py_ssize_t res;
3110
3111 res = sizeof(permutationsobject);
3112 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3113 res += po->r * sizeof(Py_ssize_t);
3114 return PyLong_FromSsize_t(res);
3115}
3116
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003117static int
3118permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3119{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003120 Py_VISIT(po->pool);
3121 Py_VISIT(po->result);
3122 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003123}
3124
3125static PyObject *
3126permutations_next(permutationsobject *po)
3127{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003128 PyObject *elem;
3129 PyObject *oldelem;
3130 PyObject *pool = po->pool;
3131 Py_ssize_t *indices = po->indices;
3132 Py_ssize_t *cycles = po->cycles;
3133 PyObject *result = po->result;
3134 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3135 Py_ssize_t r = po->r;
3136 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003137
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003138 if (po->stopped)
3139 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003140
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003141 if (result == NULL) {
3142 /* On the first pass, initialize result tuple using the indices */
3143 result = PyTuple_New(r);
3144 if (result == NULL)
3145 goto empty;
3146 po->result = result;
3147 for (i=0; i<r ; i++) {
3148 index = indices[i];
3149 elem = PyTuple_GET_ITEM(pool, index);
3150 Py_INCREF(elem);
3151 PyTuple_SET_ITEM(result, i, elem);
3152 }
3153 } else {
3154 if (n == 0)
3155 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003156
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003157 /* Copy the previous result tuple or re-use it if available */
3158 if (Py_REFCNT(result) > 1) {
3159 PyObject *old_result = result;
3160 result = PyTuple_New(r);
3161 if (result == NULL)
3162 goto empty;
3163 po->result = result;
3164 for (i=0; i<r ; i++) {
3165 elem = PyTuple_GET_ITEM(old_result, i);
3166 Py_INCREF(elem);
3167 PyTuple_SET_ITEM(result, i, elem);
3168 }
3169 Py_DECREF(old_result);
3170 }
3171 /* Now, we've got the only copy so we can update it in-place */
3172 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003173
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003174 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3175 for (i=r-1 ; i>=0 ; i--) {
3176 cycles[i] -= 1;
3177 if (cycles[i] == 0) {
3178 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3179 index = indices[i];
3180 for (j=i ; j<n-1 ; j++)
3181 indices[j] = indices[j+1];
3182 indices[n-1] = index;
3183 cycles[i] = n - i;
3184 } else {
3185 j = cycles[i];
3186 index = indices[i];
3187 indices[i] = indices[n-j];
3188 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003189
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003190 for (k=i; k<r ; k++) {
3191 /* start with i, the leftmost element that changed */
3192 /* yield tuple(pool[k] for k in indices[:r]) */
3193 index = indices[k];
3194 elem = PyTuple_GET_ITEM(pool, index);
3195 Py_INCREF(elem);
3196 oldelem = PyTuple_GET_ITEM(result, k);
3197 PyTuple_SET_ITEM(result, k, elem);
3198 Py_DECREF(oldelem);
3199 }
3200 break;
3201 }
3202 }
3203 /* If i is negative, then the cycles have all
3204 rolled-over and we're done. */
3205 if (i < 0)
3206 goto empty;
3207 }
3208 Py_INCREF(result);
3209 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003210
3211empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003212 po->stopped = 1;
3213 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003214}
3215
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003216static PyObject *
3217permutations_reduce(permutationsobject *po)
3218{
3219 if (po->result == NULL) {
3220 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3221 } else if (po->stopped) {
3222 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3223 } else {
3224 PyObject *indices=NULL, *cycles=NULL;
3225 Py_ssize_t n, i;
3226
3227 /* we must pickle the indices and cycles and use them for setstate */
3228 n = PyTuple_GET_SIZE(po->pool);
3229 indices = PyTuple_New(n);
3230 if (indices == NULL)
3231 goto err;
3232 for (i=0; i<n; i++){
3233 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3234 if (!index)
3235 goto err;
3236 PyTuple_SET_ITEM(indices, i, index);
3237 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003238
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003239 cycles = PyTuple_New(po->r);
3240 if (cycles == NULL)
3241 goto err;
3242 for (i=0; i<po->r; i++)
3243 {
3244 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3245 if (!index)
3246 goto err;
3247 PyTuple_SET_ITEM(cycles, i, index);
3248 }
3249 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3250 po->pool, po->r,
3251 indices, cycles);
3252 err:
3253 Py_XDECREF(indices);
3254 Py_XDECREF(cycles);
3255 return NULL;
3256 }
3257}
3258
3259static PyObject *
3260permutations_setstate(permutationsobject *po, PyObject *state)
3261{
3262 PyObject *indices, *cycles, *result;
3263 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003264
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003265 if (!PyArg_ParseTuple(state, "O!O!",
3266 &PyTuple_Type, &indices,
3267 &PyTuple_Type, &cycles))
3268 return NULL;
3269
3270 n = PyTuple_GET_SIZE(po->pool);
3271 if (PyTuple_GET_SIZE(indices) != n ||
3272 PyTuple_GET_SIZE(cycles) != po->r)
3273 {
3274 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3275 return NULL;
3276 }
3277
3278 for (i=0; i<n; i++)
3279 {
3280 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3281 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3282 if (index < 0 && PyErr_Occurred())
3283 return NULL; /* not an integer */
3284 /* clamp the index */
3285 if (index < 0)
3286 index = 0;
3287 else if (index > n-1)
3288 index = n-1;
3289 po->indices[i] = index;
3290 }
3291
3292 for (i=0; i<po->r; i++)
3293 {
3294 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3295 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3296 if (index < 0 && PyErr_Occurred())
3297 return NULL; /* not an integer */
3298 if (index < 1)
3299 index = 1;
3300 else if (index > n-i)
3301 index = n-i;
3302 po->cycles[i] = index;
3303 }
3304 result = PyTuple_New(po->r);
3305 if (result == NULL)
3306 return NULL;
3307 for (i=0; i<po->r; i++) {
3308 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3309 Py_INCREF(element);
3310 PyTuple_SET_ITEM(result, i, element);
3311 }
3312 Py_CLEAR(po->result);
3313 po->result = result;
3314 Py_RETURN_NONE;
3315}
3316
3317static PyMethodDef permuations_methods[] = {
3318 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3319 reduce_doc},
3320 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3321 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003322 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3323 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003324 {NULL, NULL} /* sentinel */
3325};
3326
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003327PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003328"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003329\n\
3330Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003331permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003332
3333static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003334 PyVarObject_HEAD_INIT(NULL, 0)
3335 "itertools.permutations", /* tp_name */
3336 sizeof(permutationsobject), /* tp_basicsize */
3337 0, /* tp_itemsize */
3338 /* methods */
3339 (destructor)permutations_dealloc, /* tp_dealloc */
3340 0, /* tp_print */
3341 0, /* tp_getattr */
3342 0, /* tp_setattr */
3343 0, /* tp_reserved */
3344 0, /* tp_repr */
3345 0, /* tp_as_number */
3346 0, /* tp_as_sequence */
3347 0, /* tp_as_mapping */
3348 0, /* tp_hash */
3349 0, /* tp_call */
3350 0, /* tp_str */
3351 PyObject_GenericGetAttr, /* tp_getattro */
3352 0, /* tp_setattro */
3353 0, /* tp_as_buffer */
3354 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3355 Py_TPFLAGS_BASETYPE, /* tp_flags */
3356 permutations_doc, /* tp_doc */
3357 (traverseproc)permutations_traverse, /* tp_traverse */
3358 0, /* tp_clear */
3359 0, /* tp_richcompare */
3360 0, /* tp_weaklistoffset */
3361 PyObject_SelfIter, /* tp_iter */
3362 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003363 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003364 0, /* tp_members */
3365 0, /* tp_getset */
3366 0, /* tp_base */
3367 0, /* tp_dict */
3368 0, /* tp_descr_get */
3369 0, /* tp_descr_set */
3370 0, /* tp_dictoffset */
3371 0, /* tp_init */
3372 0, /* tp_alloc */
3373 permutations_new, /* tp_new */
3374 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003375};
3376
Raymond Hettinger482ba772010-12-01 22:48:00 +00003377/* accumulate object ************************************************************/
3378
3379typedef struct {
3380 PyObject_HEAD
3381 PyObject *total;
3382 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003383 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003384} accumulateobject;
3385
3386static PyTypeObject accumulate_type;
3387
3388static PyObject *
3389accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3390{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003391 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003392 PyObject *iterable;
3393 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003394 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003395 accumulateobject *lz;
3396
Raymond Hettinger5d446132011-03-27 18:52:10 -07003397 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3398 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003399 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003400
3401 /* Get iterator. */
3402 it = PyObject_GetIter(iterable);
3403 if (it == NULL)
3404 return NULL;
3405
Raymond Hettinger482ba772010-12-01 22:48:00 +00003406 /* create accumulateobject structure */
3407 lz = (accumulateobject *)type->tp_alloc(type, 0);
3408 if (lz == NULL) {
3409 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003410 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003411 }
3412
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003413 if (binop != Py_None) {
3414 Py_XINCREF(binop);
3415 lz->binop = binop;
3416 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003417 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003418 lz->it = it;
3419 return (PyObject *)lz;
3420}
3421
3422static void
3423accumulate_dealloc(accumulateobject *lz)
3424{
3425 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003426 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003427 Py_XDECREF(lz->total);
3428 Py_XDECREF(lz->it);
3429 Py_TYPE(lz)->tp_free(lz);
3430}
3431
3432static int
3433accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3434{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003435 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003436 Py_VISIT(lz->it);
3437 Py_VISIT(lz->total);
3438 return 0;
3439}
3440
3441static PyObject *
3442accumulate_next(accumulateobject *lz)
3443{
3444 PyObject *val, *oldtotal, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003445
Raymond Hettinger482ba772010-12-01 22:48:00 +00003446 val = PyIter_Next(lz->it);
3447 if (val == NULL)
3448 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003449
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003450 if (lz->total == NULL) {
3451 Py_INCREF(val);
3452 lz->total = val;
3453 return lz->total;
3454 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003455
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003456 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003457 newtotal = PyNumber_Add(lz->total, val);
3458 else
3459 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003460 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003461 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003462 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003463
3464 oldtotal = lz->total;
3465 lz->total = newtotal;
3466 Py_DECREF(oldtotal);
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003467
Raymond Hettinger482ba772010-12-01 22:48:00 +00003468 Py_INCREF(newtotal);
3469 return newtotal;
3470}
3471
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003472static PyObject *
3473accumulate_reduce(accumulateobject *lz)
3474{
3475 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3476 lz->it, lz->binop?lz->binop:Py_None,
3477 lz->total?lz->total:Py_None);
3478 }
3479
3480static PyObject *
3481accumulate_setstate(accumulateobject *lz, PyObject *state)
3482{
3483 Py_CLEAR(lz->total);
3484 lz->total = state;
3485 Py_INCREF(lz->total);
3486 Py_RETURN_NONE;
3487}
3488
3489static PyMethodDef accumulate_methods[] = {
3490 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3491 reduce_doc},
3492 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3493 setstate_doc},
3494 {NULL, NULL} /* sentinel */
3495};
3496
Raymond Hettinger482ba772010-12-01 22:48:00 +00003497PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003498"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003499\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003500Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003501
3502static PyTypeObject accumulate_type = {
3503 PyVarObject_HEAD_INIT(NULL, 0)
3504 "itertools.accumulate", /* tp_name */
3505 sizeof(accumulateobject), /* tp_basicsize */
3506 0, /* tp_itemsize */
3507 /* methods */
3508 (destructor)accumulate_dealloc, /* tp_dealloc */
3509 0, /* tp_print */
3510 0, /* tp_getattr */
3511 0, /* tp_setattr */
3512 0, /* tp_reserved */
3513 0, /* tp_repr */
3514 0, /* tp_as_number */
3515 0, /* tp_as_sequence */
3516 0, /* tp_as_mapping */
3517 0, /* tp_hash */
3518 0, /* tp_call */
3519 0, /* tp_str */
3520 PyObject_GenericGetAttr, /* tp_getattro */
3521 0, /* tp_setattro */
3522 0, /* tp_as_buffer */
3523 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3524 Py_TPFLAGS_BASETYPE, /* tp_flags */
3525 accumulate_doc, /* tp_doc */
3526 (traverseproc)accumulate_traverse, /* tp_traverse */
3527 0, /* tp_clear */
3528 0, /* tp_richcompare */
3529 0, /* tp_weaklistoffset */
3530 PyObject_SelfIter, /* tp_iter */
3531 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003532 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003533 0, /* tp_members */
3534 0, /* tp_getset */
3535 0, /* tp_base */
3536 0, /* tp_dict */
3537 0, /* tp_descr_get */
3538 0, /* tp_descr_set */
3539 0, /* tp_dictoffset */
3540 0, /* tp_init */
3541 0, /* tp_alloc */
3542 accumulate_new, /* tp_new */
3543 PyObject_GC_Del, /* tp_free */
3544};
3545
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003546
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003547/* compress object ************************************************************/
3548
3549/* Equivalent to:
3550
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003551 def compress(data, selectors):
3552 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3553 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003554*/
3555
3556typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003557 PyObject_HEAD
3558 PyObject *data;
3559 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003560} compressobject;
3561
3562static PyTypeObject compress_type;
3563
3564static PyObject *
3565compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3566{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003567 PyObject *seq1, *seq2;
3568 PyObject *data=NULL, *selectors=NULL;
3569 compressobject *lz;
3570 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003571
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003572 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3573 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003574
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003575 data = PyObject_GetIter(seq1);
3576 if (data == NULL)
3577 goto fail;
3578 selectors = PyObject_GetIter(seq2);
3579 if (selectors == NULL)
3580 goto fail;
3581
3582 /* create compressobject structure */
3583 lz = (compressobject *)type->tp_alloc(type, 0);
3584 if (lz == NULL)
3585 goto fail;
3586 lz->data = data;
3587 lz->selectors = selectors;
3588 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003589
3590fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003591 Py_XDECREF(data);
3592 Py_XDECREF(selectors);
3593 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003594}
3595
3596static void
3597compress_dealloc(compressobject *lz)
3598{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003599 PyObject_GC_UnTrack(lz);
3600 Py_XDECREF(lz->data);
3601 Py_XDECREF(lz->selectors);
3602 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003603}
3604
3605static int
3606compress_traverse(compressobject *lz, visitproc visit, void *arg)
3607{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003608 Py_VISIT(lz->data);
3609 Py_VISIT(lz->selectors);
3610 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003611}
3612
3613static PyObject *
3614compress_next(compressobject *lz)
3615{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003616 PyObject *data = lz->data, *selectors = lz->selectors;
3617 PyObject *datum, *selector;
3618 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3619 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3620 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003621
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003622 while (1) {
3623 /* Steps: get datum, get selector, evaluate selector.
3624 Order is important (to match the pure python version
3625 in terms of which input gets a chance to raise an
3626 exception first).
3627 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003628
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003629 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003630 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003631 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003632
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003633 selector = selectornext(selectors);
3634 if (selector == NULL) {
3635 Py_DECREF(datum);
3636 return NULL;
3637 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003638
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003639 ok = PyObject_IsTrue(selector);
3640 Py_DECREF(selector);
3641 if (ok == 1)
3642 return datum;
3643 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003644 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003645 return NULL;
3646 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003647}
3648
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003649static PyObject *
3650compress_reduce(compressobject *lz)
3651{
3652 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3653 lz->data, lz->selectors);
3654 }
3655
3656static PyMethodDef compress_methods[] = {
3657 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3658 reduce_doc},
3659 {NULL, NULL} /* sentinel */
3660};
3661
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003662PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003663"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003664\n\
3665Return data elements corresponding to true selector elements.\n\
3666Forms a shorter iterator from selected data elements using the\n\
3667selectors to choose the data elements.");
3668
3669static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003670 PyVarObject_HEAD_INIT(NULL, 0)
3671 "itertools.compress", /* tp_name */
3672 sizeof(compressobject), /* tp_basicsize */
3673 0, /* tp_itemsize */
3674 /* methods */
3675 (destructor)compress_dealloc, /* tp_dealloc */
3676 0, /* tp_print */
3677 0, /* tp_getattr */
3678 0, /* tp_setattr */
3679 0, /* tp_reserved */
3680 0, /* tp_repr */
3681 0, /* tp_as_number */
3682 0, /* tp_as_sequence */
3683 0, /* tp_as_mapping */
3684 0, /* tp_hash */
3685 0, /* tp_call */
3686 0, /* tp_str */
3687 PyObject_GenericGetAttr, /* tp_getattro */
3688 0, /* tp_setattro */
3689 0, /* tp_as_buffer */
3690 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3691 Py_TPFLAGS_BASETYPE, /* tp_flags */
3692 compress_doc, /* tp_doc */
3693 (traverseproc)compress_traverse, /* tp_traverse */
3694 0, /* tp_clear */
3695 0, /* tp_richcompare */
3696 0, /* tp_weaklistoffset */
3697 PyObject_SelfIter, /* tp_iter */
3698 (iternextfunc)compress_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003699 compress_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003700 0, /* tp_members */
3701 0, /* tp_getset */
3702 0, /* tp_base */
3703 0, /* tp_dict */
3704 0, /* tp_descr_get */
3705 0, /* tp_descr_set */
3706 0, /* tp_dictoffset */
3707 0, /* tp_init */
3708 0, /* tp_alloc */
3709 compress_new, /* tp_new */
3710 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003711};
3712
3713
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003714/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003715
3716typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003717 PyObject_HEAD
3718 PyObject *func;
3719 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003720} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003721
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003722static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003723
3724static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003725filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003726{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003727 PyObject *func, *seq;
3728 PyObject *it;
3729 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003730
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003731 if (type == &filterfalse_type &&
3732 !_PyArg_NoKeywords("filterfalse()", kwds))
3733 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003734
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003735 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3736 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003737
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003738 /* Get iterator. */
3739 it = PyObject_GetIter(seq);
3740 if (it == NULL)
3741 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003742
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003743 /* create filterfalseobject structure */
3744 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3745 if (lz == NULL) {
3746 Py_DECREF(it);
3747 return NULL;
3748 }
3749 Py_INCREF(func);
3750 lz->func = func;
3751 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003752
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003753 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003754}
3755
3756static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003757filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003758{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003759 PyObject_GC_UnTrack(lz);
3760 Py_XDECREF(lz->func);
3761 Py_XDECREF(lz->it);
3762 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003763}
3764
3765static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003766filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003767{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003768 Py_VISIT(lz->it);
3769 Py_VISIT(lz->func);
3770 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003771}
3772
3773static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003774filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003775{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003776 PyObject *item;
3777 PyObject *it = lz->it;
3778 long ok;
3779 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003780
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003781 iternext = *Py_TYPE(it)->tp_iternext;
3782 for (;;) {
3783 item = iternext(it);
3784 if (item == NULL)
3785 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003786
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003787 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3788 ok = PyObject_IsTrue(item);
3789 } else {
3790 PyObject *good;
3791 good = PyObject_CallFunctionObjArgs(lz->func,
3792 item, NULL);
3793 if (good == NULL) {
3794 Py_DECREF(item);
3795 return NULL;
3796 }
3797 ok = PyObject_IsTrue(good);
3798 Py_DECREF(good);
3799 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003800 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003801 return item;
3802 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003803 if (ok < 0)
3804 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003805 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003806}
3807
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003808static PyObject *
3809filterfalse_reduce(filterfalseobject *lz)
3810{
3811 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3812 lz->func, lz->it);
3813 }
3814
3815static PyMethodDef filterfalse_methods[] = {
3816 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3817 reduce_doc},
3818 {NULL, NULL} /* sentinel */
3819};
3820
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003821PyDoc_STRVAR(filterfalse_doc,
3822"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003823\n\
3824Return those items of sequence for which function(item) is false.\n\
3825If function is None, return the items that are false.");
3826
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003827static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003828 PyVarObject_HEAD_INIT(NULL, 0)
3829 "itertools.filterfalse", /* tp_name */
3830 sizeof(filterfalseobject), /* tp_basicsize */
3831 0, /* tp_itemsize */
3832 /* methods */
3833 (destructor)filterfalse_dealloc, /* tp_dealloc */
3834 0, /* tp_print */
3835 0, /* tp_getattr */
3836 0, /* tp_setattr */
3837 0, /* tp_reserved */
3838 0, /* tp_repr */
3839 0, /* tp_as_number */
3840 0, /* tp_as_sequence */
3841 0, /* tp_as_mapping */
3842 0, /* tp_hash */
3843 0, /* tp_call */
3844 0, /* tp_str */
3845 PyObject_GenericGetAttr, /* tp_getattro */
3846 0, /* tp_setattro */
3847 0, /* tp_as_buffer */
3848 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3849 Py_TPFLAGS_BASETYPE, /* tp_flags */
3850 filterfalse_doc, /* tp_doc */
3851 (traverseproc)filterfalse_traverse, /* tp_traverse */
3852 0, /* tp_clear */
3853 0, /* tp_richcompare */
3854 0, /* tp_weaklistoffset */
3855 PyObject_SelfIter, /* tp_iter */
3856 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003857 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003858 0, /* tp_members */
3859 0, /* tp_getset */
3860 0, /* tp_base */
3861 0, /* tp_dict */
3862 0, /* tp_descr_get */
3863 0, /* tp_descr_set */
3864 0, /* tp_dictoffset */
3865 0, /* tp_init */
3866 0, /* tp_alloc */
3867 filterfalse_new, /* tp_new */
3868 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003869};
3870
3871
3872/* count object ************************************************************/
3873
3874typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003875 PyObject_HEAD
3876 Py_ssize_t cnt;
3877 PyObject *long_cnt;
3878 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003879} countobject;
3880
Raymond Hettinger30729212009-02-12 06:28:27 +00003881/* Counting logic and invariants:
3882
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003883fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003884
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003885 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3886 Advances with: cnt += 1
3887 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003888
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003889slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003890
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003891 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3892 All counting is done with python objects (no overflows or underflows).
3893 Advances with: long_cnt += long_step
3894 Step may be zero -- effectively a slow version of repeat(cnt).
3895 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003896*/
3897
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003898static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003899
3900static PyObject *
3901count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3902{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003903 countobject *lz;
3904 int slow_mode = 0;
3905 Py_ssize_t cnt = 0;
3906 PyObject *long_cnt = NULL;
3907 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003908 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003909 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003910
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003911 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3912 kwlist, &long_cnt, &long_step))
3913 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003914
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003915 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3916 (long_step != NULL && !PyNumber_Check(long_step))) {
3917 PyErr_SetString(PyExc_TypeError, "a number is required");
3918 return NULL;
3919 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003920
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003921 if (long_cnt != NULL) {
3922 cnt = PyLong_AsSsize_t(long_cnt);
3923 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3924 PyErr_Clear();
3925 slow_mode = 1;
3926 }
3927 Py_INCREF(long_cnt);
3928 } else {
3929 cnt = 0;
3930 long_cnt = PyLong_FromLong(0);
3931 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003932
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003933 /* If not specified, step defaults to 1 */
3934 if (long_step == NULL) {
3935 long_step = PyLong_FromLong(1);
3936 if (long_step == NULL) {
3937 Py_DECREF(long_cnt);
3938 return NULL;
3939 }
3940 } else
3941 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003942
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003943 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003944
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003945 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003946 step = PyLong_AsLong(long_step);
3947 if (step != 1) {
3948 slow_mode = 1;
3949 if (step == -1 && PyErr_Occurred())
3950 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003951 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003952
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003953 if (slow_mode)
3954 cnt = PY_SSIZE_T_MAX;
3955 else
3956 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003957
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003958 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3959 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3960 assert(slow_mode ||
3961 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003962
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003963 /* create countobject structure */
3964 lz = (countobject *)type->tp_alloc(type, 0);
3965 if (lz == NULL) {
3966 Py_XDECREF(long_cnt);
3967 return NULL;
3968 }
3969 lz->cnt = cnt;
3970 lz->long_cnt = long_cnt;
3971 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003972
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003973 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003974}
3975
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003976static void
3977count_dealloc(countobject *lz)
3978{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003979 PyObject_GC_UnTrack(lz);
3980 Py_XDECREF(lz->long_cnt);
3981 Py_XDECREF(lz->long_step);
3982 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003983}
3984
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003985static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003986count_traverse(countobject *lz, visitproc visit, void *arg)
3987{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003988 Py_VISIT(lz->long_cnt);
3989 Py_VISIT(lz->long_step);
3990 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003991}
3992
3993static PyObject *
3994count_nextlong(countobject *lz)
3995{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003996 PyObject *long_cnt;
3997 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003998
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003999 long_cnt = lz->long_cnt;
4000 if (long_cnt == NULL) {
4001 /* Switch to slow_mode */
4002 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4003 if (long_cnt == NULL)
4004 return NULL;
4005 }
4006 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004007
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004008 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4009 if (stepped_up == NULL)
4010 return NULL;
4011 lz->long_cnt = stepped_up;
4012 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004013}
4014
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004015static PyObject *
4016count_next(countobject *lz)
4017{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004018 if (lz->cnt == PY_SSIZE_T_MAX)
4019 return count_nextlong(lz);
4020 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004021}
4022
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004023static PyObject *
4024count_repr(countobject *lz)
4025{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004026 if (lz->cnt != PY_SSIZE_T_MAX)
4027 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004028
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004029 if (PyLong_Check(lz->long_step)) {
4030 long step = PyLong_AsLong(lz->long_step);
4031 if (step == -1 && PyErr_Occurred()) {
4032 PyErr_Clear();
4033 }
4034 if (step == 1) {
4035 /* Don't display step when it is an integer equal to 1 */
4036 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
4037 }
4038 }
4039 return PyUnicode_FromFormat("count(%R, %R)",
4040 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004041}
4042
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004043static PyObject *
4044count_reduce(countobject *lz)
4045{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004046 if (lz->cnt == PY_SSIZE_T_MAX)
4047 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4048 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004049}
4050
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004051static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004052 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004053 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004054 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004055};
4056
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004057PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00004058 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004059\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004060Return a count object whose .__next__() method returns consecutive values.\n\
4061Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00004062 def count(firstval=0, step=1):\n\
Eli Benderskyc554f722013-09-03 06:37:19 -07004063 x = firstval\n\
4064 while 1:\n\
4065 yield x\n\
4066 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004067
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004068static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004069 PyVarObject_HEAD_INIT(NULL, 0)
4070 "itertools.count", /* tp_name */
4071 sizeof(countobject), /* tp_basicsize */
4072 0, /* tp_itemsize */
4073 /* methods */
4074 (destructor)count_dealloc, /* tp_dealloc */
4075 0, /* tp_print */
4076 0, /* tp_getattr */
4077 0, /* tp_setattr */
4078 0, /* tp_reserved */
4079 (reprfunc)count_repr, /* tp_repr */
4080 0, /* tp_as_number */
4081 0, /* tp_as_sequence */
4082 0, /* tp_as_mapping */
4083 0, /* tp_hash */
4084 0, /* tp_call */
4085 0, /* tp_str */
4086 PyObject_GenericGetAttr, /* tp_getattro */
4087 0, /* tp_setattro */
4088 0, /* tp_as_buffer */
4089 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4090 Py_TPFLAGS_BASETYPE, /* tp_flags */
4091 count_doc, /* tp_doc */
4092 (traverseproc)count_traverse, /* tp_traverse */
4093 0, /* tp_clear */
4094 0, /* tp_richcompare */
4095 0, /* tp_weaklistoffset */
4096 PyObject_SelfIter, /* tp_iter */
4097 (iternextfunc)count_next, /* tp_iternext */
4098 count_methods, /* tp_methods */
4099 0, /* tp_members */
4100 0, /* tp_getset */
4101 0, /* tp_base */
4102 0, /* tp_dict */
4103 0, /* tp_descr_get */
4104 0, /* tp_descr_set */
4105 0, /* tp_dictoffset */
4106 0, /* tp_init */
4107 0, /* tp_alloc */
4108 count_new, /* tp_new */
4109 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004110};
4111
4112
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004113/* repeat object ************************************************************/
4114
4115typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004116 PyObject_HEAD
4117 PyObject *element;
4118 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004119} repeatobject;
4120
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004121static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004122
4123static PyObject *
4124repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4125{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004126 repeatobject *ro;
4127 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004128 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004129 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004130
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004131 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4132 &element, &cnt))
4133 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004134
Raymond Hettinger97d35552014-06-24 21:36:58 -07004135 if (kwds != NULL)
4136 n_kwds = PyDict_Size(kwds);
4137 /* Does user supply times argument? */
4138 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004139 cnt = 0;
4140
4141 ro = (repeatobject *)type->tp_alloc(type, 0);
4142 if (ro == NULL)
4143 return NULL;
4144 Py_INCREF(element);
4145 ro->element = element;
4146 ro->cnt = cnt;
4147 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004148}
4149
4150static void
4151repeat_dealloc(repeatobject *ro)
4152{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004153 PyObject_GC_UnTrack(ro);
4154 Py_XDECREF(ro->element);
4155 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004156}
4157
4158static int
4159repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4160{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004161 Py_VISIT(ro->element);
4162 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004163}
4164
4165static PyObject *
4166repeat_next(repeatobject *ro)
4167{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004168 if (ro->cnt == 0)
4169 return NULL;
4170 if (ro->cnt > 0)
4171 ro->cnt--;
4172 Py_INCREF(ro->element);
4173 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004174}
4175
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004176static PyObject *
4177repeat_repr(repeatobject *ro)
4178{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004179 if (ro->cnt == -1)
4180 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4181 else
4182 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4183}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004184
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004185static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004186repeat_len(repeatobject *ro)
4187{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004188 if (ro->cnt == -1) {
4189 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4190 return NULL;
4191 }
4192 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004193}
4194
Armin Rigof5b3e362006-02-11 21:32:43 +00004195PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004196
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004197static PyObject *
4198repeat_reduce(repeatobject *ro)
4199{
4200 /* unpickle this so that a new repeat iterator is constructed with an
4201 * object, then call __setstate__ on it to set cnt
4202 */
4203 if (ro->cnt >= 0)
4204 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4205 else
4206 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4207}
4208
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004209static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004210 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004211 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004212 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004213};
4214
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004215PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004216"repeat(object [,times]) -> create an iterator which returns the object\n\
4217for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004218endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004219
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004220static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004221 PyVarObject_HEAD_INIT(NULL, 0)
4222 "itertools.repeat", /* tp_name */
4223 sizeof(repeatobject), /* tp_basicsize */
4224 0, /* tp_itemsize */
4225 /* methods */
4226 (destructor)repeat_dealloc, /* tp_dealloc */
4227 0, /* tp_print */
4228 0, /* tp_getattr */
4229 0, /* tp_setattr */
4230 0, /* tp_reserved */
4231 (reprfunc)repeat_repr, /* tp_repr */
4232 0, /* tp_as_number */
4233 0, /* tp_as_sequence */
4234 0, /* tp_as_mapping */
4235 0, /* tp_hash */
4236 0, /* tp_call */
4237 0, /* tp_str */
4238 PyObject_GenericGetAttr, /* tp_getattro */
4239 0, /* tp_setattro */
4240 0, /* tp_as_buffer */
4241 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4242 Py_TPFLAGS_BASETYPE, /* tp_flags */
4243 repeat_doc, /* tp_doc */
4244 (traverseproc)repeat_traverse, /* tp_traverse */
4245 0, /* tp_clear */
4246 0, /* tp_richcompare */
4247 0, /* tp_weaklistoffset */
4248 PyObject_SelfIter, /* tp_iter */
4249 (iternextfunc)repeat_next, /* tp_iternext */
4250 repeat_methods, /* tp_methods */
4251 0, /* tp_members */
4252 0, /* tp_getset */
4253 0, /* tp_base */
4254 0, /* tp_dict */
4255 0, /* tp_descr_get */
4256 0, /* tp_descr_set */
4257 0, /* tp_dictoffset */
4258 0, /* tp_init */
4259 0, /* tp_alloc */
4260 repeat_new, /* tp_new */
4261 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004262};
4263
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004264/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004265
4266#include "Python.h"
4267
4268typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004269 PyObject_HEAD
4270 Py_ssize_t tuplesize;
4271 Py_ssize_t numactive;
4272 PyObject *ittuple; /* tuple of iterators */
4273 PyObject *result;
4274 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004275} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004276
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004277static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004278
4279static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004280zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004281{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004282 ziplongestobject *lz;
4283 Py_ssize_t i;
4284 PyObject *ittuple; /* tuple of iterators */
4285 PyObject *result;
4286 PyObject *fillvalue = Py_None;
4287 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004288
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004289 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
4290 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4291 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
4292 PyErr_SetString(PyExc_TypeError,
4293 "zip_longest() got an unexpected keyword argument");
4294 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004295 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004296 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004297
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004298 /* args must be a tuple */
4299 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00004300
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004301 /* obtain iterators */
4302 ittuple = PyTuple_New(tuplesize);
4303 if (ittuple == NULL)
4304 return NULL;
4305 for (i=0; i < tuplesize; ++i) {
4306 PyObject *item = PyTuple_GET_ITEM(args, i);
4307 PyObject *it = PyObject_GetIter(item);
4308 if (it == NULL) {
4309 if (PyErr_ExceptionMatches(PyExc_TypeError))
4310 PyErr_Format(PyExc_TypeError,
4311 "zip_longest argument #%zd must support iteration",
4312 i+1);
4313 Py_DECREF(ittuple);
4314 return NULL;
4315 }
4316 PyTuple_SET_ITEM(ittuple, i, it);
4317 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004318
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004319 /* create a result holder */
4320 result = PyTuple_New(tuplesize);
4321 if (result == NULL) {
4322 Py_DECREF(ittuple);
4323 return NULL;
4324 }
4325 for (i=0 ; i < tuplesize ; i++) {
4326 Py_INCREF(Py_None);
4327 PyTuple_SET_ITEM(result, i, Py_None);
4328 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004329
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004330 /* create ziplongestobject structure */
4331 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4332 if (lz == NULL) {
4333 Py_DECREF(ittuple);
4334 Py_DECREF(result);
4335 return NULL;
4336 }
4337 lz->ittuple = ittuple;
4338 lz->tuplesize = tuplesize;
4339 lz->numactive = tuplesize;
4340 lz->result = result;
4341 Py_INCREF(fillvalue);
4342 lz->fillvalue = fillvalue;
4343 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004344}
4345
4346static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004347zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004348{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004349 PyObject_GC_UnTrack(lz);
4350 Py_XDECREF(lz->ittuple);
4351 Py_XDECREF(lz->result);
4352 Py_XDECREF(lz->fillvalue);
4353 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004354}
4355
4356static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004357zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004358{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004359 Py_VISIT(lz->ittuple);
4360 Py_VISIT(lz->result);
4361 Py_VISIT(lz->fillvalue);
4362 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004363}
4364
4365static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004366zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004367{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004368 Py_ssize_t i;
4369 Py_ssize_t tuplesize = lz->tuplesize;
4370 PyObject *result = lz->result;
4371 PyObject *it;
4372 PyObject *item;
4373 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004374
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004375 if (tuplesize == 0)
4376 return NULL;
4377 if (lz->numactive == 0)
4378 return NULL;
4379 if (Py_REFCNT(result) == 1) {
4380 Py_INCREF(result);
4381 for (i=0 ; i < tuplesize ; i++) {
4382 it = PyTuple_GET_ITEM(lz->ittuple, i);
4383 if (it == NULL) {
4384 Py_INCREF(lz->fillvalue);
4385 item = lz->fillvalue;
4386 } else {
4387 item = PyIter_Next(it);
4388 if (item == NULL) {
4389 lz->numactive -= 1;
4390 if (lz->numactive == 0 || PyErr_Occurred()) {
4391 lz->numactive = 0;
4392 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004393 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004394 } else {
4395 Py_INCREF(lz->fillvalue);
4396 item = lz->fillvalue;
4397 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4398 Py_DECREF(it);
4399 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004400 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004401 }
4402 olditem = PyTuple_GET_ITEM(result, i);
4403 PyTuple_SET_ITEM(result, i, item);
4404 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004405 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004406 } else {
4407 result = PyTuple_New(tuplesize);
4408 if (result == NULL)
4409 return NULL;
4410 for (i=0 ; i < tuplesize ; i++) {
4411 it = PyTuple_GET_ITEM(lz->ittuple, i);
4412 if (it == NULL) {
4413 Py_INCREF(lz->fillvalue);
4414 item = lz->fillvalue;
4415 } else {
4416 item = PyIter_Next(it);
4417 if (item == NULL) {
4418 lz->numactive -= 1;
4419 if (lz->numactive == 0 || PyErr_Occurred()) {
4420 lz->numactive = 0;
4421 Py_DECREF(result);
4422 return NULL;
4423 } else {
4424 Py_INCREF(lz->fillvalue);
4425 item = lz->fillvalue;
4426 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4427 Py_DECREF(it);
4428 }
4429 }
4430 }
4431 PyTuple_SET_ITEM(result, i, item);
4432 }
4433 }
4434 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004435}
4436
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004437static PyObject *
4438zip_longest_reduce(ziplongestobject *lz)
4439{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004440
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004441 /* Create a new tuple with empty sequences where appropriate to pickle.
4442 * Then use setstate to set the fillvalue
4443 */
4444 int i;
4445 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
4446 if (args == NULL)
4447 return NULL;
4448 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4449 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4450 if (elem == NULL) {
4451 elem = PyTuple_New(0);
4452 if (elem == NULL) {
4453 Py_DECREF(args);
4454 return NULL;
4455 }
4456 } else
4457 Py_INCREF(elem);
4458 PyTuple_SET_ITEM(args, i, elem);
4459 }
4460 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4461}
4462
4463static PyObject *
4464zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4465{
4466 Py_CLEAR(lz->fillvalue);
4467 lz->fillvalue = state;
4468 Py_INCREF(lz->fillvalue);
4469 Py_RETURN_NONE;
4470}
4471
4472static PyMethodDef zip_longest_methods[] = {
4473 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4474 reduce_doc},
4475 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4476 setstate_doc},
4477 {NULL, NULL} /* sentinel */
4478};
4479
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004480PyDoc_STRVAR(zip_longest_doc,
4481"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004482\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004483Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004484the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004485method continues until the longest iterable in the argument sequence\n\
4486is exhausted and then it raises StopIteration. When the shorter iterables\n\
4487are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4488defaults to None or can be specified by a keyword argument.\n\
4489");
4490
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004491static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004492 PyVarObject_HEAD_INIT(NULL, 0)
4493 "itertools.zip_longest", /* tp_name */
4494 sizeof(ziplongestobject), /* tp_basicsize */
4495 0, /* tp_itemsize */
4496 /* methods */
4497 (destructor)zip_longest_dealloc, /* tp_dealloc */
4498 0, /* tp_print */
4499 0, /* tp_getattr */
4500 0, /* tp_setattr */
4501 0, /* tp_reserved */
4502 0, /* tp_repr */
4503 0, /* tp_as_number */
4504 0, /* tp_as_sequence */
4505 0, /* tp_as_mapping */
4506 0, /* tp_hash */
4507 0, /* tp_call */
4508 0, /* tp_str */
4509 PyObject_GenericGetAttr, /* tp_getattro */
4510 0, /* tp_setattro */
4511 0, /* tp_as_buffer */
4512 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4513 Py_TPFLAGS_BASETYPE, /* tp_flags */
4514 zip_longest_doc, /* tp_doc */
4515 (traverseproc)zip_longest_traverse, /* tp_traverse */
4516 0, /* tp_clear */
4517 0, /* tp_richcompare */
4518 0, /* tp_weaklistoffset */
4519 PyObject_SelfIter, /* tp_iter */
4520 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004521 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004522 0, /* tp_members */
4523 0, /* tp_getset */
4524 0, /* tp_base */
4525 0, /* tp_dict */
4526 0, /* tp_descr_get */
4527 0, /* tp_descr_set */
4528 0, /* tp_dictoffset */
4529 0, /* tp_init */
4530 0, /* tp_alloc */
4531 zip_longest_new, /* tp_new */
4532 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004533};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004534
4535/* module level code ********************************************************/
4536
4537PyDoc_STRVAR(module_doc,
4538"Functional tools for creating and using iterators.\n\
4539\n\
4540Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004541count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004542cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004543repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004544\n\
4545Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004546accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004547chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger8df58f72013-09-09 01:29:40 -05004548chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004549compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4550dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4551groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004552filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004553islice(seq, [start,] stop [, step]) --> elements from\n\
4554 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004555starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004556tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004557takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004558zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004559\n\
4560Combinatoric generators:\n\
4561product(p, q, ... [repeat=1]) --> cartesian product\n\
4562permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004563combinations(p, r)\n\
4564combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004565");
4566
4567
Raymond Hettingerad983e72003-11-12 14:32:26 +00004568static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004569 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4570 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004571};
4572
Martin v. Löwis1a214512008-06-11 05:26:20 +00004573
4574static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004575 PyModuleDef_HEAD_INIT,
4576 "itertools",
4577 module_doc,
4578 -1,
4579 module_methods,
4580 NULL,
4581 NULL,
4582 NULL,
4583 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004584};
4585
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004586PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004587PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004588{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004589 int i;
4590 PyObject *m;
4591 char *name;
4592 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004593 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004594 &combinations_type,
4595 &cwr_type,
4596 &cycle_type,
4597 &dropwhile_type,
4598 &takewhile_type,
4599 &islice_type,
4600 &starmap_type,
4601 &chain_type,
4602 &compress_type,
4603 &filterfalse_type,
4604 &count_type,
4605 &ziplongest_type,
4606 &permutations_type,
4607 &product_type,
4608 &repeat_type,
4609 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004610 &_grouper_type,
4611 &tee_type,
4612 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004613 NULL
4614 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004615
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004616 Py_TYPE(&teedataobject_type) = &PyType_Type;
4617 m = PyModule_Create(&itertoolsmodule);
4618 if (m == NULL)
4619 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004620
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004621 for (i=0 ; typelist[i] != NULL ; i++) {
4622 if (PyType_Ready(typelist[i]) < 0)
4623 return NULL;
4624 name = strchr(typelist[i]->tp_name, '.');
4625 assert (name != NULL);
4626 Py_INCREF(typelist[i]);
4627 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4628 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004629
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004630 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004631}