blob: f5fa3fb3bb32a64c89f332b6ce55a134ab9eb3d6 [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 Hettinger8df58f72013-09-09 01:29:40 -05007 Copyright (c) 2003-2013 Python Software Foundation.
Raymond Hettinger96ef8112003-02-01 00:10:11 +00008 All rights reserved.
9*/
10
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000011
12/* groupby object ***********************************************************/
13
14typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000015 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000021} groupbyobject;
22
23static PyTypeObject groupby_type;
24static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26static PyObject *
27groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000029 static char *kwargs[] = {"iterable", "key", NULL};
30 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000032
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000033 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
34 &it, &keyfunc))
35 return NULL;
36
37 gbo = (groupbyobject *)type->tp_alloc(type, 0);
38 if (gbo == NULL)
39 return NULL;
40 gbo->tgtkey = NULL;
41 gbo->currkey = NULL;
42 gbo->currvalue = NULL;
43 gbo->keyfunc = keyfunc;
44 Py_INCREF(keyfunc);
45 gbo->it = PyObject_GetIter(it);
46 if (gbo->it == NULL) {
47 Py_DECREF(gbo);
48 return NULL;
49 }
50 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000051}
52
53static void
54groupby_dealloc(groupbyobject *gbo)
55{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000056 PyObject_GC_UnTrack(gbo);
57 Py_XDECREF(gbo->it);
58 Py_XDECREF(gbo->keyfunc);
59 Py_XDECREF(gbo->tgtkey);
60 Py_XDECREF(gbo->currkey);
61 Py_XDECREF(gbo->currvalue);
62 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000063}
64
65static int
66groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
67{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000068 Py_VISIT(gbo->it);
69 Py_VISIT(gbo->keyfunc);
70 Py_VISIT(gbo->tgtkey);
71 Py_VISIT(gbo->currkey);
72 Py_VISIT(gbo->currvalue);
73 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000074}
75
76static PyObject *
77groupby_next(groupbyobject *gbo)
78{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000081 /* skip to next iteration group */
82 for (;;) {
83 if (gbo->currkey == NULL)
84 /* pass */;
85 else if (gbo->tgtkey == NULL)
86 break;
87 else {
88 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000089
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000090 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
91 gbo->currkey, Py_EQ);
92 if (rcmp == -1)
93 return NULL;
94 else if (rcmp == 0)
95 break;
96 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000097
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000098 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000101
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000102 if (gbo->keyfunc == Py_None) {
103 newkey = newvalue;
104 Py_INCREF(newvalue);
105 } else {
106 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
107 newvalue, NULL);
108 if (newkey == NULL) {
109 Py_DECREF(newvalue);
110 return NULL;
111 }
112 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000113
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000114 tmp = gbo->currkey;
115 gbo->currkey = newkey;
116 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000117
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000118 tmp = gbo->currvalue;
119 gbo->currvalue = newvalue;
120 Py_XDECREF(tmp);
121 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000123 Py_INCREF(gbo->currkey);
124 tmp = gbo->tgtkey;
125 gbo->tgtkey = gbo->currkey;
126 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000127
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000131
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000135}
136
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000137static PyObject *
138groupby_reduce(groupbyobject *lz)
139{
140 /* reduce as a 'new' call with an optional 'setstate' if groupby
141 * has started
142 */
143 PyObject *value;
144 if (lz->tgtkey && lz->currkey && lz->currvalue)
145 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
146 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
147 else
148 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
149 lz->it, lz->keyfunc);
150
151 return value;
152}
153
154PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
155
156static PyObject *
157groupby_setstate(groupbyobject *lz, PyObject *state)
158{
159 PyObject *currkey, *currvalue, *tgtkey;
160 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey))
161 return NULL;
162 Py_CLEAR(lz->currkey);
163 lz->currkey = currkey;
164 Py_INCREF(lz->currkey);
165 Py_CLEAR(lz->currvalue);
166 lz->currvalue = currvalue;
167 Py_INCREF(lz->currvalue);
168 Py_CLEAR(lz->tgtkey);
169 lz->tgtkey = tgtkey;
170 Py_INCREF(lz->tgtkey);
171 Py_RETURN_NONE;
172}
173
174PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
175
176static PyMethodDef groupby_methods[] = {
177 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
178 reduce_doc},
179 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
180 setstate_doc},
181 {NULL, NULL} /* sentinel */
182};
183
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000184PyDoc_STRVAR(groupby_doc,
185"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
186(key, sub-iterator) grouped by each value of key(value).\n");
187
188static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000189 PyVarObject_HEAD_INIT(NULL, 0)
190 "itertools.groupby", /* tp_name */
191 sizeof(groupbyobject), /* tp_basicsize */
192 0, /* tp_itemsize */
193 /* methods */
194 (destructor)groupby_dealloc, /* tp_dealloc */
195 0, /* tp_print */
196 0, /* tp_getattr */
197 0, /* tp_setattr */
198 0, /* tp_reserved */
199 0, /* tp_repr */
200 0, /* tp_as_number */
201 0, /* tp_as_sequence */
202 0, /* tp_as_mapping */
203 0, /* tp_hash */
204 0, /* tp_call */
205 0, /* tp_str */
206 PyObject_GenericGetAttr, /* tp_getattro */
207 0, /* tp_setattro */
208 0, /* tp_as_buffer */
209 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
210 Py_TPFLAGS_BASETYPE, /* tp_flags */
211 groupby_doc, /* tp_doc */
212 (traverseproc)groupby_traverse, /* tp_traverse */
213 0, /* tp_clear */
214 0, /* tp_richcompare */
215 0, /* tp_weaklistoffset */
216 PyObject_SelfIter, /* tp_iter */
217 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000218 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000219 0, /* tp_members */
220 0, /* tp_getset */
221 0, /* tp_base */
222 0, /* tp_dict */
223 0, /* tp_descr_get */
224 0, /* tp_descr_set */
225 0, /* tp_dictoffset */
226 0, /* tp_init */
227 0, /* tp_alloc */
228 groupby_new, /* tp_new */
229 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000230};
231
232
233/* _grouper object (internal) ************************************************/
234
235typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000236 PyObject_HEAD
237 PyObject *parent;
238 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000239} _grouperobject;
240
241static PyTypeObject _grouper_type;
242
243static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000244_grouper_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
245{
246 PyObject *parent, *tgtkey;
247
248 if (!PyArg_ParseTuple(args, "O!O", &groupby_type, &parent, &tgtkey))
249 return NULL;
250
251 return _grouper_create((groupbyobject*) parent, tgtkey);
252}
253
254static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000255_grouper_create(groupbyobject *parent, PyObject *tgtkey)
256{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000257 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000259 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
260 if (igo == NULL)
261 return NULL;
262 igo->parent = (PyObject *)parent;
263 Py_INCREF(parent);
264 igo->tgtkey = tgtkey;
265 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000266
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000267 PyObject_GC_Track(igo);
268 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000269}
270
271static void
272_grouper_dealloc(_grouperobject *igo)
273{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000274 PyObject_GC_UnTrack(igo);
275 Py_DECREF(igo->parent);
276 Py_DECREF(igo->tgtkey);
277 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000278}
279
280static int
281_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
282{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000283 Py_VISIT(igo->parent);
284 Py_VISIT(igo->tgtkey);
285 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000286}
287
288static PyObject *
289_grouper_next(_grouperobject *igo)
290{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000291 groupbyobject *gbo = (groupbyobject *)igo->parent;
292 PyObject *newvalue, *newkey, *r;
293 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000294
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000295 if (gbo->currvalue == NULL) {
296 newvalue = PyIter_Next(gbo->it);
297 if (newvalue == NULL)
298 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000299
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000300 if (gbo->keyfunc == Py_None) {
301 newkey = newvalue;
302 Py_INCREF(newvalue);
303 } else {
304 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
305 newvalue, NULL);
306 if (newkey == NULL) {
307 Py_DECREF(newvalue);
308 return NULL;
309 }
310 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000311
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000312 assert(gbo->currkey == NULL);
313 gbo->currkey = newkey;
314 gbo->currvalue = newvalue;
315 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000316
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000317 assert(gbo->currkey != NULL);
318 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
319 if (rcmp <= 0)
320 /* got any error or current group is end */
321 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000322
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000323 r = gbo->currvalue;
324 gbo->currvalue = NULL;
325 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000326
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000327 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000328}
329
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000330static PyObject *
331_grouper_reduce(_grouperobject *lz)
332{
333 return Py_BuildValue("O(OO)", Py_TYPE(lz),
334 lz->parent, lz->tgtkey);
335}
336
337static PyMethodDef _grouper_methods[] = {
338 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
339 reduce_doc},
340 {NULL, NULL} /* sentinel */
341};
342
343
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000344static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000345 PyVarObject_HEAD_INIT(NULL, 0)
346 "itertools._grouper", /* tp_name */
347 sizeof(_grouperobject), /* tp_basicsize */
348 0, /* tp_itemsize */
349 /* methods */
350 (destructor)_grouper_dealloc, /* tp_dealloc */
351 0, /* tp_print */
352 0, /* tp_getattr */
353 0, /* tp_setattr */
354 0, /* tp_reserved */
355 0, /* tp_repr */
356 0, /* tp_as_number */
357 0, /* tp_as_sequence */
358 0, /* tp_as_mapping */
359 0, /* tp_hash */
360 0, /* tp_call */
361 0, /* tp_str */
362 PyObject_GenericGetAttr, /* tp_getattro */
363 0, /* tp_setattro */
364 0, /* tp_as_buffer */
365 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
366 0, /* tp_doc */
367 (traverseproc)_grouper_traverse,/* tp_traverse */
368 0, /* tp_clear */
369 0, /* tp_richcompare */
370 0, /* tp_weaklistoffset */
371 PyObject_SelfIter, /* tp_iter */
372 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000373 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000374 0, /* tp_members */
375 0, /* tp_getset */
376 0, /* tp_base */
377 0, /* tp_dict */
378 0, /* tp_descr_get */
379 0, /* tp_descr_set */
380 0, /* tp_dictoffset */
381 0, /* tp_init */
382 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000383 _grouper_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000384 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000385};
386
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000387
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000388
Raymond Hettingerad983e72003-11-12 14:32:26 +0000389/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000390
Raymond Hettingerad983e72003-11-12 14:32:26 +0000391/* The teedataobject pre-allocates space for LINKCELLS number of objects.
392 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000393 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000394 the other structure members including PyHEAD overhead. The larger the
395 value, the less memory overhead per object and the less time spent
396 allocating/deallocating new links. The smaller the number, the less
397 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000398*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000399#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000400
401typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000402 PyObject_HEAD
403 PyObject *it;
Antoine Pitroub4a46cb2013-09-20 22:19:22 +0200404 int numread; /* 0 <= numread <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000405 PyObject *nextlink;
406 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000407} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000408
409typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000410 PyObject_HEAD
411 teedataobject *dataobj;
Antoine Pitroub4a46cb2013-09-20 22:19:22 +0200412 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000413 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000414} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415
Raymond Hettingerad983e72003-11-12 14:32:26 +0000416static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417
418static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000419teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000420{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000421 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000422
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000423 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
424 if (tdo == NULL)
425 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000426
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000427 tdo->numread = 0;
428 tdo->nextlink = NULL;
429 Py_INCREF(it);
430 tdo->it = it;
431 PyObject_GC_Track(tdo);
432 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000433}
434
435static PyObject *
436teedataobject_jumplink(teedataobject *tdo)
437{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000438 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000439 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000440 Py_XINCREF(tdo->nextlink);
441 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000442}
443
444static PyObject *
445teedataobject_getitem(teedataobject *tdo, int i)
446{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000447 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000448
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000449 assert(i < LINKCELLS);
450 if (i < tdo->numread)
451 value = tdo->values[i];
452 else {
453 /* this is the lead iterator, so fetch more data */
454 assert(i == tdo->numread);
455 value = PyIter_Next(tdo->it);
456 if (value == NULL)
457 return NULL;
458 tdo->numread++;
459 tdo->values[i] = value;
460 }
461 Py_INCREF(value);
462 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000463}
464
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000465static int
466teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
467{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000468 int i;
469 Py_VISIT(tdo->it);
470 for (i = 0; i < tdo->numread; i++)
471 Py_VISIT(tdo->values[i]);
472 Py_VISIT(tdo->nextlink);
473 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000474}
475
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200476static void
477teedataobject_safe_decref(PyObject *obj)
478{
479 while (obj && Py_TYPE(obj) == &teedataobject_type &&
480 Py_REFCNT(obj) == 1) {
481 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
482 ((teedataobject *)obj)->nextlink = NULL;
483 Py_DECREF(obj);
484 obj = nextlink;
485 }
486 Py_XDECREF(obj);
487}
488
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000489static int
490teedataobject_clear(teedataobject *tdo)
491{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000492 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200493 PyObject *tmp;
494
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000495 Py_CLEAR(tdo->it);
496 for (i=0 ; i<tdo->numread ; i++)
497 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200498 tmp = tdo->nextlink;
499 tdo->nextlink = NULL;
500 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000501 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000502}
503
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000504static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000505teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000506{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000507 PyObject_GC_UnTrack(tdo);
508 teedataobject_clear(tdo);
509 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000510}
511
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000512static PyObject *
513teedataobject_reduce(teedataobject *tdo)
514{
515 int i;
516 /* create a temporary list of already iterated values */
517 PyObject *values = PyList_New(tdo->numread);
518 if (!values)
519 return NULL;
520 for (i=0 ; i<tdo->numread ; i++) {
521 Py_INCREF(tdo->values[i]);
522 PyList_SET_ITEM(values, i, tdo->values[i]);
523 }
524 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
525 values,
526 tdo->nextlink ? tdo->nextlink : Py_None);
527}
528
529static PyTypeObject teedataobject_type;
530
531static PyObject *
532teedataobject_new(PyTypeObject *type, PyObject *args, PyObject *kw)
533{
534 teedataobject *tdo;
535 PyObject *it, *values, *next;
536 Py_ssize_t i, len;
537
538 assert(type == &teedataobject_type);
539 if (!PyArg_ParseTuple(args, "OO!O", &it, &PyList_Type, &values, &next))
540 return NULL;
541
542 tdo = (teedataobject *)teedataobject_newinternal(it);
543 if (!tdo)
544 return NULL;
545
546 len = PyList_GET_SIZE(values);
547 if (len > LINKCELLS)
548 goto err;
549 for (i=0; i<len; i++) {
550 tdo->values[i] = PyList_GET_ITEM(values, i);
551 Py_INCREF(tdo->values[i]);
552 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200553 /* len <= LINKCELLS < INT_MAX */
554 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000555
556 if (len == LINKCELLS) {
557 if (next != Py_None) {
558 if (Py_TYPE(next) != &teedataobject_type)
559 goto err;
560 assert(tdo->nextlink == NULL);
561 Py_INCREF(next);
562 tdo->nextlink = next;
563 }
564 } else {
565 if (next != Py_None)
566 goto err; /* shouldn't have a next if we are not full */
567 }
568 return (PyObject*)tdo;
569
570err:
571 Py_XDECREF(tdo);
572 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
573 return NULL;
574}
575
576static PyMethodDef teedataobject_methods[] = {
577 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
578 reduce_doc},
579 {NULL, NULL} /* sentinel */
580};
581
Raymond Hettingerad983e72003-11-12 14:32:26 +0000582PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000583
Raymond Hettingerad983e72003-11-12 14:32:26 +0000584static PyTypeObject teedataobject_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000585 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000586 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000587 sizeof(teedataobject), /* tp_basicsize */
588 0, /* tp_itemsize */
589 /* methods */
590 (destructor)teedataobject_dealloc, /* tp_dealloc */
591 0, /* tp_print */
592 0, /* tp_getattr */
593 0, /* tp_setattr */
594 0, /* tp_reserved */
595 0, /* tp_repr */
596 0, /* tp_as_number */
597 0, /* tp_as_sequence */
598 0, /* tp_as_mapping */
599 0, /* tp_hash */
600 0, /* tp_call */
601 0, /* tp_str */
602 PyObject_GenericGetAttr, /* tp_getattro */
603 0, /* tp_setattro */
604 0, /* tp_as_buffer */
605 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
606 teedataobject_doc, /* tp_doc */
607 (traverseproc)teedataobject_traverse, /* tp_traverse */
608 (inquiry)teedataobject_clear, /* tp_clear */
609 0, /* tp_richcompare */
610 0, /* tp_weaklistoffset */
611 0, /* tp_iter */
612 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000613 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000614 0, /* tp_members */
615 0, /* tp_getset */
616 0, /* tp_base */
617 0, /* tp_dict */
618 0, /* tp_descr_get */
619 0, /* tp_descr_set */
620 0, /* tp_dictoffset */
621 0, /* tp_init */
622 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000623 teedataobject_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000624 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000625};
626
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000627
628static PyTypeObject tee_type;
629
630static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000631tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000632{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000633 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000634
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000635 if (to->index >= LINKCELLS) {
636 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200637 if (link == NULL)
638 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000639 Py_DECREF(to->dataobj);
640 to->dataobj = (teedataobject *)link;
641 to->index = 0;
642 }
643 value = teedataobject_getitem(to->dataobj, to->index);
644 if (value == NULL)
645 return NULL;
646 to->index++;
647 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000648}
649
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000650static int
651tee_traverse(teeobject *to, visitproc visit, void *arg)
652{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000653 Py_VISIT((PyObject *)to->dataobj);
654 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000655}
656
Raymond Hettingerad983e72003-11-12 14:32:26 +0000657static PyObject *
658tee_copy(teeobject *to)
659{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000660 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000661
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000662 newto = PyObject_GC_New(teeobject, &tee_type);
663 if (newto == NULL)
664 return NULL;
665 Py_INCREF(to->dataobj);
666 newto->dataobj = to->dataobj;
667 newto->index = to->index;
668 newto->weakreflist = NULL;
669 PyObject_GC_Track(newto);
670 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000671}
672
673PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
674
675static PyObject *
676tee_fromiterable(PyObject *iterable)
677{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000678 teeobject *to;
679 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000680
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000681 it = PyObject_GetIter(iterable);
682 if (it == NULL)
683 return NULL;
684 if (PyObject_TypeCheck(it, &tee_type)) {
685 to = (teeobject *)tee_copy((teeobject *)it);
686 goto done;
687 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000688
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000689 to = PyObject_GC_New(teeobject, &tee_type);
690 if (to == NULL)
691 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000692 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000693 if (!to->dataobj) {
694 PyObject_GC_Del(to);
695 to = NULL;
696 goto done;
697 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000698
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000699 to->index = 0;
700 to->weakreflist = NULL;
701 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000702done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000703 Py_XDECREF(it);
704 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000705}
706
707static PyObject *
708tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
709{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000710 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000711
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000712 if (!PyArg_UnpackTuple(args, "_tee", 1, 1, &iterable))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000713 return NULL;
714 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000715}
716
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000717static int
718tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000719{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000720 if (to->weakreflist != NULL)
721 PyObject_ClearWeakRefs((PyObject *) to);
722 Py_CLEAR(to->dataobj);
723 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000724}
725
726static void
727tee_dealloc(teeobject *to)
728{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000729 PyObject_GC_UnTrack(to);
730 tee_clear(to);
731 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000732}
733
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000734static PyObject *
735tee_reduce(teeobject *to)
736{
737 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
738}
739
740static PyObject *
741tee_setstate(teeobject *to, PyObject *state)
742{
743 teedataobject *tdo;
744 int index;
745 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index))
746 return NULL;
747 if (index < 0 || index > LINKCELLS) {
748 PyErr_SetString(PyExc_ValueError, "Index out of range");
749 return NULL;
750 }
751 Py_CLEAR(to->dataobj);
752 to->dataobj = tdo;
753 Py_INCREF(to->dataobj);
754 to->index = index;
755 Py_RETURN_NONE;
756}
757
Raymond Hettingerad983e72003-11-12 14:32:26 +0000758PyDoc_STRVAR(teeobject_doc,
759"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000760
Raymond Hettingerad983e72003-11-12 14:32:26 +0000761static PyMethodDef tee_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000762 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000763 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
764 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000765 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000766};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000767
768static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000769 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000770 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000771 sizeof(teeobject), /* tp_basicsize */
772 0, /* tp_itemsize */
773 /* methods */
774 (destructor)tee_dealloc, /* tp_dealloc */
775 0, /* tp_print */
776 0, /* tp_getattr */
777 0, /* tp_setattr */
778 0, /* tp_reserved */
779 0, /* tp_repr */
780 0, /* tp_as_number */
781 0, /* tp_as_sequence */
782 0, /* tp_as_mapping */
783 0, /* tp_hash */
784 0, /* tp_call */
785 0, /* tp_str */
786 0, /* tp_getattro */
787 0, /* tp_setattro */
788 0, /* tp_as_buffer */
789 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
790 teeobject_doc, /* tp_doc */
791 (traverseproc)tee_traverse, /* tp_traverse */
792 (inquiry)tee_clear, /* tp_clear */
793 0, /* tp_richcompare */
794 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
795 PyObject_SelfIter, /* tp_iter */
796 (iternextfunc)tee_next, /* tp_iternext */
797 tee_methods, /* tp_methods */
798 0, /* tp_members */
799 0, /* tp_getset */
800 0, /* tp_base */
801 0, /* tp_dict */
802 0, /* tp_descr_get */
803 0, /* tp_descr_set */
804 0, /* tp_dictoffset */
805 0, /* tp_init */
806 0, /* tp_alloc */
807 tee_new, /* tp_new */
808 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000809};
810
Raymond Hettingerad983e72003-11-12 14:32:26 +0000811static PyObject *
812tee(PyObject *self, PyObject *args)
813{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000814 Py_ssize_t i, n=2;
815 PyObject *it, *iterable, *copyable, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200816 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000817
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000818 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
819 return NULL;
820 if (n < 0) {
821 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
822 return NULL;
823 }
824 result = PyTuple_New(n);
825 if (result == NULL)
826 return NULL;
827 if (n == 0)
828 return result;
829 it = PyObject_GetIter(iterable);
830 if (it == NULL) {
831 Py_DECREF(result);
832 return NULL;
833 }
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200834 if (!_PyObject_HasAttrId(it, &PyId___copy__)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000835 copyable = tee_fromiterable(it);
836 Py_DECREF(it);
837 if (copyable == NULL) {
838 Py_DECREF(result);
839 return NULL;
840 }
841 } else
842 copyable = it;
843 PyTuple_SET_ITEM(result, 0, copyable);
844 for (i=1 ; i<n ; i++) {
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200845
846 copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000847 if (copyable == NULL) {
848 Py_DECREF(result);
849 return NULL;
850 }
851 PyTuple_SET_ITEM(result, i, copyable);
852 }
853 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000854}
855
856PyDoc_STRVAR(tee_doc,
857"tee(iterable, n=2) --> tuple of n independent iterators.");
858
859
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000860/* cycle object **********************************************************/
861
862typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000863 PyObject_HEAD
864 PyObject *it;
865 PyObject *saved;
866 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000867} cycleobject;
868
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000869static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000870
871static PyObject *
872cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
873{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000874 PyObject *it;
875 PyObject *iterable;
876 PyObject *saved;
877 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000878
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000879 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
880 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000881
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000882 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
883 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000884
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000885 /* Get iterator. */
886 it = PyObject_GetIter(iterable);
887 if (it == NULL)
888 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000889
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000890 saved = PyList_New(0);
891 if (saved == NULL) {
892 Py_DECREF(it);
893 return NULL;
894 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000895
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000896 /* create cycleobject structure */
897 lz = (cycleobject *)type->tp_alloc(type, 0);
898 if (lz == NULL) {
899 Py_DECREF(it);
900 Py_DECREF(saved);
901 return NULL;
902 }
903 lz->it = it;
904 lz->saved = saved;
905 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000906
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000907 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000908}
909
910static void
911cycle_dealloc(cycleobject *lz)
912{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000913 PyObject_GC_UnTrack(lz);
914 Py_XDECREF(lz->saved);
915 Py_XDECREF(lz->it);
916 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000917}
918
919static int
920cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
921{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000922 Py_VISIT(lz->it);
923 Py_VISIT(lz->saved);
924 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000925}
926
927static PyObject *
928cycle_next(cycleobject *lz)
929{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000930 PyObject *item;
931 PyObject *it;
932 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000933
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000934 while (1) {
935 item = PyIter_Next(lz->it);
936 if (item != NULL) {
937 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
938 Py_DECREF(item);
939 return NULL;
940 }
941 return item;
942 }
943 if (PyErr_Occurred()) {
944 if (PyErr_ExceptionMatches(PyExc_StopIteration))
945 PyErr_Clear();
946 else
947 return NULL;
948 }
949 if (PyList_Size(lz->saved) == 0)
950 return NULL;
951 it = PyObject_GetIter(lz->saved);
952 if (it == NULL)
953 return NULL;
954 tmp = lz->it;
955 lz->it = it;
956 lz->firstpass = 1;
957 Py_DECREF(tmp);
958 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000959}
960
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000961static PyObject *
962cycle_reduce(cycleobject *lz)
963{
964 /* Create a new cycle with the iterator tuple, then set
965 * the saved state on it.
966 */
Serhiy Storchakad269b5e2013-01-25 13:38:56 +0200967 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz),
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000968 lz->it, lz->saved, lz->firstpass);
969 }
970
971static PyObject *
972cycle_setstate(cycleobject *lz, PyObject *state)
973{
974 PyObject *saved=NULL;
975 int firstpass;
976 if (!PyArg_ParseTuple(state, "Oi", &saved, &firstpass))
977 return NULL;
978 Py_CLEAR(lz->saved);
979 lz->saved = saved;
980 Py_XINCREF(lz->saved);
981 lz->firstpass = firstpass != 0;
982 Py_RETURN_NONE;
983}
984
985static PyMethodDef cycle_methods[] = {
986 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
987 reduce_doc},
988 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
989 setstate_doc},
990 {NULL, NULL} /* sentinel */
991};
992
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000993PyDoc_STRVAR(cycle_doc,
994"cycle(iterable) --> cycle object\n\
995\n\
996Return elements from the iterable until it is exhausted.\n\
997Then repeat the sequence indefinitely.");
998
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000999static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001000 PyVarObject_HEAD_INIT(NULL, 0)
1001 "itertools.cycle", /* tp_name */
1002 sizeof(cycleobject), /* tp_basicsize */
1003 0, /* tp_itemsize */
1004 /* methods */
1005 (destructor)cycle_dealloc, /* tp_dealloc */
1006 0, /* tp_print */
1007 0, /* tp_getattr */
1008 0, /* tp_setattr */
1009 0, /* tp_reserved */
1010 0, /* tp_repr */
1011 0, /* tp_as_number */
1012 0, /* tp_as_sequence */
1013 0, /* tp_as_mapping */
1014 0, /* tp_hash */
1015 0, /* tp_call */
1016 0, /* tp_str */
1017 PyObject_GenericGetAttr, /* tp_getattro */
1018 0, /* tp_setattro */
1019 0, /* tp_as_buffer */
1020 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1021 Py_TPFLAGS_BASETYPE, /* tp_flags */
1022 cycle_doc, /* tp_doc */
1023 (traverseproc)cycle_traverse, /* tp_traverse */
1024 0, /* tp_clear */
1025 0, /* tp_richcompare */
1026 0, /* tp_weaklistoffset */
1027 PyObject_SelfIter, /* tp_iter */
1028 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001029 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001030 0, /* tp_members */
1031 0, /* tp_getset */
1032 0, /* tp_base */
1033 0, /* tp_dict */
1034 0, /* tp_descr_get */
1035 0, /* tp_descr_set */
1036 0, /* tp_dictoffset */
1037 0, /* tp_init */
1038 0, /* tp_alloc */
1039 cycle_new, /* tp_new */
1040 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001041};
1042
1043
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001044/* dropwhile object **********************************************************/
1045
1046typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001047 PyObject_HEAD
1048 PyObject *func;
1049 PyObject *it;
1050 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001051} dropwhileobject;
1052
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001053static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001054
1055static PyObject *
1056dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1057{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001058 PyObject *func, *seq;
1059 PyObject *it;
1060 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001061
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001062 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
1063 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001064
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001065 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
1066 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001067
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001068 /* Get iterator. */
1069 it = PyObject_GetIter(seq);
1070 if (it == NULL)
1071 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001072
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001073 /* create dropwhileobject structure */
1074 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1075 if (lz == NULL) {
1076 Py_DECREF(it);
1077 return NULL;
1078 }
1079 Py_INCREF(func);
1080 lz->func = func;
1081 lz->it = it;
1082 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001083
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001084 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001085}
1086
1087static void
1088dropwhile_dealloc(dropwhileobject *lz)
1089{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001090 PyObject_GC_UnTrack(lz);
1091 Py_XDECREF(lz->func);
1092 Py_XDECREF(lz->it);
1093 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001094}
1095
1096static int
1097dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1098{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001099 Py_VISIT(lz->it);
1100 Py_VISIT(lz->func);
1101 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001102}
1103
1104static PyObject *
1105dropwhile_next(dropwhileobject *lz)
1106{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001107 PyObject *item, *good;
1108 PyObject *it = lz->it;
1109 long ok;
1110 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001111
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001112 iternext = *Py_TYPE(it)->tp_iternext;
1113 for (;;) {
1114 item = iternext(it);
1115 if (item == NULL)
1116 return NULL;
1117 if (lz->start == 1)
1118 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001119
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001120 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1121 if (good == NULL) {
1122 Py_DECREF(item);
1123 return NULL;
1124 }
1125 ok = PyObject_IsTrue(good);
1126 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001127 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001128 lz->start = 1;
1129 return item;
1130 }
1131 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001132 if (ok < 0)
1133 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001134 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001135}
1136
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001137static PyObject *
1138dropwhile_reduce(dropwhileobject *lz)
1139{
1140 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1141 lz->func, lz->it, lz->start);
1142}
1143
1144static PyObject *
1145dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1146{
1147 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001148 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001149 return NULL;
1150 lz->start = start;
1151 Py_RETURN_NONE;
1152}
1153
1154static PyMethodDef dropwhile_methods[] = {
1155 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1156 reduce_doc},
1157 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1158 setstate_doc},
1159 {NULL, NULL} /* sentinel */
1160};
1161
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001162PyDoc_STRVAR(dropwhile_doc,
1163"dropwhile(predicate, iterable) --> dropwhile object\n\
1164\n\
1165Drop items from the iterable while predicate(item) is true.\n\
1166Afterwards, return every element until the iterable is exhausted.");
1167
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001168static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001169 PyVarObject_HEAD_INIT(NULL, 0)
1170 "itertools.dropwhile", /* tp_name */
1171 sizeof(dropwhileobject), /* tp_basicsize */
1172 0, /* tp_itemsize */
1173 /* methods */
1174 (destructor)dropwhile_dealloc, /* tp_dealloc */
1175 0, /* tp_print */
1176 0, /* tp_getattr */
1177 0, /* tp_setattr */
1178 0, /* tp_reserved */
1179 0, /* tp_repr */
1180 0, /* tp_as_number */
1181 0, /* tp_as_sequence */
1182 0, /* tp_as_mapping */
1183 0, /* tp_hash */
1184 0, /* tp_call */
1185 0, /* tp_str */
1186 PyObject_GenericGetAttr, /* tp_getattro */
1187 0, /* tp_setattro */
1188 0, /* tp_as_buffer */
1189 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1190 Py_TPFLAGS_BASETYPE, /* tp_flags */
1191 dropwhile_doc, /* tp_doc */
1192 (traverseproc)dropwhile_traverse, /* tp_traverse */
1193 0, /* tp_clear */
1194 0, /* tp_richcompare */
1195 0, /* tp_weaklistoffset */
1196 PyObject_SelfIter, /* tp_iter */
1197 (iternextfunc)dropwhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001198 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001199 0, /* tp_members */
1200 0, /* tp_getset */
1201 0, /* tp_base */
1202 0, /* tp_dict */
1203 0, /* tp_descr_get */
1204 0, /* tp_descr_set */
1205 0, /* tp_dictoffset */
1206 0, /* tp_init */
1207 0, /* tp_alloc */
1208 dropwhile_new, /* tp_new */
1209 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001210};
1211
1212
1213/* takewhile object **********************************************************/
1214
1215typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001216 PyObject_HEAD
1217 PyObject *func;
1218 PyObject *it;
1219 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001220} takewhileobject;
1221
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001222static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001223
1224static PyObject *
1225takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1226{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001227 PyObject *func, *seq;
1228 PyObject *it;
1229 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001230
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001231 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1232 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001233
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001234 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1235 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001236
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001237 /* Get iterator. */
1238 it = PyObject_GetIter(seq);
1239 if (it == NULL)
1240 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001241
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001242 /* create takewhileobject structure */
1243 lz = (takewhileobject *)type->tp_alloc(type, 0);
1244 if (lz == NULL) {
1245 Py_DECREF(it);
1246 return NULL;
1247 }
1248 Py_INCREF(func);
1249 lz->func = func;
1250 lz->it = it;
1251 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001252
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001253 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001254}
1255
1256static void
1257takewhile_dealloc(takewhileobject *lz)
1258{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001259 PyObject_GC_UnTrack(lz);
1260 Py_XDECREF(lz->func);
1261 Py_XDECREF(lz->it);
1262 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001263}
1264
1265static int
1266takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1267{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001268 Py_VISIT(lz->it);
1269 Py_VISIT(lz->func);
1270 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001271}
1272
1273static PyObject *
1274takewhile_next(takewhileobject *lz)
1275{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001276 PyObject *item, *good;
1277 PyObject *it = lz->it;
1278 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001279
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001280 if (lz->stop == 1)
1281 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001282
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001283 item = (*Py_TYPE(it)->tp_iternext)(it);
1284 if (item == NULL)
1285 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001286
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001287 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1288 if (good == NULL) {
1289 Py_DECREF(item);
1290 return NULL;
1291 }
1292 ok = PyObject_IsTrue(good);
1293 Py_DECREF(good);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001294 if (ok == 1)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001295 return item;
1296 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001297 if (ok == 0)
1298 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001299 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001300}
1301
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001302static PyObject *
1303takewhile_reduce(takewhileobject *lz)
1304{
1305 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1306 lz->func, lz->it, lz->stop);
1307}
1308
1309static PyObject *
1310takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1311{
1312 int stop = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001313 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001314 return NULL;
1315 lz->stop = stop;
1316 Py_RETURN_NONE;
1317}
1318
1319static PyMethodDef takewhile_reduce_methods[] = {
1320 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1321 reduce_doc},
1322 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1323 setstate_doc},
1324 {NULL, NULL} /* sentinel */
1325};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001326PyDoc_STRVAR(takewhile_doc,
1327"takewhile(predicate, iterable) --> takewhile object\n\
1328\n\
1329Return successive entries from an iterable as long as the \n\
1330predicate evaluates to true for each entry.");
1331
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001332static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001333 PyVarObject_HEAD_INIT(NULL, 0)
1334 "itertools.takewhile", /* tp_name */
1335 sizeof(takewhileobject), /* tp_basicsize */
1336 0, /* tp_itemsize */
1337 /* methods */
1338 (destructor)takewhile_dealloc, /* tp_dealloc */
1339 0, /* tp_print */
1340 0, /* tp_getattr */
1341 0, /* tp_setattr */
1342 0, /* tp_reserved */
1343 0, /* tp_repr */
1344 0, /* tp_as_number */
1345 0, /* tp_as_sequence */
1346 0, /* tp_as_mapping */
1347 0, /* tp_hash */
1348 0, /* tp_call */
1349 0, /* tp_str */
1350 PyObject_GenericGetAttr, /* tp_getattro */
1351 0, /* tp_setattro */
1352 0, /* tp_as_buffer */
1353 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1354 Py_TPFLAGS_BASETYPE, /* tp_flags */
1355 takewhile_doc, /* tp_doc */
1356 (traverseproc)takewhile_traverse, /* tp_traverse */
1357 0, /* tp_clear */
1358 0, /* tp_richcompare */
1359 0, /* tp_weaklistoffset */
1360 PyObject_SelfIter, /* tp_iter */
1361 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001362 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001363 0, /* tp_members */
1364 0, /* tp_getset */
1365 0, /* tp_base */
1366 0, /* tp_dict */
1367 0, /* tp_descr_get */
1368 0, /* tp_descr_set */
1369 0, /* tp_dictoffset */
1370 0, /* tp_init */
1371 0, /* tp_alloc */
1372 takewhile_new, /* tp_new */
1373 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001374};
1375
1376
1377/* islice object ************************************************************/
1378
1379typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001380 PyObject_HEAD
1381 PyObject *it;
1382 Py_ssize_t next;
1383 Py_ssize_t stop;
1384 Py_ssize_t step;
1385 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001386} isliceobject;
1387
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001388static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001389
1390static PyObject *
1391islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1392{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001393 PyObject *seq;
1394 Py_ssize_t start=0, stop=-1, step=1;
1395 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1396 Py_ssize_t numargs;
1397 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001398
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001399 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1400 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001401
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001402 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1403 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001404
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001405 numargs = PyTuple_Size(args);
1406 if (numargs == 2) {
1407 if (a1 != Py_None) {
1408 stop = PyLong_AsSsize_t(a1);
1409 if (stop == -1) {
1410 if (PyErr_Occurred())
1411 PyErr_Clear();
1412 PyErr_SetString(PyExc_ValueError,
1413 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1414 return NULL;
1415 }
1416 }
1417 } else {
1418 if (a1 != Py_None)
1419 start = PyLong_AsSsize_t(a1);
1420 if (start == -1 && PyErr_Occurred())
1421 PyErr_Clear();
1422 if (a2 != Py_None) {
1423 stop = PyLong_AsSsize_t(a2);
1424 if (stop == -1) {
1425 if (PyErr_Occurred())
1426 PyErr_Clear();
1427 PyErr_SetString(PyExc_ValueError,
1428 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1429 return NULL;
1430 }
1431 }
1432 }
1433 if (start<0 || stop<-1) {
1434 PyErr_SetString(PyExc_ValueError,
1435 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1436 return NULL;
1437 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001438
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001439 if (a3 != NULL) {
1440 if (a3 != Py_None)
1441 step = PyLong_AsSsize_t(a3);
1442 if (step == -1 && PyErr_Occurred())
1443 PyErr_Clear();
1444 }
1445 if (step<1) {
1446 PyErr_SetString(PyExc_ValueError,
1447 "Step for islice() must be a positive integer or None.");
1448 return NULL;
1449 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001450
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001451 /* Get iterator. */
1452 it = PyObject_GetIter(seq);
1453 if (it == NULL)
1454 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001455
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001456 /* create isliceobject structure */
1457 lz = (isliceobject *)type->tp_alloc(type, 0);
1458 if (lz == NULL) {
1459 Py_DECREF(it);
1460 return NULL;
1461 }
1462 lz->it = it;
1463 lz->next = start;
1464 lz->stop = stop;
1465 lz->step = step;
1466 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001467
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001468 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469}
1470
1471static void
1472islice_dealloc(isliceobject *lz)
1473{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001474 PyObject_GC_UnTrack(lz);
1475 Py_XDECREF(lz->it);
1476 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001477}
1478
1479static int
1480islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1481{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001482 Py_VISIT(lz->it);
1483 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001484}
1485
1486static PyObject *
1487islice_next(isliceobject *lz)
1488{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001489 PyObject *item;
1490 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001491 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001492 Py_ssize_t oldnext;
1493 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001494
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001495 if (it == NULL)
1496 return NULL;
1497
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001498 iternext = *Py_TYPE(it)->tp_iternext;
1499 while (lz->cnt < lz->next) {
1500 item = iternext(it);
1501 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001502 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001503 Py_DECREF(item);
1504 lz->cnt++;
1505 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001506 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001507 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001508 item = iternext(it);
1509 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001510 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001511 lz->cnt++;
1512 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001513 /* The (size_t) cast below avoids the danger of undefined
1514 behaviour from signed integer overflow. */
1515 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001516 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1517 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001518 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001519
1520empty:
1521 Py_CLEAR(lz->it);
1522 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001523}
1524
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001525static PyObject *
1526islice_reduce(isliceobject *lz)
1527{
1528 /* When unpickled, generate a new object with the same bounds,
1529 * then 'setstate' with the next and count
1530 */
1531 PyObject *stop;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001532 if (lz->it == NULL) {
1533 PyObject *empty_list;
1534 PyObject *empty_it;
1535 empty_list = PyList_New(0);
1536 if (empty_list == NULL)
1537 return NULL;
1538 empty_it = PyObject_GetIter(empty_list);
1539 Py_DECREF(empty_list);
1540 if (empty_it == NULL)
1541 return NULL;
1542 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1543 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001544 if (lz->stop == -1) {
1545 stop = Py_None;
1546 Py_INCREF(stop);
1547 } else {
1548 stop = PyLong_FromSsize_t(lz->stop);
1549 if (stop == NULL)
1550 return NULL;
1551 }
1552 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1553 lz->it, lz->next, stop, lz->step,
1554 lz->cnt);
1555}
1556
1557static PyObject *
1558islice_setstate(isliceobject *lz, PyObject *state)
1559{
1560 Py_ssize_t cnt = PyLong_AsSsize_t(state);
1561 if (cnt == -1 && PyErr_Occurred())
1562 return NULL;
1563 lz->cnt = cnt;
1564 Py_RETURN_NONE;
1565}
1566
1567static PyMethodDef islice_methods[] = {
1568 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1569 reduce_doc},
1570 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1571 setstate_doc},
1572 {NULL, NULL} /* sentinel */
1573};
1574
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001575PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001576"islice(iterable, stop) --> islice object\n\
1577islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001578\n\
1579Return an iterator whose next() method returns selected values from an\n\
1580iterable. If start is specified, will skip all preceding elements;\n\
1581otherwise, start defaults to zero. Step defaults to one. If\n\
1582specified as another value, step determines how many values are \n\
1583skipped between successive calls. Works like a slice() on a list\n\
1584but returns an iterator.");
1585
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001586static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001587 PyVarObject_HEAD_INIT(NULL, 0)
1588 "itertools.islice", /* tp_name */
1589 sizeof(isliceobject), /* tp_basicsize */
1590 0, /* tp_itemsize */
1591 /* methods */
1592 (destructor)islice_dealloc, /* tp_dealloc */
1593 0, /* tp_print */
1594 0, /* tp_getattr */
1595 0, /* tp_setattr */
1596 0, /* tp_reserved */
1597 0, /* tp_repr */
1598 0, /* tp_as_number */
1599 0, /* tp_as_sequence */
1600 0, /* tp_as_mapping */
1601 0, /* tp_hash */
1602 0, /* tp_call */
1603 0, /* tp_str */
1604 PyObject_GenericGetAttr, /* tp_getattro */
1605 0, /* tp_setattro */
1606 0, /* tp_as_buffer */
1607 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1608 Py_TPFLAGS_BASETYPE, /* tp_flags */
1609 islice_doc, /* tp_doc */
1610 (traverseproc)islice_traverse, /* tp_traverse */
1611 0, /* tp_clear */
1612 0, /* tp_richcompare */
1613 0, /* tp_weaklistoffset */
1614 PyObject_SelfIter, /* tp_iter */
1615 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001616 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001617 0, /* tp_members */
1618 0, /* tp_getset */
1619 0, /* tp_base */
1620 0, /* tp_dict */
1621 0, /* tp_descr_get */
1622 0, /* tp_descr_set */
1623 0, /* tp_dictoffset */
1624 0, /* tp_init */
1625 0, /* tp_alloc */
1626 islice_new, /* tp_new */
1627 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001628};
1629
1630
1631/* starmap object ************************************************************/
1632
1633typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001634 PyObject_HEAD
1635 PyObject *func;
1636 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001637} starmapobject;
1638
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001639static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001640
1641static PyObject *
1642starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1643{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001644 PyObject *func, *seq;
1645 PyObject *it;
1646 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001647
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001648 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1649 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001650
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001651 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1652 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001653
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001654 /* Get iterator. */
1655 it = PyObject_GetIter(seq);
1656 if (it == NULL)
1657 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001658
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001659 /* create starmapobject structure */
1660 lz = (starmapobject *)type->tp_alloc(type, 0);
1661 if (lz == NULL) {
1662 Py_DECREF(it);
1663 return NULL;
1664 }
1665 Py_INCREF(func);
1666 lz->func = func;
1667 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001668
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001669 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001670}
1671
1672static void
1673starmap_dealloc(starmapobject *lz)
1674{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001675 PyObject_GC_UnTrack(lz);
1676 Py_XDECREF(lz->func);
1677 Py_XDECREF(lz->it);
1678 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001679}
1680
1681static int
1682starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1683{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001684 Py_VISIT(lz->it);
1685 Py_VISIT(lz->func);
1686 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001687}
1688
1689static PyObject *
1690starmap_next(starmapobject *lz)
1691{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001692 PyObject *args;
1693 PyObject *result;
1694 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001695
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001696 args = (*Py_TYPE(it)->tp_iternext)(it);
1697 if (args == NULL)
1698 return NULL;
1699 if (!PyTuple_CheckExact(args)) {
1700 PyObject *newargs = PySequence_Tuple(args);
1701 Py_DECREF(args);
1702 if (newargs == NULL)
1703 return NULL;
1704 args = newargs;
1705 }
1706 result = PyObject_Call(lz->func, args, NULL);
1707 Py_DECREF(args);
1708 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001709}
1710
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001711static PyObject *
1712starmap_reduce(starmapobject *lz)
1713{
1714 /* Just pickle the iterator */
1715 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1716}
1717
1718static PyMethodDef starmap_methods[] = {
1719 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1720 reduce_doc},
1721 {NULL, NULL} /* sentinel */
1722};
1723
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001724PyDoc_STRVAR(starmap_doc,
1725"starmap(function, sequence) --> starmap object\n\
1726\n\
1727Return an iterator whose values are returned from the function evaluated\n\
1728with a argument tuple taken from the given sequence.");
1729
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001730static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001731 PyVarObject_HEAD_INIT(NULL, 0)
1732 "itertools.starmap", /* tp_name */
1733 sizeof(starmapobject), /* tp_basicsize */
1734 0, /* tp_itemsize */
1735 /* methods */
1736 (destructor)starmap_dealloc, /* tp_dealloc */
1737 0, /* tp_print */
1738 0, /* tp_getattr */
1739 0, /* tp_setattr */
1740 0, /* tp_reserved */
1741 0, /* tp_repr */
1742 0, /* tp_as_number */
1743 0, /* tp_as_sequence */
1744 0, /* tp_as_mapping */
1745 0, /* tp_hash */
1746 0, /* tp_call */
1747 0, /* tp_str */
1748 PyObject_GenericGetAttr, /* tp_getattro */
1749 0, /* tp_setattro */
1750 0, /* tp_as_buffer */
1751 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1752 Py_TPFLAGS_BASETYPE, /* tp_flags */
1753 starmap_doc, /* tp_doc */
1754 (traverseproc)starmap_traverse, /* tp_traverse */
1755 0, /* tp_clear */
1756 0, /* tp_richcompare */
1757 0, /* tp_weaklistoffset */
1758 PyObject_SelfIter, /* tp_iter */
1759 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001760 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001761 0, /* tp_members */
1762 0, /* tp_getset */
1763 0, /* tp_base */
1764 0, /* tp_dict */
1765 0, /* tp_descr_get */
1766 0, /* tp_descr_set */
1767 0, /* tp_dictoffset */
1768 0, /* tp_init */
1769 0, /* tp_alloc */
1770 starmap_new, /* tp_new */
1771 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001772};
1773
1774
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001775/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001776
1777typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001778 PyObject_HEAD
1779 PyObject *source; /* Iterator over input iterables */
1780 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001781} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001782
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001783static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001784
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001785static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001786chain_new_internal(PyTypeObject *type, PyObject *source)
1787{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001788 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001789
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001790 lz = (chainobject *)type->tp_alloc(type, 0);
1791 if (lz == NULL) {
1792 Py_DECREF(source);
1793 return NULL;
1794 }
1795
1796 lz->source = source;
1797 lz->active = NULL;
1798 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001799}
1800
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001801static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001802chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001803{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001804 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001805
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001806 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1807 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001808
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001809 source = PyObject_GetIter(args);
1810 if (source == NULL)
1811 return NULL;
1812
1813 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001814}
1815
1816static PyObject *
1817chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1818{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001819 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001820
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001821 source = PyObject_GetIter(arg);
1822 if (source == NULL)
1823 return NULL;
1824
1825 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001826}
1827
1828static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001829chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001830{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001831 PyObject_GC_UnTrack(lz);
1832 Py_XDECREF(lz->active);
1833 Py_XDECREF(lz->source);
1834 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001835}
1836
Raymond Hettinger2012f172003-02-07 05:32:58 +00001837static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001838chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001839{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001840 Py_VISIT(lz->source);
1841 Py_VISIT(lz->active);
1842 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001843}
1844
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001845static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001846chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001847{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001848 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001849
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001850 if (lz->source == NULL)
1851 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001852
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001853 if (lz->active == NULL) {
1854 PyObject *iterable = PyIter_Next(lz->source);
1855 if (iterable == NULL) {
1856 Py_CLEAR(lz->source);
1857 return NULL; /* no more input sources */
1858 }
1859 lz->active = PyObject_GetIter(iterable);
1860 Py_DECREF(iterable);
1861 if (lz->active == NULL) {
1862 Py_CLEAR(lz->source);
1863 return NULL; /* input not iterable */
1864 }
1865 }
1866 item = PyIter_Next(lz->active);
1867 if (item != NULL)
1868 return item;
1869 if (PyErr_Occurred()) {
1870 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1871 PyErr_Clear();
1872 else
1873 return NULL; /* input raised an exception */
1874 }
1875 Py_CLEAR(lz->active);
1876 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001877}
1878
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001879static PyObject *
1880chain_reduce(chainobject *lz)
1881{
1882 if (lz->source) {
1883 /* we can't pickle function objects (itertools.from_iterable) so
1884 * we must use setstate to replace the iterable. One day we
1885 * will fix pickling of functions
1886 */
1887 if (lz->active) {
1888 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1889 } else {
1890 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1891 }
1892 } else {
1893 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1894 }
1895 return NULL;
1896}
1897
1898static PyObject *
1899chain_setstate(chainobject *lz, PyObject *state)
1900{
1901 PyObject *source, *active=NULL;
1902 if (! PyArg_ParseTuple(state, "O|O", &source, &active))
1903 return NULL;
1904
1905 Py_CLEAR(lz->source);
1906 lz->source = source;
1907 Py_INCREF(lz->source);
1908 Py_CLEAR(lz->active);
1909 lz->active = active;
1910 Py_XINCREF(lz->active);
1911 Py_RETURN_NONE;
1912}
1913
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001914PyDoc_STRVAR(chain_doc,
1915"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001916\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001917Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001918first iterable until it is exhausted, then elements from the next\n\
1919iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001920
Christian Heimesf16baeb2008-02-29 14:57:44 +00001921PyDoc_STRVAR(chain_from_iterable_doc,
1922"chain.from_iterable(iterable) --> chain object\n\
1923\n\
1924Alternate chain() contructor taking a single iterable argument\n\
1925that evaluates lazily.");
1926
1927static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001928 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1929 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001930 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1931 reduce_doc},
1932 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1933 setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001934 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001935};
1936
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001937static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001938 PyVarObject_HEAD_INIT(NULL, 0)
1939 "itertools.chain", /* tp_name */
1940 sizeof(chainobject), /* tp_basicsize */
1941 0, /* tp_itemsize */
1942 /* methods */
1943 (destructor)chain_dealloc, /* tp_dealloc */
1944 0, /* tp_print */
1945 0, /* tp_getattr */
1946 0, /* tp_setattr */
1947 0, /* tp_reserved */
1948 0, /* tp_repr */
1949 0, /* tp_as_number */
1950 0, /* tp_as_sequence */
1951 0, /* tp_as_mapping */
1952 0, /* tp_hash */
1953 0, /* tp_call */
1954 0, /* tp_str */
1955 PyObject_GenericGetAttr, /* tp_getattro */
1956 0, /* tp_setattro */
1957 0, /* tp_as_buffer */
1958 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1959 Py_TPFLAGS_BASETYPE, /* tp_flags */
1960 chain_doc, /* tp_doc */
1961 (traverseproc)chain_traverse, /* tp_traverse */
1962 0, /* tp_clear */
1963 0, /* tp_richcompare */
1964 0, /* tp_weaklistoffset */
1965 PyObject_SelfIter, /* tp_iter */
1966 (iternextfunc)chain_next, /* tp_iternext */
1967 chain_methods, /* tp_methods */
1968 0, /* tp_members */
1969 0, /* tp_getset */
1970 0, /* tp_base */
1971 0, /* tp_dict */
1972 0, /* tp_descr_get */
1973 0, /* tp_descr_set */
1974 0, /* tp_dictoffset */
1975 0, /* tp_init */
1976 0, /* tp_alloc */
1977 chain_new, /* tp_new */
1978 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001979};
1980
1981
Christian Heimesc3f30c42008-02-22 16:37:40 +00001982/* product object ************************************************************/
1983
1984typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001985 PyObject_HEAD
1986 PyObject *pools; /* tuple of pool tuples */
1987 Py_ssize_t *indices; /* one index per pool */
1988 PyObject *result; /* most recently returned result tuple */
1989 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001990} productobject;
1991
1992static PyTypeObject product_type;
1993
1994static PyObject *
1995product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1996{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001997 productobject *lz;
1998 Py_ssize_t nargs, npools, repeat=1;
1999 PyObject *pools = NULL;
2000 Py_ssize_t *indices = NULL;
2001 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002002
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002003 if (kwds != NULL) {
2004 char *kwlist[] = {"repeat", 0};
2005 PyObject *tmpargs = PyTuple_New(0);
2006 if (tmpargs == NULL)
2007 return NULL;
2008 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
2009 Py_DECREF(tmpargs);
2010 return NULL;
2011 }
2012 Py_DECREF(tmpargs);
2013 if (repeat < 0) {
2014 PyErr_SetString(PyExc_ValueError,
2015 "repeat argument cannot be negative");
2016 return NULL;
2017 }
2018 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002019
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002020 assert(PyTuple_CheckExact(args));
2021 if (repeat == 0) {
2022 nargs = 0;
2023 } else {
2024 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002025 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002026 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2027 return NULL;
2028 }
2029 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002030 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002031
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002032 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002033 if (indices == NULL) {
2034 PyErr_NoMemory();
2035 goto error;
2036 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002037
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002038 pools = PyTuple_New(npools);
2039 if (pools == NULL)
2040 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002041
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002042 for (i=0; i < nargs ; ++i) {
2043 PyObject *item = PyTuple_GET_ITEM(args, i);
2044 PyObject *pool = PySequence_Tuple(item);
2045 if (pool == NULL)
2046 goto error;
2047 PyTuple_SET_ITEM(pools, i, pool);
2048 indices[i] = 0;
2049 }
2050 for ( ; i < npools; ++i) {
2051 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2052 Py_INCREF(pool);
2053 PyTuple_SET_ITEM(pools, i, pool);
2054 indices[i] = 0;
2055 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002056
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002057 /* create productobject structure */
2058 lz = (productobject *)type->tp_alloc(type, 0);
2059 if (lz == NULL)
2060 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002061
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002062 lz->pools = pools;
2063 lz->indices = indices;
2064 lz->result = NULL;
2065 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002066
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002067 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002068
2069error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002070 if (indices != NULL)
2071 PyMem_Free(indices);
2072 Py_XDECREF(pools);
2073 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002074}
2075
2076static void
2077product_dealloc(productobject *lz)
2078{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002079 PyObject_GC_UnTrack(lz);
2080 Py_XDECREF(lz->pools);
2081 Py_XDECREF(lz->result);
2082 if (lz->indices != NULL)
2083 PyMem_Free(lz->indices);
2084 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002085}
2086
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002087static PyObject *
2088product_sizeof(productobject *lz, void *unused)
2089{
2090 Py_ssize_t res;
2091
2092 res = sizeof(productobject);
2093 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2094 return PyLong_FromSsize_t(res);
2095}
2096
2097PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2098
Christian Heimesc3f30c42008-02-22 16:37:40 +00002099static int
2100product_traverse(productobject *lz, visitproc visit, void *arg)
2101{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002102 Py_VISIT(lz->pools);
2103 Py_VISIT(lz->result);
2104 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002105}
2106
2107static PyObject *
2108product_next(productobject *lz)
2109{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002110 PyObject *pool;
2111 PyObject *elem;
2112 PyObject *oldelem;
2113 PyObject *pools = lz->pools;
2114 PyObject *result = lz->result;
2115 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2116 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002117
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002118 if (lz->stopped)
2119 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002120
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002121 if (result == NULL) {
2122 /* On the first pass, return an initial tuple filled with the
2123 first element from each pool. */
2124 result = PyTuple_New(npools);
2125 if (result == NULL)
2126 goto empty;
2127 lz->result = result;
2128 for (i=0; i < npools; i++) {
2129 pool = PyTuple_GET_ITEM(pools, i);
2130 if (PyTuple_GET_SIZE(pool) == 0)
2131 goto empty;
2132 elem = PyTuple_GET_ITEM(pool, 0);
2133 Py_INCREF(elem);
2134 PyTuple_SET_ITEM(result, i, elem);
2135 }
2136 } else {
2137 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002138
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002139 /* Copy the previous result tuple or re-use it if available */
2140 if (Py_REFCNT(result) > 1) {
2141 PyObject *old_result = result;
2142 result = PyTuple_New(npools);
2143 if (result == NULL)
2144 goto empty;
2145 lz->result = result;
2146 for (i=0; i < npools; i++) {
2147 elem = PyTuple_GET_ITEM(old_result, i);
2148 Py_INCREF(elem);
2149 PyTuple_SET_ITEM(result, i, elem);
2150 }
2151 Py_DECREF(old_result);
2152 }
2153 /* Now, we've got the only copy so we can update it in-place */
2154 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002155
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002156 /* Update the pool indices right-to-left. Only advance to the
2157 next pool when the previous one rolls-over */
2158 for (i=npools-1 ; i >= 0 ; i--) {
2159 pool = PyTuple_GET_ITEM(pools, i);
2160 indices[i]++;
2161 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2162 /* Roll-over and advance to next pool */
2163 indices[i] = 0;
2164 elem = PyTuple_GET_ITEM(pool, 0);
2165 Py_INCREF(elem);
2166 oldelem = PyTuple_GET_ITEM(result, i);
2167 PyTuple_SET_ITEM(result, i, elem);
2168 Py_DECREF(oldelem);
2169 } else {
2170 /* No rollover. Just increment and stop here. */
2171 elem = PyTuple_GET_ITEM(pool, indices[i]);
2172 Py_INCREF(elem);
2173 oldelem = PyTuple_GET_ITEM(result, i);
2174 PyTuple_SET_ITEM(result, i, elem);
2175 Py_DECREF(oldelem);
2176 break;
2177 }
2178 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002179
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002180 /* If i is negative, then the indices have all rolled-over
2181 and we're done. */
2182 if (i < 0)
2183 goto empty;
2184 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002185
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002186 Py_INCREF(result);
2187 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002188
2189empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002190 lz->stopped = 1;
2191 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002192}
2193
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002194static PyObject *
2195product_reduce(productobject *lz)
2196{
2197 if (lz->stopped) {
2198 return Py_BuildValue("O(())", Py_TYPE(lz));
2199 } else if (lz->result == NULL) {
2200 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2201 } else {
2202 PyObject *indices;
2203 Py_ssize_t n, i;
2204
2205 /* we must pickle the indices use them for setstate, and
2206 * additionally indicate that the iterator has started
2207 */
2208 n = PyTuple_GET_SIZE(lz->pools);
2209 indices = PyTuple_New(n);
2210 if (indices == NULL)
2211 return NULL;
2212 for (i=0; i<n; i++){
2213 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2214 if (!index) {
2215 Py_DECREF(indices);
2216 return NULL;
2217 }
2218 PyTuple_SET_ITEM(indices, i, index);
2219 }
2220 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2221 }
2222}
2223
2224static PyObject *
2225product_setstate(productobject *lz, PyObject *state)
2226{
2227 PyObject *result;
2228 Py_ssize_t n, i;
2229
2230 n = PyTuple_GET_SIZE(lz->pools);
2231 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2232 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2233 return NULL;
2234 }
2235 for (i=0; i<n; i++)
2236 {
2237 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2238 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2239 if (index < 0 && PyErr_Occurred())
2240 return NULL; /* not an integer */
2241 /* clamp the index */
2242 if (index < 0)
2243 index = 0;
2244 else if (index > n-1)
2245 index = n-1;
2246 lz->indices[i] = index;
2247 }
2248
2249 result = PyTuple_New(n);
2250 if (!result)
2251 return NULL;
2252 for (i=0; i<n; i++) {
2253 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2254 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2255 Py_INCREF(element);
2256 PyTuple_SET_ITEM(result, i, element);
2257 }
2258 Py_CLEAR(lz->result);
2259 lz->result = result;
2260 Py_RETURN_NONE;
2261}
2262
2263static PyMethodDef product_methods[] = {
2264 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2265 reduce_doc},
2266 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2267 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002268 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2269 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002270 {NULL, NULL} /* sentinel */
2271};
2272
Christian Heimesc3f30c42008-02-22 16:37:40 +00002273PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002274"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002275\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002276Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002277For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2278The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2279cycle in a manner similar to an odometer (with the rightmost element changing\n\
2280on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002281To compute the product of an iterable with itself, specify the number\n\
2282of repetitions with the optional repeat keyword argument. For example,\n\
2283product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002284product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2285product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2286
2287static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002288 PyVarObject_HEAD_INIT(NULL, 0)
2289 "itertools.product", /* tp_name */
2290 sizeof(productobject), /* tp_basicsize */
2291 0, /* tp_itemsize */
2292 /* methods */
2293 (destructor)product_dealloc, /* tp_dealloc */
2294 0, /* tp_print */
2295 0, /* tp_getattr */
2296 0, /* tp_setattr */
2297 0, /* tp_reserved */
2298 0, /* tp_repr */
2299 0, /* tp_as_number */
2300 0, /* tp_as_sequence */
2301 0, /* tp_as_mapping */
2302 0, /* tp_hash */
2303 0, /* tp_call */
2304 0, /* tp_str */
2305 PyObject_GenericGetAttr, /* tp_getattro */
2306 0, /* tp_setattro */
2307 0, /* tp_as_buffer */
2308 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2309 Py_TPFLAGS_BASETYPE, /* tp_flags */
2310 product_doc, /* tp_doc */
2311 (traverseproc)product_traverse, /* tp_traverse */
2312 0, /* tp_clear */
2313 0, /* tp_richcompare */
2314 0, /* tp_weaklistoffset */
2315 PyObject_SelfIter, /* tp_iter */
2316 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002317 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002318 0, /* tp_members */
2319 0, /* tp_getset */
2320 0, /* tp_base */
2321 0, /* tp_dict */
2322 0, /* tp_descr_get */
2323 0, /* tp_descr_set */
2324 0, /* tp_dictoffset */
2325 0, /* tp_init */
2326 0, /* tp_alloc */
2327 product_new, /* tp_new */
2328 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002329};
2330
2331
Christian Heimes380f7f22008-02-28 11:19:05 +00002332/* combinations object ************************************************************/
2333
2334typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002335 PyObject_HEAD
2336 PyObject *pool; /* input converted to a tuple */
2337 Py_ssize_t *indices; /* one index per result element */
2338 PyObject *result; /* most recently returned result tuple */
2339 Py_ssize_t r; /* size of result tuple */
2340 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002341} combinationsobject;
2342
2343static PyTypeObject combinations_type;
2344
2345static PyObject *
2346combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2347{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002348 combinationsobject *co;
2349 Py_ssize_t n;
2350 Py_ssize_t r;
2351 PyObject *pool = NULL;
2352 PyObject *iterable = NULL;
2353 Py_ssize_t *indices = NULL;
2354 Py_ssize_t i;
2355 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002356
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002357 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2358 &iterable, &r))
2359 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002360
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002361 pool = PySequence_Tuple(iterable);
2362 if (pool == NULL)
2363 goto error;
2364 n = PyTuple_GET_SIZE(pool);
2365 if (r < 0) {
2366 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2367 goto error;
2368 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002369
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002370 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002371 if (indices == NULL) {
2372 PyErr_NoMemory();
2373 goto error;
2374 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002375
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002376 for (i=0 ; i<r ; i++)
2377 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002378
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002379 /* create combinationsobject structure */
2380 co = (combinationsobject *)type->tp_alloc(type, 0);
2381 if (co == NULL)
2382 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002383
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002384 co->pool = pool;
2385 co->indices = indices;
2386 co->result = NULL;
2387 co->r = r;
2388 co->stopped = r > n ? 1 : 0;
2389
2390 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002391
2392error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002393 if (indices != NULL)
2394 PyMem_Free(indices);
2395 Py_XDECREF(pool);
2396 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002397}
2398
2399static void
2400combinations_dealloc(combinationsobject *co)
2401{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002402 PyObject_GC_UnTrack(co);
2403 Py_XDECREF(co->pool);
2404 Py_XDECREF(co->result);
2405 if (co->indices != NULL)
2406 PyMem_Free(co->indices);
2407 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002408}
2409
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002410static PyObject *
2411combinations_sizeof(combinationsobject *co, void *unused)
2412{
2413 Py_ssize_t res;
2414
2415 res = sizeof(combinationsobject);
2416 res += co->r * sizeof(Py_ssize_t);
2417 return PyLong_FromSsize_t(res);
2418}
2419
Christian Heimes380f7f22008-02-28 11:19:05 +00002420static int
2421combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2422{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002423 Py_VISIT(co->pool);
2424 Py_VISIT(co->result);
2425 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002426}
2427
2428static PyObject *
2429combinations_next(combinationsobject *co)
2430{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002431 PyObject *elem;
2432 PyObject *oldelem;
2433 PyObject *pool = co->pool;
2434 Py_ssize_t *indices = co->indices;
2435 PyObject *result = co->result;
2436 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2437 Py_ssize_t r = co->r;
2438 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002439
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002440 if (co->stopped)
2441 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002442
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002443 if (result == NULL) {
2444 /* On the first pass, initialize result tuple using the indices */
2445 result = PyTuple_New(r);
2446 if (result == NULL)
2447 goto empty;
2448 co->result = result;
2449 for (i=0; i<r ; i++) {
2450 index = indices[i];
2451 elem = PyTuple_GET_ITEM(pool, index);
2452 Py_INCREF(elem);
2453 PyTuple_SET_ITEM(result, i, elem);
2454 }
2455 } else {
2456 /* Copy the previous result tuple or re-use it if available */
2457 if (Py_REFCNT(result) > 1) {
2458 PyObject *old_result = result;
2459 result = PyTuple_New(r);
2460 if (result == NULL)
2461 goto empty;
2462 co->result = result;
2463 for (i=0; i<r ; i++) {
2464 elem = PyTuple_GET_ITEM(old_result, i);
2465 Py_INCREF(elem);
2466 PyTuple_SET_ITEM(result, i, elem);
2467 }
2468 Py_DECREF(old_result);
2469 }
2470 /* Now, we've got the only copy so we can update it in-place
2471 * CPython's empty tuple is a singleton and cached in
2472 * PyTuple's freelist.
2473 */
2474 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002475
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002476 /* Scan indices right-to-left until finding one that is not
2477 at its maximum (i + n - r). */
2478 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2479 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002480
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002481 /* If i is negative, then the indices are all at
2482 their maximum value and we're done. */
2483 if (i < 0)
2484 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002485
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002486 /* Increment the current index which we know is not at its
2487 maximum. Then move back to the right setting each index
2488 to its lowest possible value (one higher than the index
2489 to its left -- this maintains the sort order invariant). */
2490 indices[i]++;
2491 for (j=i+1 ; j<r ; j++)
2492 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002493
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002494 /* Update the result tuple for the new indices
2495 starting with i, the leftmost index that changed */
2496 for ( ; i<r ; i++) {
2497 index = indices[i];
2498 elem = PyTuple_GET_ITEM(pool, index);
2499 Py_INCREF(elem);
2500 oldelem = PyTuple_GET_ITEM(result, i);
2501 PyTuple_SET_ITEM(result, i, elem);
2502 Py_DECREF(oldelem);
2503 }
2504 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002505
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002506 Py_INCREF(result);
2507 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002508
2509empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002510 co->stopped = 1;
2511 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002512}
2513
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002514static PyObject *
2515combinations_reduce(combinationsobject *lz)
2516{
2517 if (lz->result == NULL) {
2518 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2519 } else if (lz->stopped) {
2520 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2521 } else {
2522 PyObject *indices;
2523 Py_ssize_t i;
2524
2525 /* we must pickle the indices and use them for setstate */
2526 indices = PyTuple_New(lz->r);
2527 if (!indices)
2528 return NULL;
2529 for (i=0; i<lz->r; i++)
2530 {
2531 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2532 if (!index) {
2533 Py_DECREF(indices);
2534 return NULL;
2535 }
2536 PyTuple_SET_ITEM(indices, i, index);
2537 }
2538
2539 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2540 }
2541}
2542
2543static PyObject *
2544combinations_setstate(combinationsobject *lz, PyObject *state)
2545{
2546 PyObject *result;
2547 Py_ssize_t i;
2548 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2549
2550 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2551 {
2552 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2553 return NULL;
2554 }
2555
2556 for (i=0; i<lz->r; i++)
2557 {
2558 Py_ssize_t max;
2559 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2560 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2561 if (index == -1 && PyErr_Occurred())
2562 return NULL; /* not an integer */
2563 max = i + n - lz->r;
2564 /* clamp the index (beware of negative max) */
2565 if (index > max)
2566 index = max;
2567 if (index < 0)
2568 index = 0;
2569 lz->indices[i] = index;
2570 }
2571
2572 result = PyTuple_New(lz->r);
2573 if (result == NULL)
2574 return NULL;
2575 for (i=0; i<lz->r; i++) {
2576 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2577 Py_INCREF(element);
2578 PyTuple_SET_ITEM(result, i, element);
2579 }
2580
2581 Py_CLEAR(lz->result);
2582 lz->result = result;
2583 Py_RETURN_NONE;
2584}
2585
2586static PyMethodDef combinations_methods[] = {
2587 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2588 reduce_doc},
2589 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2590 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002591 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2592 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002593 {NULL, NULL} /* sentinel */
2594};
2595
Christian Heimes380f7f22008-02-28 11:19:05 +00002596PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002597"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002598\n\
2599Return successive r-length combinations of elements in the iterable.\n\n\
2600combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2601
2602static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002603 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002604 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002605 sizeof(combinationsobject), /* tp_basicsize */
2606 0, /* tp_itemsize */
2607 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002608 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002609 0, /* tp_print */
2610 0, /* tp_getattr */
2611 0, /* tp_setattr */
2612 0, /* tp_reserved */
2613 0, /* tp_repr */
2614 0, /* tp_as_number */
2615 0, /* tp_as_sequence */
2616 0, /* tp_as_mapping */
2617 0, /* tp_hash */
2618 0, /* tp_call */
2619 0, /* tp_str */
2620 PyObject_GenericGetAttr, /* tp_getattro */
2621 0, /* tp_setattro */
2622 0, /* tp_as_buffer */
2623 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2624 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002625 combinations_doc, /* tp_doc */
2626 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002627 0, /* tp_clear */
2628 0, /* tp_richcompare */
2629 0, /* tp_weaklistoffset */
2630 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002631 (iternextfunc)combinations_next, /* tp_iternext */
2632 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002633 0, /* tp_members */
2634 0, /* tp_getset */
2635 0, /* tp_base */
2636 0, /* tp_dict */
2637 0, /* tp_descr_get */
2638 0, /* tp_descr_set */
2639 0, /* tp_dictoffset */
2640 0, /* tp_init */
2641 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002642 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002643 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002644};
2645
2646
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002647/* combinations with replacement object *******************************************/
2648
2649/* Equivalent to:
2650
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002651 def combinations_with_replacement(iterable, r):
2652 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2653 # number items returned: (n+r-1)! / r! / (n-1)!
2654 pool = tuple(iterable)
2655 n = len(pool)
2656 indices = [0] * r
2657 yield tuple(pool[i] for i in indices)
2658 while 1:
2659 for i in reversed(range(r)):
2660 if indices[i] != n - 1:
2661 break
2662 else:
2663 return
2664 indices[i:] = [indices[i] + 1] * (r - i)
2665 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002666
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002667 def combinations_with_replacement2(iterable, r):
2668 'Alternate version that filters from product()'
2669 pool = tuple(iterable)
2670 n = len(pool)
2671 for indices in product(range(n), repeat=r):
2672 if sorted(indices) == list(indices):
2673 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002674*/
2675typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002676 PyObject_HEAD
2677 PyObject *pool; /* input converted to a tuple */
2678 Py_ssize_t *indices; /* one index per result element */
2679 PyObject *result; /* most recently returned result tuple */
2680 Py_ssize_t r; /* size of result tuple */
2681 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002682} cwrobject;
2683
2684static PyTypeObject cwr_type;
2685
2686static PyObject *
2687cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2688{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002689 cwrobject *co;
2690 Py_ssize_t n;
2691 Py_ssize_t r;
2692 PyObject *pool = NULL;
2693 PyObject *iterable = NULL;
2694 Py_ssize_t *indices = NULL;
2695 Py_ssize_t i;
2696 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002697
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002698 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2699 &iterable, &r))
2700 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002701
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002702 pool = PySequence_Tuple(iterable);
2703 if (pool == NULL)
2704 goto error;
2705 n = PyTuple_GET_SIZE(pool);
2706 if (r < 0) {
2707 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2708 goto error;
2709 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002710
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002711 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002712 if (indices == NULL) {
2713 PyErr_NoMemory();
2714 goto error;
2715 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002716
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002717 for (i=0 ; i<r ; i++)
2718 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002719
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002720 /* create cwrobject structure */
2721 co = (cwrobject *)type->tp_alloc(type, 0);
2722 if (co == NULL)
2723 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002724
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002725 co->pool = pool;
2726 co->indices = indices;
2727 co->result = NULL;
2728 co->r = r;
2729 co->stopped = !n && r;
2730
2731 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002732
2733error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002734 if (indices != NULL)
2735 PyMem_Free(indices);
2736 Py_XDECREF(pool);
2737 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002738}
2739
2740static void
2741cwr_dealloc(cwrobject *co)
2742{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002743 PyObject_GC_UnTrack(co);
2744 Py_XDECREF(co->pool);
2745 Py_XDECREF(co->result);
2746 if (co->indices != NULL)
2747 PyMem_Free(co->indices);
2748 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002749}
2750
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002751static PyObject *
2752cwr_sizeof(cwrobject *co, void *unused)
2753{
2754 Py_ssize_t res;
2755
2756 res = sizeof(cwrobject);
2757 res += co->r * sizeof(Py_ssize_t);
2758 return PyLong_FromSsize_t(res);
2759}
2760
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002761static int
2762cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2763{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002764 Py_VISIT(co->pool);
2765 Py_VISIT(co->result);
2766 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002767}
2768
2769static PyObject *
2770cwr_next(cwrobject *co)
2771{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002772 PyObject *elem;
2773 PyObject *oldelem;
2774 PyObject *pool = co->pool;
2775 Py_ssize_t *indices = co->indices;
2776 PyObject *result = co->result;
2777 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2778 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002779 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002780
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002781 if (co->stopped)
2782 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002783
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002784 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002785 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002786 result = PyTuple_New(r);
2787 if (result == NULL)
2788 goto empty;
2789 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002790 if (n > 0) {
2791 elem = PyTuple_GET_ITEM(pool, 0);
2792 for (i=0; i<r ; i++) {
2793 assert(indices[i] == 0);
2794 Py_INCREF(elem);
2795 PyTuple_SET_ITEM(result, i, elem);
2796 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002797 }
2798 } else {
2799 /* Copy the previous result tuple or re-use it if available */
2800 if (Py_REFCNT(result) > 1) {
2801 PyObject *old_result = result;
2802 result = PyTuple_New(r);
2803 if (result == NULL)
2804 goto empty;
2805 co->result = result;
2806 for (i=0; i<r ; i++) {
2807 elem = PyTuple_GET_ITEM(old_result, i);
2808 Py_INCREF(elem);
2809 PyTuple_SET_ITEM(result, i, elem);
2810 }
2811 Py_DECREF(old_result);
2812 }
2813 /* Now, we've got the only copy so we can update it in-place CPython's
2814 empty tuple is a singleton and cached in PyTuple's freelist. */
2815 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002816
Tim Peters9edb1682013-09-03 11:49:31 -05002817 /* Scan indices right-to-left until finding one that is not
2818 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002819 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2820 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002821
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002822 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002823 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002824 if (i < 0)
2825 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002826
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002827 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002828 maximum. Then set all to the right to the same value. */
2829 index = indices[i] + 1;
2830 assert(index < n);
2831 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002832 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002833 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002834 Py_INCREF(elem);
2835 oldelem = PyTuple_GET_ITEM(result, i);
2836 PyTuple_SET_ITEM(result, i, elem);
2837 Py_DECREF(oldelem);
2838 }
2839 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002840
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002841 Py_INCREF(result);
2842 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002843
2844empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002845 co->stopped = 1;
2846 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002847}
2848
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002849static PyObject *
2850cwr_reduce(cwrobject *lz)
2851{
2852 if (lz->result == NULL) {
2853 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2854 } else if (lz->stopped) {
2855 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2856 } else {
2857 PyObject *indices;
2858 Py_ssize_t i;
2859
2860 /* we must pickle the indices and use them for setstate */
2861 indices = PyTuple_New(lz->r);
2862 if (!indices)
2863 return NULL;
2864 for (i=0; i<lz->r; i++)
2865 {
2866 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2867 if (!index) {
2868 Py_DECREF(indices);
2869 return NULL;
2870 }
2871 PyTuple_SET_ITEM(indices, i, index);
2872 }
2873
2874 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2875 }
2876}
2877
2878static PyObject *
2879cwr_setstate(cwrobject *lz, PyObject *state)
2880{
2881 PyObject *result;
2882 Py_ssize_t n, i;
2883
2884 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2885 {
2886 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2887 return NULL;
2888 }
2889
2890 n = PyTuple_GET_SIZE(lz->pool);
2891 for (i=0; i<lz->r; i++)
2892 {
2893 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2894 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2895 if (index < 0 && PyErr_Occurred())
2896 return NULL; /* not an integer */
2897 /* clamp the index */
2898 if (index < 0)
2899 index = 0;
2900 else if (index > n-1)
2901 index = n-1;
2902 lz->indices[i] = index;
2903 }
2904 result = PyTuple_New(lz->r);
2905 if (result == NULL)
2906 return NULL;
2907 for (i=0; i<lz->r; i++) {
2908 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2909 Py_INCREF(element);
2910 PyTuple_SET_ITEM(result, i, element);
2911 }
2912 Py_CLEAR(lz->result);
2913 lz->result = result;
2914 Py_RETURN_NONE;
2915}
2916
2917static PyMethodDef cwr_methods[] = {
2918 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2919 reduce_doc},
2920 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2921 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002922 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
2923 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002924 {NULL, NULL} /* sentinel */
2925};
2926
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002927PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002928"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002929\n\
2930Return successive r-length combinations of elements in the iterable\n\
2931allowing individual elements to have successive repeats.\n\
2932combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2933
2934static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002935 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002936 "itertools.combinations_with_replacement", /* tp_name */
2937 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002938 0, /* tp_itemsize */
2939 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002940 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002941 0, /* tp_print */
2942 0, /* tp_getattr */
2943 0, /* tp_setattr */
2944 0, /* tp_reserved */
2945 0, /* tp_repr */
2946 0, /* tp_as_number */
2947 0, /* tp_as_sequence */
2948 0, /* tp_as_mapping */
2949 0, /* tp_hash */
2950 0, /* tp_call */
2951 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002952 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002953 0, /* tp_setattro */
2954 0, /* tp_as_buffer */
2955 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002956 Py_TPFLAGS_BASETYPE, /* tp_flags */
2957 cwr_doc, /* tp_doc */
2958 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002959 0, /* tp_clear */
2960 0, /* tp_richcompare */
2961 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002962 PyObject_SelfIter, /* tp_iter */
2963 (iternextfunc)cwr_next, /* tp_iternext */
2964 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002965 0, /* tp_members */
2966 0, /* tp_getset */
2967 0, /* tp_base */
2968 0, /* tp_dict */
2969 0, /* tp_descr_get */
2970 0, /* tp_descr_set */
2971 0, /* tp_dictoffset */
2972 0, /* tp_init */
2973 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002974 cwr_new, /* tp_new */
2975 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002976};
2977
2978
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002979/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002980
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002981def permutations(iterable, r=None):
2982 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2983 pool = tuple(iterable)
2984 n = len(pool)
2985 r = n if r is None else r
2986 indices = range(n)
2987 cycles = range(n-r+1, n+1)[::-1]
2988 yield tuple(pool[i] for i in indices[:r])
2989 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03002990 for i in reversed(range(r)):
2991 cycles[i] -= 1
2992 if cycles[i] == 0:
2993 indices[i:] = indices[i+1:] + indices[i:i+1]
2994 cycles[i] = n - i
2995 else:
2996 j = cycles[i]
2997 indices[i], indices[-j] = indices[-j], indices[i]
2998 yield tuple(pool[i] for i in indices[:r])
2999 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003000 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003001 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003002*/
3003
3004typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003005 PyObject_HEAD
3006 PyObject *pool; /* input converted to a tuple */
3007 Py_ssize_t *indices; /* one index per element in the pool */
3008 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3009 PyObject *result; /* most recently returned result tuple */
3010 Py_ssize_t r; /* size of result tuple */
3011 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003012} permutationsobject;
3013
3014static PyTypeObject permutations_type;
3015
3016static PyObject *
3017permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3018{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003019 permutationsobject *po;
3020 Py_ssize_t n;
3021 Py_ssize_t r;
3022 PyObject *robj = Py_None;
3023 PyObject *pool = NULL;
3024 PyObject *iterable = NULL;
3025 Py_ssize_t *indices = NULL;
3026 Py_ssize_t *cycles = NULL;
3027 Py_ssize_t i;
3028 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003029
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003030 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
3031 &iterable, &robj))
3032 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003033
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003034 pool = PySequence_Tuple(iterable);
3035 if (pool == NULL)
3036 goto error;
3037 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003038
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003039 r = n;
3040 if (robj != Py_None) {
3041 if (!PyLong_Check(robj)) {
3042 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3043 goto error;
3044 }
3045 r = PyLong_AsSsize_t(robj);
3046 if (r == -1 && PyErr_Occurred())
3047 goto error;
3048 }
3049 if (r < 0) {
3050 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3051 goto error;
3052 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003053
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003054 indices = PyMem_New(Py_ssize_t, n);
3055 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003056 if (indices == NULL || cycles == NULL) {
3057 PyErr_NoMemory();
3058 goto error;
3059 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003060
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003061 for (i=0 ; i<n ; i++)
3062 indices[i] = i;
3063 for (i=0 ; i<r ; i++)
3064 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003065
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003066 /* create permutationsobject structure */
3067 po = (permutationsobject *)type->tp_alloc(type, 0);
3068 if (po == NULL)
3069 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003070
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003071 po->pool = pool;
3072 po->indices = indices;
3073 po->cycles = cycles;
3074 po->result = NULL;
3075 po->r = r;
3076 po->stopped = r > n ? 1 : 0;
3077
3078 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003079
3080error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003081 if (indices != NULL)
3082 PyMem_Free(indices);
3083 if (cycles != NULL)
3084 PyMem_Free(cycles);
3085 Py_XDECREF(pool);
3086 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003087}
3088
3089static void
3090permutations_dealloc(permutationsobject *po)
3091{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003092 PyObject_GC_UnTrack(po);
3093 Py_XDECREF(po->pool);
3094 Py_XDECREF(po->result);
3095 PyMem_Free(po->indices);
3096 PyMem_Free(po->cycles);
3097 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003098}
3099
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003100static PyObject *
3101permutations_sizeof(permutationsobject *po, void *unused)
3102{
3103 Py_ssize_t res;
3104
3105 res = sizeof(permutationsobject);
3106 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3107 res += po->r * sizeof(Py_ssize_t);
3108 return PyLong_FromSsize_t(res);
3109}
3110
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003111static int
3112permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3113{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003114 Py_VISIT(po->pool);
3115 Py_VISIT(po->result);
3116 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003117}
3118
3119static PyObject *
3120permutations_next(permutationsobject *po)
3121{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003122 PyObject *elem;
3123 PyObject *oldelem;
3124 PyObject *pool = po->pool;
3125 Py_ssize_t *indices = po->indices;
3126 Py_ssize_t *cycles = po->cycles;
3127 PyObject *result = po->result;
3128 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3129 Py_ssize_t r = po->r;
3130 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003131
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003132 if (po->stopped)
3133 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003134
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003135 if (result == NULL) {
3136 /* On the first pass, initialize result tuple using the indices */
3137 result = PyTuple_New(r);
3138 if (result == NULL)
3139 goto empty;
3140 po->result = result;
3141 for (i=0; i<r ; i++) {
3142 index = indices[i];
3143 elem = PyTuple_GET_ITEM(pool, index);
3144 Py_INCREF(elem);
3145 PyTuple_SET_ITEM(result, i, elem);
3146 }
3147 } else {
3148 if (n == 0)
3149 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003150
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003151 /* Copy the previous result tuple or re-use it if available */
3152 if (Py_REFCNT(result) > 1) {
3153 PyObject *old_result = result;
3154 result = PyTuple_New(r);
3155 if (result == NULL)
3156 goto empty;
3157 po->result = result;
3158 for (i=0; i<r ; i++) {
3159 elem = PyTuple_GET_ITEM(old_result, i);
3160 Py_INCREF(elem);
3161 PyTuple_SET_ITEM(result, i, elem);
3162 }
3163 Py_DECREF(old_result);
3164 }
3165 /* Now, we've got the only copy so we can update it in-place */
3166 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003167
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003168 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3169 for (i=r-1 ; i>=0 ; i--) {
3170 cycles[i] -= 1;
3171 if (cycles[i] == 0) {
3172 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3173 index = indices[i];
3174 for (j=i ; j<n-1 ; j++)
3175 indices[j] = indices[j+1];
3176 indices[n-1] = index;
3177 cycles[i] = n - i;
3178 } else {
3179 j = cycles[i];
3180 index = indices[i];
3181 indices[i] = indices[n-j];
3182 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003183
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003184 for (k=i; k<r ; k++) {
3185 /* start with i, the leftmost element that changed */
3186 /* yield tuple(pool[k] for k in indices[:r]) */
3187 index = indices[k];
3188 elem = PyTuple_GET_ITEM(pool, index);
3189 Py_INCREF(elem);
3190 oldelem = PyTuple_GET_ITEM(result, k);
3191 PyTuple_SET_ITEM(result, k, elem);
3192 Py_DECREF(oldelem);
3193 }
3194 break;
3195 }
3196 }
3197 /* If i is negative, then the cycles have all
3198 rolled-over and we're done. */
3199 if (i < 0)
3200 goto empty;
3201 }
3202 Py_INCREF(result);
3203 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003204
3205empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003206 po->stopped = 1;
3207 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003208}
3209
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003210static PyObject *
3211permutations_reduce(permutationsobject *po)
3212{
3213 if (po->result == NULL) {
3214 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3215 } else if (po->stopped) {
3216 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3217 } else {
3218 PyObject *indices=NULL, *cycles=NULL;
3219 Py_ssize_t n, i;
3220
3221 /* we must pickle the indices and cycles and use them for setstate */
3222 n = PyTuple_GET_SIZE(po->pool);
3223 indices = PyTuple_New(n);
3224 if (indices == NULL)
3225 goto err;
3226 for (i=0; i<n; i++){
3227 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3228 if (!index)
3229 goto err;
3230 PyTuple_SET_ITEM(indices, i, index);
3231 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003232
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003233 cycles = PyTuple_New(po->r);
3234 if (cycles == NULL)
3235 goto err;
3236 for (i=0; i<po->r; i++)
3237 {
3238 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3239 if (!index)
3240 goto err;
3241 PyTuple_SET_ITEM(cycles, i, index);
3242 }
3243 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3244 po->pool, po->r,
3245 indices, cycles);
3246 err:
3247 Py_XDECREF(indices);
3248 Py_XDECREF(cycles);
3249 return NULL;
3250 }
3251}
3252
3253static PyObject *
3254permutations_setstate(permutationsobject *po, PyObject *state)
3255{
3256 PyObject *indices, *cycles, *result;
3257 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003258
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003259 if (!PyArg_ParseTuple(state, "O!O!",
3260 &PyTuple_Type, &indices,
3261 &PyTuple_Type, &cycles))
3262 return NULL;
3263
3264 n = PyTuple_GET_SIZE(po->pool);
3265 if (PyTuple_GET_SIZE(indices) != n ||
3266 PyTuple_GET_SIZE(cycles) != po->r)
3267 {
3268 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3269 return NULL;
3270 }
3271
3272 for (i=0; i<n; i++)
3273 {
3274 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3275 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3276 if (index < 0 && PyErr_Occurred())
3277 return NULL; /* not an integer */
3278 /* clamp the index */
3279 if (index < 0)
3280 index = 0;
3281 else if (index > n-1)
3282 index = n-1;
3283 po->indices[i] = index;
3284 }
3285
3286 for (i=0; i<po->r; i++)
3287 {
3288 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3289 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3290 if (index < 0 && PyErr_Occurred())
3291 return NULL; /* not an integer */
3292 if (index < 1)
3293 index = 1;
3294 else if (index > n-i)
3295 index = n-i;
3296 po->cycles[i] = index;
3297 }
3298 result = PyTuple_New(po->r);
3299 if (result == NULL)
3300 return NULL;
3301 for (i=0; i<po->r; i++) {
3302 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3303 Py_INCREF(element);
3304 PyTuple_SET_ITEM(result, i, element);
3305 }
3306 Py_CLEAR(po->result);
3307 po->result = result;
3308 Py_RETURN_NONE;
3309}
3310
3311static PyMethodDef permuations_methods[] = {
3312 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3313 reduce_doc},
3314 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3315 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003316 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3317 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003318 {NULL, NULL} /* sentinel */
3319};
3320
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003321PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003322"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003323\n\
3324Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003325permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003326
3327static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003328 PyVarObject_HEAD_INIT(NULL, 0)
3329 "itertools.permutations", /* tp_name */
3330 sizeof(permutationsobject), /* tp_basicsize */
3331 0, /* tp_itemsize */
3332 /* methods */
3333 (destructor)permutations_dealloc, /* tp_dealloc */
3334 0, /* tp_print */
3335 0, /* tp_getattr */
3336 0, /* tp_setattr */
3337 0, /* tp_reserved */
3338 0, /* tp_repr */
3339 0, /* tp_as_number */
3340 0, /* tp_as_sequence */
3341 0, /* tp_as_mapping */
3342 0, /* tp_hash */
3343 0, /* tp_call */
3344 0, /* tp_str */
3345 PyObject_GenericGetAttr, /* tp_getattro */
3346 0, /* tp_setattro */
3347 0, /* tp_as_buffer */
3348 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3349 Py_TPFLAGS_BASETYPE, /* tp_flags */
3350 permutations_doc, /* tp_doc */
3351 (traverseproc)permutations_traverse, /* tp_traverse */
3352 0, /* tp_clear */
3353 0, /* tp_richcompare */
3354 0, /* tp_weaklistoffset */
3355 PyObject_SelfIter, /* tp_iter */
3356 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003357 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003358 0, /* tp_members */
3359 0, /* tp_getset */
3360 0, /* tp_base */
3361 0, /* tp_dict */
3362 0, /* tp_descr_get */
3363 0, /* tp_descr_set */
3364 0, /* tp_dictoffset */
3365 0, /* tp_init */
3366 0, /* tp_alloc */
3367 permutations_new, /* tp_new */
3368 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003369};
3370
Raymond Hettinger482ba772010-12-01 22:48:00 +00003371/* accumulate object ************************************************************/
3372
3373typedef struct {
3374 PyObject_HEAD
3375 PyObject *total;
3376 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003377 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003378} accumulateobject;
3379
3380static PyTypeObject accumulate_type;
3381
3382static PyObject *
3383accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3384{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003385 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003386 PyObject *iterable;
3387 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003388 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003389 accumulateobject *lz;
3390
Raymond Hettinger5d446132011-03-27 18:52:10 -07003391 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3392 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003393 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003394
3395 /* Get iterator. */
3396 it = PyObject_GetIter(iterable);
3397 if (it == NULL)
3398 return NULL;
3399
Raymond Hettinger482ba772010-12-01 22:48:00 +00003400 /* create accumulateobject structure */
3401 lz = (accumulateobject *)type->tp_alloc(type, 0);
3402 if (lz == NULL) {
3403 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003404 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003405 }
3406
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003407 if (binop != Py_None) {
3408 Py_XINCREF(binop);
3409 lz->binop = binop;
3410 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003411 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003412 lz->it = it;
3413 return (PyObject *)lz;
3414}
3415
3416static void
3417accumulate_dealloc(accumulateobject *lz)
3418{
3419 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003420 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003421 Py_XDECREF(lz->total);
3422 Py_XDECREF(lz->it);
3423 Py_TYPE(lz)->tp_free(lz);
3424}
3425
3426static int
3427accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3428{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003429 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003430 Py_VISIT(lz->it);
3431 Py_VISIT(lz->total);
3432 return 0;
3433}
3434
3435static PyObject *
3436accumulate_next(accumulateobject *lz)
3437{
3438 PyObject *val, *oldtotal, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003439
Raymond Hettinger482ba772010-12-01 22:48:00 +00003440 val = PyIter_Next(lz->it);
3441 if (val == NULL)
3442 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003443
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003444 if (lz->total == NULL) {
3445 Py_INCREF(val);
3446 lz->total = val;
3447 return lz->total;
3448 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003449
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003450 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003451 newtotal = PyNumber_Add(lz->total, val);
3452 else
3453 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003454 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003455 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003456 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003457
3458 oldtotal = lz->total;
3459 lz->total = newtotal;
3460 Py_DECREF(oldtotal);
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003461
Raymond Hettinger482ba772010-12-01 22:48:00 +00003462 Py_INCREF(newtotal);
3463 return newtotal;
3464}
3465
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003466static PyObject *
3467accumulate_reduce(accumulateobject *lz)
3468{
3469 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3470 lz->it, lz->binop?lz->binop:Py_None,
3471 lz->total?lz->total:Py_None);
3472 }
3473
3474static PyObject *
3475accumulate_setstate(accumulateobject *lz, PyObject *state)
3476{
3477 Py_CLEAR(lz->total);
3478 lz->total = state;
3479 Py_INCREF(lz->total);
3480 Py_RETURN_NONE;
3481}
3482
3483static PyMethodDef accumulate_methods[] = {
3484 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3485 reduce_doc},
3486 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3487 setstate_doc},
3488 {NULL, NULL} /* sentinel */
3489};
3490
Raymond Hettinger482ba772010-12-01 22:48:00 +00003491PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003492"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003493\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003494Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003495
3496static PyTypeObject accumulate_type = {
3497 PyVarObject_HEAD_INIT(NULL, 0)
3498 "itertools.accumulate", /* tp_name */
3499 sizeof(accumulateobject), /* tp_basicsize */
3500 0, /* tp_itemsize */
3501 /* methods */
3502 (destructor)accumulate_dealloc, /* tp_dealloc */
3503 0, /* tp_print */
3504 0, /* tp_getattr */
3505 0, /* tp_setattr */
3506 0, /* tp_reserved */
3507 0, /* tp_repr */
3508 0, /* tp_as_number */
3509 0, /* tp_as_sequence */
3510 0, /* tp_as_mapping */
3511 0, /* tp_hash */
3512 0, /* tp_call */
3513 0, /* tp_str */
3514 PyObject_GenericGetAttr, /* tp_getattro */
3515 0, /* tp_setattro */
3516 0, /* tp_as_buffer */
3517 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3518 Py_TPFLAGS_BASETYPE, /* tp_flags */
3519 accumulate_doc, /* tp_doc */
3520 (traverseproc)accumulate_traverse, /* tp_traverse */
3521 0, /* tp_clear */
3522 0, /* tp_richcompare */
3523 0, /* tp_weaklistoffset */
3524 PyObject_SelfIter, /* tp_iter */
3525 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003526 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003527 0, /* tp_members */
3528 0, /* tp_getset */
3529 0, /* tp_base */
3530 0, /* tp_dict */
3531 0, /* tp_descr_get */
3532 0, /* tp_descr_set */
3533 0, /* tp_dictoffset */
3534 0, /* tp_init */
3535 0, /* tp_alloc */
3536 accumulate_new, /* tp_new */
3537 PyObject_GC_Del, /* tp_free */
3538};
3539
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003540
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003541/* compress object ************************************************************/
3542
3543/* Equivalent to:
3544
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003545 def compress(data, selectors):
3546 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3547 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003548*/
3549
3550typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003551 PyObject_HEAD
3552 PyObject *data;
3553 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003554} compressobject;
3555
3556static PyTypeObject compress_type;
3557
3558static PyObject *
3559compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3560{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003561 PyObject *seq1, *seq2;
3562 PyObject *data=NULL, *selectors=NULL;
3563 compressobject *lz;
3564 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003565
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003566 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3567 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003568
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003569 data = PyObject_GetIter(seq1);
3570 if (data == NULL)
3571 goto fail;
3572 selectors = PyObject_GetIter(seq2);
3573 if (selectors == NULL)
3574 goto fail;
3575
3576 /* create compressobject structure */
3577 lz = (compressobject *)type->tp_alloc(type, 0);
3578 if (lz == NULL)
3579 goto fail;
3580 lz->data = data;
3581 lz->selectors = selectors;
3582 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003583
3584fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003585 Py_XDECREF(data);
3586 Py_XDECREF(selectors);
3587 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003588}
3589
3590static void
3591compress_dealloc(compressobject *lz)
3592{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003593 PyObject_GC_UnTrack(lz);
3594 Py_XDECREF(lz->data);
3595 Py_XDECREF(lz->selectors);
3596 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003597}
3598
3599static int
3600compress_traverse(compressobject *lz, visitproc visit, void *arg)
3601{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003602 Py_VISIT(lz->data);
3603 Py_VISIT(lz->selectors);
3604 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003605}
3606
3607static PyObject *
3608compress_next(compressobject *lz)
3609{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003610 PyObject *data = lz->data, *selectors = lz->selectors;
3611 PyObject *datum, *selector;
3612 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3613 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3614 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003615
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003616 while (1) {
3617 /* Steps: get datum, get selector, evaluate selector.
3618 Order is important (to match the pure python version
3619 in terms of which input gets a chance to raise an
3620 exception first).
3621 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003622
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003623 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003624 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003625 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003626
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003627 selector = selectornext(selectors);
3628 if (selector == NULL) {
3629 Py_DECREF(datum);
3630 return NULL;
3631 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003632
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003633 ok = PyObject_IsTrue(selector);
3634 Py_DECREF(selector);
3635 if (ok == 1)
3636 return datum;
3637 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003638 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003639 return NULL;
3640 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003641}
3642
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003643static PyObject *
3644compress_reduce(compressobject *lz)
3645{
3646 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3647 lz->data, lz->selectors);
3648 }
3649
3650static PyMethodDef compress_methods[] = {
3651 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3652 reduce_doc},
3653 {NULL, NULL} /* sentinel */
3654};
3655
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003656PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003657"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003658\n\
3659Return data elements corresponding to true selector elements.\n\
3660Forms a shorter iterator from selected data elements using the\n\
3661selectors to choose the data elements.");
3662
3663static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003664 PyVarObject_HEAD_INIT(NULL, 0)
3665 "itertools.compress", /* tp_name */
3666 sizeof(compressobject), /* tp_basicsize */
3667 0, /* tp_itemsize */
3668 /* methods */
3669 (destructor)compress_dealloc, /* tp_dealloc */
3670 0, /* tp_print */
3671 0, /* tp_getattr */
3672 0, /* tp_setattr */
3673 0, /* tp_reserved */
3674 0, /* tp_repr */
3675 0, /* tp_as_number */
3676 0, /* tp_as_sequence */
3677 0, /* tp_as_mapping */
3678 0, /* tp_hash */
3679 0, /* tp_call */
3680 0, /* tp_str */
3681 PyObject_GenericGetAttr, /* tp_getattro */
3682 0, /* tp_setattro */
3683 0, /* tp_as_buffer */
3684 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3685 Py_TPFLAGS_BASETYPE, /* tp_flags */
3686 compress_doc, /* tp_doc */
3687 (traverseproc)compress_traverse, /* tp_traverse */
3688 0, /* tp_clear */
3689 0, /* tp_richcompare */
3690 0, /* tp_weaklistoffset */
3691 PyObject_SelfIter, /* tp_iter */
3692 (iternextfunc)compress_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003693 compress_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003694 0, /* tp_members */
3695 0, /* tp_getset */
3696 0, /* tp_base */
3697 0, /* tp_dict */
3698 0, /* tp_descr_get */
3699 0, /* tp_descr_set */
3700 0, /* tp_dictoffset */
3701 0, /* tp_init */
3702 0, /* tp_alloc */
3703 compress_new, /* tp_new */
3704 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003705};
3706
3707
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003708/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003709
3710typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003711 PyObject_HEAD
3712 PyObject *func;
3713 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003714} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003715
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003716static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003717
3718static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003719filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003720{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003721 PyObject *func, *seq;
3722 PyObject *it;
3723 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003724
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003725 if (type == &filterfalse_type &&
3726 !_PyArg_NoKeywords("filterfalse()", kwds))
3727 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003728
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003729 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3730 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003731
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003732 /* Get iterator. */
3733 it = PyObject_GetIter(seq);
3734 if (it == NULL)
3735 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003736
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003737 /* create filterfalseobject structure */
3738 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3739 if (lz == NULL) {
3740 Py_DECREF(it);
3741 return NULL;
3742 }
3743 Py_INCREF(func);
3744 lz->func = func;
3745 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003746
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003747 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003748}
3749
3750static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003751filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003752{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003753 PyObject_GC_UnTrack(lz);
3754 Py_XDECREF(lz->func);
3755 Py_XDECREF(lz->it);
3756 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003757}
3758
3759static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003760filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003761{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003762 Py_VISIT(lz->it);
3763 Py_VISIT(lz->func);
3764 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003765}
3766
3767static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003768filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003769{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003770 PyObject *item;
3771 PyObject *it = lz->it;
3772 long ok;
3773 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003774
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003775 iternext = *Py_TYPE(it)->tp_iternext;
3776 for (;;) {
3777 item = iternext(it);
3778 if (item == NULL)
3779 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003780
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003781 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3782 ok = PyObject_IsTrue(item);
3783 } else {
3784 PyObject *good;
3785 good = PyObject_CallFunctionObjArgs(lz->func,
3786 item, NULL);
3787 if (good == NULL) {
3788 Py_DECREF(item);
3789 return NULL;
3790 }
3791 ok = PyObject_IsTrue(good);
3792 Py_DECREF(good);
3793 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003794 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003795 return item;
3796 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003797 if (ok < 0)
3798 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003799 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003800}
3801
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003802static PyObject *
3803filterfalse_reduce(filterfalseobject *lz)
3804{
3805 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3806 lz->func, lz->it);
3807 }
3808
3809static PyMethodDef filterfalse_methods[] = {
3810 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3811 reduce_doc},
3812 {NULL, NULL} /* sentinel */
3813};
3814
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003815PyDoc_STRVAR(filterfalse_doc,
3816"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003817\n\
3818Return those items of sequence for which function(item) is false.\n\
3819If function is None, return the items that are false.");
3820
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003821static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003822 PyVarObject_HEAD_INIT(NULL, 0)
3823 "itertools.filterfalse", /* tp_name */
3824 sizeof(filterfalseobject), /* tp_basicsize */
3825 0, /* tp_itemsize */
3826 /* methods */
3827 (destructor)filterfalse_dealloc, /* tp_dealloc */
3828 0, /* tp_print */
3829 0, /* tp_getattr */
3830 0, /* tp_setattr */
3831 0, /* tp_reserved */
3832 0, /* tp_repr */
3833 0, /* tp_as_number */
3834 0, /* tp_as_sequence */
3835 0, /* tp_as_mapping */
3836 0, /* tp_hash */
3837 0, /* tp_call */
3838 0, /* tp_str */
3839 PyObject_GenericGetAttr, /* tp_getattro */
3840 0, /* tp_setattro */
3841 0, /* tp_as_buffer */
3842 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3843 Py_TPFLAGS_BASETYPE, /* tp_flags */
3844 filterfalse_doc, /* tp_doc */
3845 (traverseproc)filterfalse_traverse, /* tp_traverse */
3846 0, /* tp_clear */
3847 0, /* tp_richcompare */
3848 0, /* tp_weaklistoffset */
3849 PyObject_SelfIter, /* tp_iter */
3850 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003851 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003852 0, /* tp_members */
3853 0, /* tp_getset */
3854 0, /* tp_base */
3855 0, /* tp_dict */
3856 0, /* tp_descr_get */
3857 0, /* tp_descr_set */
3858 0, /* tp_dictoffset */
3859 0, /* tp_init */
3860 0, /* tp_alloc */
3861 filterfalse_new, /* tp_new */
3862 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003863};
3864
3865
3866/* count object ************************************************************/
3867
3868typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003869 PyObject_HEAD
3870 Py_ssize_t cnt;
3871 PyObject *long_cnt;
3872 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003873} countobject;
3874
Raymond Hettinger30729212009-02-12 06:28:27 +00003875/* Counting logic and invariants:
3876
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003877fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003878
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003879 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3880 Advances with: cnt += 1
3881 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003882
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003883slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
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 != NULL);
3886 All counting is done with python objects (no overflows or underflows).
3887 Advances with: long_cnt += long_step
3888 Step may be zero -- effectively a slow version of repeat(cnt).
3889 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003890*/
3891
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003892static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003893
3894static PyObject *
3895count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3896{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003897 countobject *lz;
3898 int slow_mode = 0;
3899 Py_ssize_t cnt = 0;
3900 PyObject *long_cnt = NULL;
3901 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003902 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003903 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003904
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003905 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3906 kwlist, &long_cnt, &long_step))
3907 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003908
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003909 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3910 (long_step != NULL && !PyNumber_Check(long_step))) {
3911 PyErr_SetString(PyExc_TypeError, "a number is required");
3912 return NULL;
3913 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003914
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003915 if (long_cnt != NULL) {
3916 cnt = PyLong_AsSsize_t(long_cnt);
3917 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3918 PyErr_Clear();
3919 slow_mode = 1;
3920 }
3921 Py_INCREF(long_cnt);
3922 } else {
3923 cnt = 0;
3924 long_cnt = PyLong_FromLong(0);
3925 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003926
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003927 /* If not specified, step defaults to 1 */
3928 if (long_step == NULL) {
3929 long_step = PyLong_FromLong(1);
3930 if (long_step == NULL) {
3931 Py_DECREF(long_cnt);
3932 return NULL;
3933 }
3934 } else
3935 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003936
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003937 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003938
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003939 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003940 step = PyLong_AsLong(long_step);
3941 if (step != 1) {
3942 slow_mode = 1;
3943 if (step == -1 && PyErr_Occurred())
3944 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003945 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003946
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003947 if (slow_mode)
3948 cnt = PY_SSIZE_T_MAX;
3949 else
3950 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003951
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003952 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3953 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3954 assert(slow_mode ||
3955 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003956
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003957 /* create countobject structure */
3958 lz = (countobject *)type->tp_alloc(type, 0);
3959 if (lz == NULL) {
3960 Py_XDECREF(long_cnt);
3961 return NULL;
3962 }
3963 lz->cnt = cnt;
3964 lz->long_cnt = long_cnt;
3965 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003966
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003967 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003968}
3969
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003970static void
3971count_dealloc(countobject *lz)
3972{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003973 PyObject_GC_UnTrack(lz);
3974 Py_XDECREF(lz->long_cnt);
3975 Py_XDECREF(lz->long_step);
3976 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003977}
3978
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003979static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003980count_traverse(countobject *lz, visitproc visit, void *arg)
3981{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003982 Py_VISIT(lz->long_cnt);
3983 Py_VISIT(lz->long_step);
3984 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003985}
3986
3987static PyObject *
3988count_nextlong(countobject *lz)
3989{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003990 PyObject *long_cnt;
3991 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003992
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003993 long_cnt = lz->long_cnt;
3994 if (long_cnt == NULL) {
3995 /* Switch to slow_mode */
3996 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3997 if (long_cnt == NULL)
3998 return NULL;
3999 }
4000 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004001
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004002 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4003 if (stepped_up == NULL)
4004 return NULL;
4005 lz->long_cnt = stepped_up;
4006 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004007}
4008
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004009static PyObject *
4010count_next(countobject *lz)
4011{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004012 if (lz->cnt == PY_SSIZE_T_MAX)
4013 return count_nextlong(lz);
4014 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004015}
4016
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004017static PyObject *
4018count_repr(countobject *lz)
4019{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004020 if (lz->cnt != PY_SSIZE_T_MAX)
4021 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004022
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004023 if (PyLong_Check(lz->long_step)) {
4024 long step = PyLong_AsLong(lz->long_step);
4025 if (step == -1 && PyErr_Occurred()) {
4026 PyErr_Clear();
4027 }
4028 if (step == 1) {
4029 /* Don't display step when it is an integer equal to 1 */
4030 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
4031 }
4032 }
4033 return PyUnicode_FromFormat("count(%R, %R)",
4034 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004035}
4036
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004037static PyObject *
4038count_reduce(countobject *lz)
4039{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004040 if (lz->cnt == PY_SSIZE_T_MAX)
4041 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4042 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004043}
4044
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004045static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004046 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004047 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004048 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004049};
4050
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004051PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00004052 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004053\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004054Return a count object whose .__next__() method returns consecutive values.\n\
4055Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00004056 def count(firstval=0, step=1):\n\
Eli Benderskyc554f722013-09-03 06:37:19 -07004057 x = firstval\n\
4058 while 1:\n\
4059 yield x\n\
4060 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004061
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004062static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004063 PyVarObject_HEAD_INIT(NULL, 0)
4064 "itertools.count", /* tp_name */
4065 sizeof(countobject), /* tp_basicsize */
4066 0, /* tp_itemsize */
4067 /* methods */
4068 (destructor)count_dealloc, /* tp_dealloc */
4069 0, /* tp_print */
4070 0, /* tp_getattr */
4071 0, /* tp_setattr */
4072 0, /* tp_reserved */
4073 (reprfunc)count_repr, /* tp_repr */
4074 0, /* tp_as_number */
4075 0, /* tp_as_sequence */
4076 0, /* tp_as_mapping */
4077 0, /* tp_hash */
4078 0, /* tp_call */
4079 0, /* tp_str */
4080 PyObject_GenericGetAttr, /* tp_getattro */
4081 0, /* tp_setattro */
4082 0, /* tp_as_buffer */
4083 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4084 Py_TPFLAGS_BASETYPE, /* tp_flags */
4085 count_doc, /* tp_doc */
4086 (traverseproc)count_traverse, /* tp_traverse */
4087 0, /* tp_clear */
4088 0, /* tp_richcompare */
4089 0, /* tp_weaklistoffset */
4090 PyObject_SelfIter, /* tp_iter */
4091 (iternextfunc)count_next, /* tp_iternext */
4092 count_methods, /* tp_methods */
4093 0, /* tp_members */
4094 0, /* tp_getset */
4095 0, /* tp_base */
4096 0, /* tp_dict */
4097 0, /* tp_descr_get */
4098 0, /* tp_descr_set */
4099 0, /* tp_dictoffset */
4100 0, /* tp_init */
4101 0, /* tp_alloc */
4102 count_new, /* tp_new */
4103 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004104};
4105
4106
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004107/* repeat object ************************************************************/
4108
4109typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004110 PyObject_HEAD
4111 PyObject *element;
4112 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004113} repeatobject;
4114
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004115static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004116
4117static PyObject *
4118repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4119{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004120 repeatobject *ro;
4121 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004122 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004123 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004124
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004125 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4126 &element, &cnt))
4127 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004128
Raymond Hettinger97d35552014-06-24 21:36:58 -07004129 if (kwds != NULL)
4130 n_kwds = PyDict_Size(kwds);
4131 /* Does user supply times argument? */
4132 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004133 cnt = 0;
4134
4135 ro = (repeatobject *)type->tp_alloc(type, 0);
4136 if (ro == NULL)
4137 return NULL;
4138 Py_INCREF(element);
4139 ro->element = element;
4140 ro->cnt = cnt;
4141 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004142}
4143
4144static void
4145repeat_dealloc(repeatobject *ro)
4146{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004147 PyObject_GC_UnTrack(ro);
4148 Py_XDECREF(ro->element);
4149 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004150}
4151
4152static int
4153repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4154{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004155 Py_VISIT(ro->element);
4156 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004157}
4158
4159static PyObject *
4160repeat_next(repeatobject *ro)
4161{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004162 if (ro->cnt == 0)
4163 return NULL;
4164 if (ro->cnt > 0)
4165 ro->cnt--;
4166 Py_INCREF(ro->element);
4167 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004168}
4169
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004170static PyObject *
4171repeat_repr(repeatobject *ro)
4172{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004173 if (ro->cnt == -1)
4174 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4175 else
4176 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4177}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004178
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004179static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004180repeat_len(repeatobject *ro)
4181{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004182 if (ro->cnt == -1) {
4183 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4184 return NULL;
4185 }
4186 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004187}
4188
Armin Rigof5b3e362006-02-11 21:32:43 +00004189PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004190
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004191static PyObject *
4192repeat_reduce(repeatobject *ro)
4193{
4194 /* unpickle this so that a new repeat iterator is constructed with an
4195 * object, then call __setstate__ on it to set cnt
4196 */
4197 if (ro->cnt >= 0)
4198 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4199 else
4200 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4201}
4202
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004203static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004204 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004205 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004206 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004207};
4208
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004209PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004210"repeat(object [,times]) -> create an iterator which returns the object\n\
4211for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004212endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004213
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004214static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004215 PyVarObject_HEAD_INIT(NULL, 0)
4216 "itertools.repeat", /* tp_name */
4217 sizeof(repeatobject), /* tp_basicsize */
4218 0, /* tp_itemsize */
4219 /* methods */
4220 (destructor)repeat_dealloc, /* tp_dealloc */
4221 0, /* tp_print */
4222 0, /* tp_getattr */
4223 0, /* tp_setattr */
4224 0, /* tp_reserved */
4225 (reprfunc)repeat_repr, /* tp_repr */
4226 0, /* tp_as_number */
4227 0, /* tp_as_sequence */
4228 0, /* tp_as_mapping */
4229 0, /* tp_hash */
4230 0, /* tp_call */
4231 0, /* tp_str */
4232 PyObject_GenericGetAttr, /* tp_getattro */
4233 0, /* tp_setattro */
4234 0, /* tp_as_buffer */
4235 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4236 Py_TPFLAGS_BASETYPE, /* tp_flags */
4237 repeat_doc, /* tp_doc */
4238 (traverseproc)repeat_traverse, /* tp_traverse */
4239 0, /* tp_clear */
4240 0, /* tp_richcompare */
4241 0, /* tp_weaklistoffset */
4242 PyObject_SelfIter, /* tp_iter */
4243 (iternextfunc)repeat_next, /* tp_iternext */
4244 repeat_methods, /* tp_methods */
4245 0, /* tp_members */
4246 0, /* tp_getset */
4247 0, /* tp_base */
4248 0, /* tp_dict */
4249 0, /* tp_descr_get */
4250 0, /* tp_descr_set */
4251 0, /* tp_dictoffset */
4252 0, /* tp_init */
4253 0, /* tp_alloc */
4254 repeat_new, /* tp_new */
4255 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004256};
4257
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004258/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004259
4260#include "Python.h"
4261
4262typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004263 PyObject_HEAD
4264 Py_ssize_t tuplesize;
4265 Py_ssize_t numactive;
4266 PyObject *ittuple; /* tuple of iterators */
4267 PyObject *result;
4268 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004269} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004270
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004271static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004272
4273static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004274zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004275{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004276 ziplongestobject *lz;
4277 Py_ssize_t i;
4278 PyObject *ittuple; /* tuple of iterators */
4279 PyObject *result;
4280 PyObject *fillvalue = Py_None;
4281 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004282
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004283 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
4284 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4285 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
4286 PyErr_SetString(PyExc_TypeError,
4287 "zip_longest() got an unexpected keyword argument");
4288 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004289 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004290 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004291
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004292 /* args must be a tuple */
4293 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00004294
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004295 /* obtain iterators */
4296 ittuple = PyTuple_New(tuplesize);
4297 if (ittuple == NULL)
4298 return NULL;
4299 for (i=0; i < tuplesize; ++i) {
4300 PyObject *item = PyTuple_GET_ITEM(args, i);
4301 PyObject *it = PyObject_GetIter(item);
4302 if (it == NULL) {
4303 if (PyErr_ExceptionMatches(PyExc_TypeError))
4304 PyErr_Format(PyExc_TypeError,
4305 "zip_longest argument #%zd must support iteration",
4306 i+1);
4307 Py_DECREF(ittuple);
4308 return NULL;
4309 }
4310 PyTuple_SET_ITEM(ittuple, i, it);
4311 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004312
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004313 /* create a result holder */
4314 result = PyTuple_New(tuplesize);
4315 if (result == NULL) {
4316 Py_DECREF(ittuple);
4317 return NULL;
4318 }
4319 for (i=0 ; i < tuplesize ; i++) {
4320 Py_INCREF(Py_None);
4321 PyTuple_SET_ITEM(result, i, Py_None);
4322 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004323
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004324 /* create ziplongestobject structure */
4325 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4326 if (lz == NULL) {
4327 Py_DECREF(ittuple);
4328 Py_DECREF(result);
4329 return NULL;
4330 }
4331 lz->ittuple = ittuple;
4332 lz->tuplesize = tuplesize;
4333 lz->numactive = tuplesize;
4334 lz->result = result;
4335 Py_INCREF(fillvalue);
4336 lz->fillvalue = fillvalue;
4337 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004338}
4339
4340static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004341zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004342{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004343 PyObject_GC_UnTrack(lz);
4344 Py_XDECREF(lz->ittuple);
4345 Py_XDECREF(lz->result);
4346 Py_XDECREF(lz->fillvalue);
4347 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004348}
4349
4350static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004351zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004352{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004353 Py_VISIT(lz->ittuple);
4354 Py_VISIT(lz->result);
4355 Py_VISIT(lz->fillvalue);
4356 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004357}
4358
4359static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004360zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004361{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004362 Py_ssize_t i;
4363 Py_ssize_t tuplesize = lz->tuplesize;
4364 PyObject *result = lz->result;
4365 PyObject *it;
4366 PyObject *item;
4367 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004368
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004369 if (tuplesize == 0)
4370 return NULL;
4371 if (lz->numactive == 0)
4372 return NULL;
4373 if (Py_REFCNT(result) == 1) {
4374 Py_INCREF(result);
4375 for (i=0 ; i < tuplesize ; i++) {
4376 it = PyTuple_GET_ITEM(lz->ittuple, i);
4377 if (it == NULL) {
4378 Py_INCREF(lz->fillvalue);
4379 item = lz->fillvalue;
4380 } else {
4381 item = PyIter_Next(it);
4382 if (item == NULL) {
4383 lz->numactive -= 1;
4384 if (lz->numactive == 0 || PyErr_Occurred()) {
4385 lz->numactive = 0;
4386 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004387 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004388 } else {
4389 Py_INCREF(lz->fillvalue);
4390 item = lz->fillvalue;
4391 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4392 Py_DECREF(it);
4393 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004394 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004395 }
4396 olditem = PyTuple_GET_ITEM(result, i);
4397 PyTuple_SET_ITEM(result, i, item);
4398 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004399 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004400 } else {
4401 result = PyTuple_New(tuplesize);
4402 if (result == NULL)
4403 return NULL;
4404 for (i=0 ; i < tuplesize ; i++) {
4405 it = PyTuple_GET_ITEM(lz->ittuple, i);
4406 if (it == NULL) {
4407 Py_INCREF(lz->fillvalue);
4408 item = lz->fillvalue;
4409 } else {
4410 item = PyIter_Next(it);
4411 if (item == NULL) {
4412 lz->numactive -= 1;
4413 if (lz->numactive == 0 || PyErr_Occurred()) {
4414 lz->numactive = 0;
4415 Py_DECREF(result);
4416 return NULL;
4417 } else {
4418 Py_INCREF(lz->fillvalue);
4419 item = lz->fillvalue;
4420 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4421 Py_DECREF(it);
4422 }
4423 }
4424 }
4425 PyTuple_SET_ITEM(result, i, item);
4426 }
4427 }
4428 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004429}
4430
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004431static PyObject *
4432zip_longest_reduce(ziplongestobject *lz)
4433{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004434
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004435 /* Create a new tuple with empty sequences where appropriate to pickle.
4436 * Then use setstate to set the fillvalue
4437 */
4438 int i;
4439 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
4440 if (args == NULL)
4441 return NULL;
4442 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4443 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4444 if (elem == NULL) {
4445 elem = PyTuple_New(0);
4446 if (elem == NULL) {
4447 Py_DECREF(args);
4448 return NULL;
4449 }
4450 } else
4451 Py_INCREF(elem);
4452 PyTuple_SET_ITEM(args, i, elem);
4453 }
4454 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4455}
4456
4457static PyObject *
4458zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4459{
4460 Py_CLEAR(lz->fillvalue);
4461 lz->fillvalue = state;
4462 Py_INCREF(lz->fillvalue);
4463 Py_RETURN_NONE;
4464}
4465
4466static PyMethodDef zip_longest_methods[] = {
4467 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4468 reduce_doc},
4469 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4470 setstate_doc},
4471 {NULL, NULL} /* sentinel */
4472};
4473
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004474PyDoc_STRVAR(zip_longest_doc,
4475"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004476\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004477Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004478the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004479method continues until the longest iterable in the argument sequence\n\
4480is exhausted and then it raises StopIteration. When the shorter iterables\n\
4481are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4482defaults to None or can be specified by a keyword argument.\n\
4483");
4484
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004485static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004486 PyVarObject_HEAD_INIT(NULL, 0)
4487 "itertools.zip_longest", /* tp_name */
4488 sizeof(ziplongestobject), /* tp_basicsize */
4489 0, /* tp_itemsize */
4490 /* methods */
4491 (destructor)zip_longest_dealloc, /* tp_dealloc */
4492 0, /* tp_print */
4493 0, /* tp_getattr */
4494 0, /* tp_setattr */
4495 0, /* tp_reserved */
4496 0, /* tp_repr */
4497 0, /* tp_as_number */
4498 0, /* tp_as_sequence */
4499 0, /* tp_as_mapping */
4500 0, /* tp_hash */
4501 0, /* tp_call */
4502 0, /* tp_str */
4503 PyObject_GenericGetAttr, /* tp_getattro */
4504 0, /* tp_setattro */
4505 0, /* tp_as_buffer */
4506 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4507 Py_TPFLAGS_BASETYPE, /* tp_flags */
4508 zip_longest_doc, /* tp_doc */
4509 (traverseproc)zip_longest_traverse, /* tp_traverse */
4510 0, /* tp_clear */
4511 0, /* tp_richcompare */
4512 0, /* tp_weaklistoffset */
4513 PyObject_SelfIter, /* tp_iter */
4514 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004515 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004516 0, /* tp_members */
4517 0, /* tp_getset */
4518 0, /* tp_base */
4519 0, /* tp_dict */
4520 0, /* tp_descr_get */
4521 0, /* tp_descr_set */
4522 0, /* tp_dictoffset */
4523 0, /* tp_init */
4524 0, /* tp_alloc */
4525 zip_longest_new, /* tp_new */
4526 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004527};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004528
4529/* module level code ********************************************************/
4530
4531PyDoc_STRVAR(module_doc,
4532"Functional tools for creating and using iterators.\n\
4533\n\
4534Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004535count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004536cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004537repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004538\n\
4539Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004540accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004541chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger8df58f72013-09-09 01:29:40 -05004542chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004543compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4544dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4545groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004546filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004547islice(seq, [start,] stop [, step]) --> elements from\n\
4548 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004549starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004550tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004551takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004552zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004553\n\
4554Combinatoric generators:\n\
4555product(p, q, ... [repeat=1]) --> cartesian product\n\
4556permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004557combinations(p, r)\n\
4558combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004559");
4560
4561
Raymond Hettingerad983e72003-11-12 14:32:26 +00004562static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004563 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4564 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004565};
4566
Martin v. Löwis1a214512008-06-11 05:26:20 +00004567
4568static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004569 PyModuleDef_HEAD_INIT,
4570 "itertools",
4571 module_doc,
4572 -1,
4573 module_methods,
4574 NULL,
4575 NULL,
4576 NULL,
4577 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004578};
4579
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004580PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004581PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004582{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004583 int i;
4584 PyObject *m;
4585 char *name;
4586 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004587 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004588 &combinations_type,
4589 &cwr_type,
4590 &cycle_type,
4591 &dropwhile_type,
4592 &takewhile_type,
4593 &islice_type,
4594 &starmap_type,
4595 &chain_type,
4596 &compress_type,
4597 &filterfalse_type,
4598 &count_type,
4599 &ziplongest_type,
4600 &permutations_type,
4601 &product_type,
4602 &repeat_type,
4603 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004604 &_grouper_type,
4605 &tee_type,
4606 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004607 NULL
4608 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004609
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004610 Py_TYPE(&teedataobject_type) = &PyType_Type;
4611 m = PyModule_Create(&itertoolsmodule);
4612 if (m == NULL)
4613 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004614
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004615 for (i=0 ; typelist[i] != NULL ; i++) {
4616 if (PyType_Ready(typelist[i]) < 0)
4617 return NULL;
4618 name = strchr(typelist[i]->tp_name, '.');
4619 assert (name != NULL);
4620 Py_INCREF(typelist[i]);
4621 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4622 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004623
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004624 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004625}