blob: 9c21fb86527753137c7d63bf791aebed774f957f [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
2#include "Python.h"
Fred Drake08ebfec2004-10-17 19:36:57 +00003#include "structmember.h"
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00005/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00006 by Raymond D. Hettinger <python@rcn.com>
Raymond Hettinger96ef8112003-02-01 00:10:11 +00007*/
8
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00009
10/* groupby object ***********************************************************/
11
12typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000013 PyObject_HEAD
14 PyObject *it;
15 PyObject *keyfunc;
16 PyObject *tgtkey;
17 PyObject *currkey;
18 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000019} groupbyobject;
20
21static PyTypeObject groupby_type;
22static PyObject *_grouper_create(groupbyobject *, PyObject *);
23
24static PyObject *
25groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
26{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000027 static char *kwargs[] = {"iterable", "key", NULL};
28 groupbyobject *gbo;
29 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000030
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000031 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
32 &it, &keyfunc))
33 return NULL;
34
35 gbo = (groupbyobject *)type->tp_alloc(type, 0);
36 if (gbo == NULL)
37 return NULL;
38 gbo->tgtkey = NULL;
39 gbo->currkey = NULL;
40 gbo->currvalue = NULL;
41 gbo->keyfunc = keyfunc;
42 Py_INCREF(keyfunc);
43 gbo->it = PyObject_GetIter(it);
44 if (gbo->it == NULL) {
45 Py_DECREF(gbo);
46 return NULL;
47 }
48 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000049}
50
51static void
52groupby_dealloc(groupbyobject *gbo)
53{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000054 PyObject_GC_UnTrack(gbo);
55 Py_XDECREF(gbo->it);
56 Py_XDECREF(gbo->keyfunc);
57 Py_XDECREF(gbo->tgtkey);
58 Py_XDECREF(gbo->currkey);
59 Py_XDECREF(gbo->currvalue);
60 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000061}
62
63static int
64groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
65{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000066 Py_VISIT(gbo->it);
67 Py_VISIT(gbo->keyfunc);
68 Py_VISIT(gbo->tgtkey);
69 Py_VISIT(gbo->currkey);
70 Py_VISIT(gbo->currvalue);
71 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000072}
73
74static PyObject *
75groupby_next(groupbyobject *gbo)
76{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000077 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000078
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000079 /* skip to next iteration group */
80 for (;;) {
81 if (gbo->currkey == NULL)
82 /* pass */;
83 else if (gbo->tgtkey == NULL)
84 break;
85 else {
86 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000087
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000088 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
89 gbo->currkey, Py_EQ);
90 if (rcmp == -1)
91 return NULL;
92 else if (rcmp == 0)
93 break;
94 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000095
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000096 newvalue = PyIter_Next(gbo->it);
97 if (newvalue == NULL)
98 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000099
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000100 if (gbo->keyfunc == Py_None) {
101 newkey = newvalue;
102 Py_INCREF(newvalue);
103 } else {
104 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
105 newvalue, NULL);
106 if (newkey == NULL) {
107 Py_DECREF(newvalue);
108 return NULL;
109 }
110 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000111
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000112 tmp = gbo->currkey;
113 gbo->currkey = newkey;
114 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000115
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000116 tmp = gbo->currvalue;
117 gbo->currvalue = newvalue;
118 Py_XDECREF(tmp);
119 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000120
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000121 Py_INCREF(gbo->currkey);
122 tmp = gbo->tgtkey;
123 gbo->tgtkey = gbo->currkey;
124 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000125
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000126 grouper = _grouper_create(gbo, gbo->tgtkey);
127 if (grouper == NULL)
128 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000129
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000130 r = PyTuple_Pack(2, gbo->currkey, grouper);
131 Py_DECREF(grouper);
132 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000133}
134
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000135static PyObject *
136groupby_reduce(groupbyobject *lz)
137{
138 /* reduce as a 'new' call with an optional 'setstate' if groupby
139 * has started
140 */
141 PyObject *value;
142 if (lz->tgtkey && lz->currkey && lz->currvalue)
143 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
144 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
145 else
146 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
147 lz->it, lz->keyfunc);
148
149 return value;
150}
151
152PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
153
154static PyObject *
155groupby_setstate(groupbyobject *lz, PyObject *state)
156{
157 PyObject *currkey, *currvalue, *tgtkey;
158 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey))
159 return NULL;
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200160 Py_INCREF(currkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300161 Py_XSETREF(lz->currkey, currkey);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200162 Py_INCREF(currvalue);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300163 Py_XSETREF(lz->currvalue, currvalue);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200164 Py_INCREF(tgtkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300165 Py_XSETREF(lz->tgtkey, tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000166 Py_RETURN_NONE;
167}
168
169PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
170
171static PyMethodDef groupby_methods[] = {
172 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
173 reduce_doc},
174 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
175 setstate_doc},
176 {NULL, NULL} /* sentinel */
177};
178
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000179PyDoc_STRVAR(groupby_doc,
180"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
181(key, sub-iterator) grouped by each value of key(value).\n");
182
183static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000184 PyVarObject_HEAD_INIT(NULL, 0)
185 "itertools.groupby", /* tp_name */
186 sizeof(groupbyobject), /* tp_basicsize */
187 0, /* tp_itemsize */
188 /* methods */
189 (destructor)groupby_dealloc, /* tp_dealloc */
190 0, /* tp_print */
191 0, /* tp_getattr */
192 0, /* tp_setattr */
193 0, /* tp_reserved */
194 0, /* tp_repr */
195 0, /* tp_as_number */
196 0, /* tp_as_sequence */
197 0, /* tp_as_mapping */
198 0, /* tp_hash */
199 0, /* tp_call */
200 0, /* tp_str */
201 PyObject_GenericGetAttr, /* tp_getattro */
202 0, /* tp_setattro */
203 0, /* tp_as_buffer */
204 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
205 Py_TPFLAGS_BASETYPE, /* tp_flags */
206 groupby_doc, /* tp_doc */
207 (traverseproc)groupby_traverse, /* tp_traverse */
208 0, /* tp_clear */
209 0, /* tp_richcompare */
210 0, /* tp_weaklistoffset */
211 PyObject_SelfIter, /* tp_iter */
212 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000213 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000214 0, /* tp_members */
215 0, /* tp_getset */
216 0, /* tp_base */
217 0, /* tp_dict */
218 0, /* tp_descr_get */
219 0, /* tp_descr_set */
220 0, /* tp_dictoffset */
221 0, /* tp_init */
222 0, /* tp_alloc */
223 groupby_new, /* tp_new */
224 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000225};
226
227
228/* _grouper object (internal) ************************************************/
229
230typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000231 PyObject_HEAD
232 PyObject *parent;
233 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000234} _grouperobject;
235
236static PyTypeObject _grouper_type;
237
238static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000239_grouper_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
240{
241 PyObject *parent, *tgtkey;
242
243 if (!PyArg_ParseTuple(args, "O!O", &groupby_type, &parent, &tgtkey))
244 return NULL;
245
246 return _grouper_create((groupbyobject*) parent, tgtkey);
247}
248
249static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000250_grouper_create(groupbyobject *parent, PyObject *tgtkey)
251{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000252 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000253
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000254 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
255 if (igo == NULL)
256 return NULL;
257 igo->parent = (PyObject *)parent;
258 Py_INCREF(parent);
259 igo->tgtkey = tgtkey;
260 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000261
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000262 PyObject_GC_Track(igo);
263 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000264}
265
266static void
267_grouper_dealloc(_grouperobject *igo)
268{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000269 PyObject_GC_UnTrack(igo);
270 Py_DECREF(igo->parent);
271 Py_DECREF(igo->tgtkey);
272 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000273}
274
275static int
276_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
277{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000278 Py_VISIT(igo->parent);
279 Py_VISIT(igo->tgtkey);
280 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000281}
282
283static PyObject *
284_grouper_next(_grouperobject *igo)
285{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000286 groupbyobject *gbo = (groupbyobject *)igo->parent;
287 PyObject *newvalue, *newkey, *r;
288 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000289
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000290 if (gbo->currvalue == NULL) {
291 newvalue = PyIter_Next(gbo->it);
292 if (newvalue == NULL)
293 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000294
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000295 if (gbo->keyfunc == Py_None) {
296 newkey = newvalue;
297 Py_INCREF(newvalue);
298 } else {
299 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
300 newvalue, NULL);
301 if (newkey == NULL) {
302 Py_DECREF(newvalue);
303 return NULL;
304 }
305 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000306
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000307 assert(gbo->currkey == NULL);
308 gbo->currkey = newkey;
309 gbo->currvalue = newvalue;
310 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000311
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000312 assert(gbo->currkey != NULL);
313 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
314 if (rcmp <= 0)
315 /* got any error or current group is end */
316 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000317
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000318 r = gbo->currvalue;
319 gbo->currvalue = NULL;
320 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000321
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000322 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000323}
324
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000325static PyObject *
326_grouper_reduce(_grouperobject *lz)
327{
328 return Py_BuildValue("O(OO)", Py_TYPE(lz),
329 lz->parent, lz->tgtkey);
330}
331
332static PyMethodDef _grouper_methods[] = {
333 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
334 reduce_doc},
335 {NULL, NULL} /* sentinel */
336};
337
338
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000339static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000340 PyVarObject_HEAD_INIT(NULL, 0)
341 "itertools._grouper", /* tp_name */
342 sizeof(_grouperobject), /* tp_basicsize */
343 0, /* tp_itemsize */
344 /* methods */
345 (destructor)_grouper_dealloc, /* tp_dealloc */
346 0, /* tp_print */
347 0, /* tp_getattr */
348 0, /* tp_setattr */
349 0, /* tp_reserved */
350 0, /* tp_repr */
351 0, /* tp_as_number */
352 0, /* tp_as_sequence */
353 0, /* tp_as_mapping */
354 0, /* tp_hash */
355 0, /* tp_call */
356 0, /* tp_str */
357 PyObject_GenericGetAttr, /* tp_getattro */
358 0, /* tp_setattro */
359 0, /* tp_as_buffer */
360 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
361 0, /* tp_doc */
362 (traverseproc)_grouper_traverse,/* tp_traverse */
363 0, /* tp_clear */
364 0, /* tp_richcompare */
365 0, /* tp_weaklistoffset */
366 PyObject_SelfIter, /* tp_iter */
367 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000368 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000369 0, /* tp_members */
370 0, /* tp_getset */
371 0, /* tp_base */
372 0, /* tp_dict */
373 0, /* tp_descr_get */
374 0, /* tp_descr_set */
375 0, /* tp_dictoffset */
376 0, /* tp_init */
377 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000378 _grouper_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000379 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000380};
381
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000382
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000383
Raymond Hettingerad983e72003-11-12 14:32:26 +0000384/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000385
Raymond Hettingerad983e72003-11-12 14:32:26 +0000386/* The teedataobject pre-allocates space for LINKCELLS number of objects.
387 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000388 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000389 the other structure members including PyHEAD overhead. The larger the
390 value, the less memory overhead per object and the less time spent
391 allocating/deallocating new links. The smaller the number, the less
392 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000393*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000394#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000395
396typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000397 PyObject_HEAD
398 PyObject *it;
Antoine Pitroub4a46cb2013-09-20 22:19:22 +0200399 int numread; /* 0 <= numread <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000400 PyObject *nextlink;
401 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000402} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000403
404typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000405 PyObject_HEAD
406 teedataobject *dataobj;
Antoine Pitroub4a46cb2013-09-20 22:19:22 +0200407 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000408 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000409} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000410
Raymond Hettingerad983e72003-11-12 14:32:26 +0000411static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000412
413static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000414teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000416 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000418 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
419 if (tdo == NULL)
420 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000421
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000422 tdo->numread = 0;
423 tdo->nextlink = NULL;
424 Py_INCREF(it);
425 tdo->it = it;
426 PyObject_GC_Track(tdo);
427 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000428}
429
430static PyObject *
431teedataobject_jumplink(teedataobject *tdo)
432{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000433 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000434 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000435 Py_XINCREF(tdo->nextlink);
436 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000437}
438
439static PyObject *
440teedataobject_getitem(teedataobject *tdo, int i)
441{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000442 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000443
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000444 assert(i < LINKCELLS);
445 if (i < tdo->numread)
446 value = tdo->values[i];
447 else {
448 /* this is the lead iterator, so fetch more data */
449 assert(i == tdo->numread);
450 value = PyIter_Next(tdo->it);
451 if (value == NULL)
452 return NULL;
453 tdo->numread++;
454 tdo->values[i] = value;
455 }
456 Py_INCREF(value);
457 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000458}
459
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000460static int
461teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
462{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000463 int i;
464 Py_VISIT(tdo->it);
465 for (i = 0; i < tdo->numread; i++)
466 Py_VISIT(tdo->values[i]);
467 Py_VISIT(tdo->nextlink);
468 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000469}
470
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200471static void
472teedataobject_safe_decref(PyObject *obj)
473{
474 while (obj && Py_TYPE(obj) == &teedataobject_type &&
475 Py_REFCNT(obj) == 1) {
476 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
477 ((teedataobject *)obj)->nextlink = NULL;
478 Py_DECREF(obj);
479 obj = nextlink;
480 }
481 Py_XDECREF(obj);
482}
483
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000484static int
485teedataobject_clear(teedataobject *tdo)
486{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000487 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200488 PyObject *tmp;
489
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000490 Py_CLEAR(tdo->it);
491 for (i=0 ; i<tdo->numread ; i++)
492 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200493 tmp = tdo->nextlink;
494 tdo->nextlink = NULL;
495 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000496 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000497}
498
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000499static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000500teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000501{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000502 PyObject_GC_UnTrack(tdo);
503 teedataobject_clear(tdo);
504 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000505}
506
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000507static PyObject *
508teedataobject_reduce(teedataobject *tdo)
509{
510 int i;
511 /* create a temporary list of already iterated values */
512 PyObject *values = PyList_New(tdo->numread);
513 if (!values)
514 return NULL;
515 for (i=0 ; i<tdo->numread ; i++) {
516 Py_INCREF(tdo->values[i]);
517 PyList_SET_ITEM(values, i, tdo->values[i]);
518 }
519 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
520 values,
521 tdo->nextlink ? tdo->nextlink : Py_None);
522}
523
524static PyTypeObject teedataobject_type;
525
526static PyObject *
527teedataobject_new(PyTypeObject *type, PyObject *args, PyObject *kw)
528{
529 teedataobject *tdo;
530 PyObject *it, *values, *next;
531 Py_ssize_t i, len;
532
533 assert(type == &teedataobject_type);
534 if (!PyArg_ParseTuple(args, "OO!O", &it, &PyList_Type, &values, &next))
535 return NULL;
536
537 tdo = (teedataobject *)teedataobject_newinternal(it);
538 if (!tdo)
539 return NULL;
540
541 len = PyList_GET_SIZE(values);
542 if (len > LINKCELLS)
543 goto err;
544 for (i=0; i<len; i++) {
545 tdo->values[i] = PyList_GET_ITEM(values, i);
546 Py_INCREF(tdo->values[i]);
547 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200548 /* len <= LINKCELLS < INT_MAX */
549 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000550
551 if (len == LINKCELLS) {
552 if (next != Py_None) {
553 if (Py_TYPE(next) != &teedataobject_type)
554 goto err;
555 assert(tdo->nextlink == NULL);
556 Py_INCREF(next);
557 tdo->nextlink = next;
558 }
559 } else {
560 if (next != Py_None)
561 goto err; /* shouldn't have a next if we are not full */
562 }
563 return (PyObject*)tdo;
564
565err:
566 Py_XDECREF(tdo);
567 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
568 return NULL;
569}
570
571static PyMethodDef teedataobject_methods[] = {
572 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
573 reduce_doc},
574 {NULL, NULL} /* sentinel */
575};
576
Raymond Hettingerad983e72003-11-12 14:32:26 +0000577PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000578
Raymond Hettingerad983e72003-11-12 14:32:26 +0000579static PyTypeObject teedataobject_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000580 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000581 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000582 sizeof(teedataobject), /* tp_basicsize */
583 0, /* tp_itemsize */
584 /* methods */
585 (destructor)teedataobject_dealloc, /* tp_dealloc */
586 0, /* tp_print */
587 0, /* tp_getattr */
588 0, /* tp_setattr */
589 0, /* tp_reserved */
590 0, /* tp_repr */
591 0, /* tp_as_number */
592 0, /* tp_as_sequence */
593 0, /* tp_as_mapping */
594 0, /* tp_hash */
595 0, /* tp_call */
596 0, /* tp_str */
597 PyObject_GenericGetAttr, /* tp_getattro */
598 0, /* tp_setattro */
599 0, /* tp_as_buffer */
600 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
601 teedataobject_doc, /* tp_doc */
602 (traverseproc)teedataobject_traverse, /* tp_traverse */
603 (inquiry)teedataobject_clear, /* tp_clear */
604 0, /* tp_richcompare */
605 0, /* tp_weaklistoffset */
606 0, /* tp_iter */
607 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000608 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000609 0, /* tp_members */
610 0, /* tp_getset */
611 0, /* tp_base */
612 0, /* tp_dict */
613 0, /* tp_descr_get */
614 0, /* tp_descr_set */
615 0, /* tp_dictoffset */
616 0, /* tp_init */
617 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000618 teedataobject_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000619 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000620};
621
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000622
623static PyTypeObject tee_type;
624
625static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000626tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000627{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000628 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000629
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000630 if (to->index >= LINKCELLS) {
631 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200632 if (link == NULL)
633 return NULL;
Serhiy Storchaka57a01d32016-04-10 18:05:40 +0300634 Py_SETREF(to->dataobj, (teedataobject *)link);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000635 to->index = 0;
636 }
637 value = teedataobject_getitem(to->dataobj, to->index);
638 if (value == NULL)
639 return NULL;
640 to->index++;
641 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000642}
643
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000644static int
645tee_traverse(teeobject *to, visitproc visit, void *arg)
646{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000647 Py_VISIT((PyObject *)to->dataobj);
648 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000649}
650
Raymond Hettingerad983e72003-11-12 14:32:26 +0000651static PyObject *
652tee_copy(teeobject *to)
653{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000654 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000655
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000656 newto = PyObject_GC_New(teeobject, &tee_type);
657 if (newto == NULL)
658 return NULL;
659 Py_INCREF(to->dataobj);
660 newto->dataobj = to->dataobj;
661 newto->index = to->index;
662 newto->weakreflist = NULL;
663 PyObject_GC_Track(newto);
664 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000665}
666
667PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
668
669static PyObject *
670tee_fromiterable(PyObject *iterable)
671{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000672 teeobject *to;
673 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000674
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000675 it = PyObject_GetIter(iterable);
676 if (it == NULL)
677 return NULL;
678 if (PyObject_TypeCheck(it, &tee_type)) {
679 to = (teeobject *)tee_copy((teeobject *)it);
680 goto done;
681 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000682
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000683 to = PyObject_GC_New(teeobject, &tee_type);
684 if (to == NULL)
685 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000686 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000687 if (!to->dataobj) {
688 PyObject_GC_Del(to);
689 to = NULL;
690 goto done;
691 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000692
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000693 to->index = 0;
694 to->weakreflist = NULL;
695 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000696done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000697 Py_XDECREF(it);
698 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000699}
700
701static PyObject *
702tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
703{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000704 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000705
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000706 if (!PyArg_UnpackTuple(args, "_tee", 1, 1, &iterable))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000707 return NULL;
708 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000709}
710
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000711static int
712tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000713{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000714 if (to->weakreflist != NULL)
715 PyObject_ClearWeakRefs((PyObject *) to);
716 Py_CLEAR(to->dataobj);
717 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000718}
719
720static void
721tee_dealloc(teeobject *to)
722{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000723 PyObject_GC_UnTrack(to);
724 tee_clear(to);
725 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000726}
727
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000728static PyObject *
729tee_reduce(teeobject *to)
730{
731 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
732}
733
734static PyObject *
735tee_setstate(teeobject *to, PyObject *state)
736{
737 teedataobject *tdo;
738 int index;
739 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index))
740 return NULL;
741 if (index < 0 || index > LINKCELLS) {
742 PyErr_SetString(PyExc_ValueError, "Index out of range");
743 return NULL;
744 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200745 Py_INCREF(tdo);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300746 Py_XSETREF(to->dataobj, tdo);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000747 to->index = index;
748 Py_RETURN_NONE;
749}
750
Raymond Hettingerad983e72003-11-12 14:32:26 +0000751PyDoc_STRVAR(teeobject_doc,
752"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000753
Raymond Hettingerad983e72003-11-12 14:32:26 +0000754static PyMethodDef tee_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000755 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000756 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
757 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000758 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000759};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000760
761static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000762 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000763 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000764 sizeof(teeobject), /* tp_basicsize */
765 0, /* tp_itemsize */
766 /* methods */
767 (destructor)tee_dealloc, /* tp_dealloc */
768 0, /* tp_print */
769 0, /* tp_getattr */
770 0, /* tp_setattr */
771 0, /* tp_reserved */
772 0, /* tp_repr */
773 0, /* tp_as_number */
774 0, /* tp_as_sequence */
775 0, /* tp_as_mapping */
776 0, /* tp_hash */
777 0, /* tp_call */
778 0, /* tp_str */
779 0, /* tp_getattro */
780 0, /* tp_setattro */
781 0, /* tp_as_buffer */
782 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
783 teeobject_doc, /* tp_doc */
784 (traverseproc)tee_traverse, /* tp_traverse */
785 (inquiry)tee_clear, /* tp_clear */
786 0, /* tp_richcompare */
787 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
788 PyObject_SelfIter, /* tp_iter */
789 (iternextfunc)tee_next, /* tp_iternext */
790 tee_methods, /* tp_methods */
791 0, /* tp_members */
792 0, /* tp_getset */
793 0, /* tp_base */
794 0, /* tp_dict */
795 0, /* tp_descr_get */
796 0, /* tp_descr_set */
797 0, /* tp_dictoffset */
798 0, /* tp_init */
799 0, /* tp_alloc */
800 tee_new, /* tp_new */
801 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000802};
803
Raymond Hettingerad983e72003-11-12 14:32:26 +0000804static PyObject *
805tee(PyObject *self, PyObject *args)
806{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000807 Py_ssize_t i, n=2;
808 PyObject *it, *iterable, *copyable, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200809 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000810
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000811 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
812 return NULL;
813 if (n < 0) {
814 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
815 return NULL;
816 }
817 result = PyTuple_New(n);
818 if (result == NULL)
819 return NULL;
820 if (n == 0)
821 return result;
822 it = PyObject_GetIter(iterable);
823 if (it == NULL) {
824 Py_DECREF(result);
825 return NULL;
826 }
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200827 if (!_PyObject_HasAttrId(it, &PyId___copy__)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000828 copyable = tee_fromiterable(it);
829 Py_DECREF(it);
830 if (copyable == NULL) {
831 Py_DECREF(result);
832 return NULL;
833 }
834 } else
835 copyable = it;
836 PyTuple_SET_ITEM(result, 0, copyable);
837 for (i=1 ; i<n ; i++) {
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200838
839 copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000840 if (copyable == NULL) {
841 Py_DECREF(result);
842 return NULL;
843 }
844 PyTuple_SET_ITEM(result, i, copyable);
845 }
846 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000847}
848
849PyDoc_STRVAR(tee_doc,
850"tee(iterable, n=2) --> tuple of n independent iterators.");
851
852
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000853/* cycle object **********************************************************/
854
855typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000856 PyObject_HEAD
857 PyObject *it;
858 PyObject *saved;
859 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000860} cycleobject;
861
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000862static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000863
864static PyObject *
865cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
866{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000867 PyObject *it;
868 PyObject *iterable;
869 PyObject *saved;
870 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000871
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000872 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
873 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000874
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000875 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
876 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000877
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000878 /* Get iterator. */
879 it = PyObject_GetIter(iterable);
880 if (it == NULL)
881 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000882
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000883 saved = PyList_New(0);
884 if (saved == NULL) {
885 Py_DECREF(it);
886 return NULL;
887 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000888
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000889 /* create cycleobject structure */
890 lz = (cycleobject *)type->tp_alloc(type, 0);
891 if (lz == NULL) {
892 Py_DECREF(it);
893 Py_DECREF(saved);
894 return NULL;
895 }
896 lz->it = it;
897 lz->saved = saved;
898 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000899
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000900 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000901}
902
903static void
904cycle_dealloc(cycleobject *lz)
905{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000906 PyObject_GC_UnTrack(lz);
907 Py_XDECREF(lz->saved);
908 Py_XDECREF(lz->it);
909 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000910}
911
912static int
913cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
914{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000915 Py_VISIT(lz->it);
916 Py_VISIT(lz->saved);
917 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000918}
919
920static PyObject *
921cycle_next(cycleobject *lz)
922{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000923 PyObject *item;
924 PyObject *it;
925 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000926
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000927 while (1) {
928 item = PyIter_Next(lz->it);
929 if (item != NULL) {
930 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
931 Py_DECREF(item);
932 return NULL;
933 }
934 return item;
935 }
936 if (PyErr_Occurred()) {
937 if (PyErr_ExceptionMatches(PyExc_StopIteration))
938 PyErr_Clear();
939 else
940 return NULL;
941 }
942 if (PyList_Size(lz->saved) == 0)
943 return NULL;
944 it = PyObject_GetIter(lz->saved);
945 if (it == NULL)
946 return NULL;
947 tmp = lz->it;
948 lz->it = it;
949 lz->firstpass = 1;
950 Py_DECREF(tmp);
951 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000952}
953
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000954static PyObject *
955cycle_reduce(cycleobject *lz)
956{
957 /* Create a new cycle with the iterator tuple, then set
958 * the saved state on it.
959 */
Serhiy Storchakad269b5e2013-01-25 13:38:56 +0200960 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz),
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000961 lz->it, lz->saved, lz->firstpass);
962 }
963
964static PyObject *
965cycle_setstate(cycleobject *lz, PyObject *state)
966{
967 PyObject *saved=NULL;
968 int firstpass;
969 if (!PyArg_ParseTuple(state, "Oi", &saved, &firstpass))
970 return NULL;
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200971 Py_XINCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300972 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000973 lz->firstpass = firstpass != 0;
974 Py_RETURN_NONE;
975}
976
977static PyMethodDef cycle_methods[] = {
978 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
979 reduce_doc},
980 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
981 setstate_doc},
982 {NULL, NULL} /* sentinel */
983};
984
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000985PyDoc_STRVAR(cycle_doc,
986"cycle(iterable) --> cycle object\n\
987\n\
988Return elements from the iterable until it is exhausted.\n\
989Then repeat the sequence indefinitely.");
990
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000991static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000992 PyVarObject_HEAD_INIT(NULL, 0)
993 "itertools.cycle", /* tp_name */
994 sizeof(cycleobject), /* tp_basicsize */
995 0, /* tp_itemsize */
996 /* methods */
997 (destructor)cycle_dealloc, /* tp_dealloc */
998 0, /* tp_print */
999 0, /* tp_getattr */
1000 0, /* tp_setattr */
1001 0, /* tp_reserved */
1002 0, /* tp_repr */
1003 0, /* tp_as_number */
1004 0, /* tp_as_sequence */
1005 0, /* tp_as_mapping */
1006 0, /* tp_hash */
1007 0, /* tp_call */
1008 0, /* tp_str */
1009 PyObject_GenericGetAttr, /* tp_getattro */
1010 0, /* tp_setattro */
1011 0, /* tp_as_buffer */
1012 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1013 Py_TPFLAGS_BASETYPE, /* tp_flags */
1014 cycle_doc, /* tp_doc */
1015 (traverseproc)cycle_traverse, /* tp_traverse */
1016 0, /* tp_clear */
1017 0, /* tp_richcompare */
1018 0, /* tp_weaklistoffset */
1019 PyObject_SelfIter, /* tp_iter */
1020 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001021 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001022 0, /* tp_members */
1023 0, /* tp_getset */
1024 0, /* tp_base */
1025 0, /* tp_dict */
1026 0, /* tp_descr_get */
1027 0, /* tp_descr_set */
1028 0, /* tp_dictoffset */
1029 0, /* tp_init */
1030 0, /* tp_alloc */
1031 cycle_new, /* tp_new */
1032 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001033};
1034
1035
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001036/* dropwhile object **********************************************************/
1037
1038typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001039 PyObject_HEAD
1040 PyObject *func;
1041 PyObject *it;
1042 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001043} dropwhileobject;
1044
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001045static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001046
1047static PyObject *
1048dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1049{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001050 PyObject *func, *seq;
1051 PyObject *it;
1052 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001054 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
1055 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001056
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001057 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
1058 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001059
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001060 /* Get iterator. */
1061 it = PyObject_GetIter(seq);
1062 if (it == NULL)
1063 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001064
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001065 /* create dropwhileobject structure */
1066 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1067 if (lz == NULL) {
1068 Py_DECREF(it);
1069 return NULL;
1070 }
1071 Py_INCREF(func);
1072 lz->func = func;
1073 lz->it = it;
1074 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001075
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001076 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001077}
1078
1079static void
1080dropwhile_dealloc(dropwhileobject *lz)
1081{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001082 PyObject_GC_UnTrack(lz);
1083 Py_XDECREF(lz->func);
1084 Py_XDECREF(lz->it);
1085 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001086}
1087
1088static int
1089dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1090{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001091 Py_VISIT(lz->it);
1092 Py_VISIT(lz->func);
1093 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001094}
1095
1096static PyObject *
1097dropwhile_next(dropwhileobject *lz)
1098{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001099 PyObject *item, *good;
1100 PyObject *it = lz->it;
1101 long ok;
1102 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001103
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001104 iternext = *Py_TYPE(it)->tp_iternext;
1105 for (;;) {
1106 item = iternext(it);
1107 if (item == NULL)
1108 return NULL;
1109 if (lz->start == 1)
1110 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001111
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001112 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1113 if (good == NULL) {
1114 Py_DECREF(item);
1115 return NULL;
1116 }
1117 ok = PyObject_IsTrue(good);
1118 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001119 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001120 lz->start = 1;
1121 return item;
1122 }
1123 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001124 if (ok < 0)
1125 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001126 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001127}
1128
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001129static PyObject *
1130dropwhile_reduce(dropwhileobject *lz)
1131{
1132 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1133 lz->func, lz->it, lz->start);
1134}
1135
1136static PyObject *
1137dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1138{
1139 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001140 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001141 return NULL;
1142 lz->start = start;
1143 Py_RETURN_NONE;
1144}
1145
1146static PyMethodDef dropwhile_methods[] = {
1147 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1148 reduce_doc},
1149 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1150 setstate_doc},
1151 {NULL, NULL} /* sentinel */
1152};
1153
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001154PyDoc_STRVAR(dropwhile_doc,
1155"dropwhile(predicate, iterable) --> dropwhile object\n\
1156\n\
1157Drop items from the iterable while predicate(item) is true.\n\
1158Afterwards, return every element until the iterable is exhausted.");
1159
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001160static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001161 PyVarObject_HEAD_INIT(NULL, 0)
1162 "itertools.dropwhile", /* tp_name */
1163 sizeof(dropwhileobject), /* tp_basicsize */
1164 0, /* tp_itemsize */
1165 /* methods */
1166 (destructor)dropwhile_dealloc, /* tp_dealloc */
1167 0, /* tp_print */
1168 0, /* tp_getattr */
1169 0, /* tp_setattr */
1170 0, /* tp_reserved */
1171 0, /* tp_repr */
1172 0, /* tp_as_number */
1173 0, /* tp_as_sequence */
1174 0, /* tp_as_mapping */
1175 0, /* tp_hash */
1176 0, /* tp_call */
1177 0, /* tp_str */
1178 PyObject_GenericGetAttr, /* tp_getattro */
1179 0, /* tp_setattro */
1180 0, /* tp_as_buffer */
1181 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1182 Py_TPFLAGS_BASETYPE, /* tp_flags */
1183 dropwhile_doc, /* tp_doc */
1184 (traverseproc)dropwhile_traverse, /* tp_traverse */
1185 0, /* tp_clear */
1186 0, /* tp_richcompare */
1187 0, /* tp_weaklistoffset */
1188 PyObject_SelfIter, /* tp_iter */
1189 (iternextfunc)dropwhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001190 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001191 0, /* tp_members */
1192 0, /* tp_getset */
1193 0, /* tp_base */
1194 0, /* tp_dict */
1195 0, /* tp_descr_get */
1196 0, /* tp_descr_set */
1197 0, /* tp_dictoffset */
1198 0, /* tp_init */
1199 0, /* tp_alloc */
1200 dropwhile_new, /* tp_new */
1201 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001202};
1203
1204
1205/* takewhile object **********************************************************/
1206
1207typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001208 PyObject_HEAD
1209 PyObject *func;
1210 PyObject *it;
1211 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001212} takewhileobject;
1213
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001214static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001215
1216static PyObject *
1217takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1218{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001219 PyObject *func, *seq;
1220 PyObject *it;
1221 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001222
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001223 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1224 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001225
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001226 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1227 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001228
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001229 /* Get iterator. */
1230 it = PyObject_GetIter(seq);
1231 if (it == NULL)
1232 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001233
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001234 /* create takewhileobject structure */
1235 lz = (takewhileobject *)type->tp_alloc(type, 0);
1236 if (lz == NULL) {
1237 Py_DECREF(it);
1238 return NULL;
1239 }
1240 Py_INCREF(func);
1241 lz->func = func;
1242 lz->it = it;
1243 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001244
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001245 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001246}
1247
1248static void
1249takewhile_dealloc(takewhileobject *lz)
1250{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001251 PyObject_GC_UnTrack(lz);
1252 Py_XDECREF(lz->func);
1253 Py_XDECREF(lz->it);
1254 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001255}
1256
1257static int
1258takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1259{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001260 Py_VISIT(lz->it);
1261 Py_VISIT(lz->func);
1262 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001263}
1264
1265static PyObject *
1266takewhile_next(takewhileobject *lz)
1267{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001268 PyObject *item, *good;
1269 PyObject *it = lz->it;
1270 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001271
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001272 if (lz->stop == 1)
1273 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001274
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001275 item = (*Py_TYPE(it)->tp_iternext)(it);
1276 if (item == NULL)
1277 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001278
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001279 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1280 if (good == NULL) {
1281 Py_DECREF(item);
1282 return NULL;
1283 }
1284 ok = PyObject_IsTrue(good);
1285 Py_DECREF(good);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001286 if (ok == 1)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001287 return item;
1288 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001289 if (ok == 0)
1290 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001291 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001292}
1293
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001294static PyObject *
1295takewhile_reduce(takewhileobject *lz)
1296{
1297 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1298 lz->func, lz->it, lz->stop);
1299}
1300
1301static PyObject *
1302takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1303{
1304 int stop = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001305 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001306 return NULL;
1307 lz->stop = stop;
1308 Py_RETURN_NONE;
1309}
1310
1311static PyMethodDef takewhile_reduce_methods[] = {
1312 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1313 reduce_doc},
1314 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1315 setstate_doc},
1316 {NULL, NULL} /* sentinel */
1317};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001318PyDoc_STRVAR(takewhile_doc,
1319"takewhile(predicate, iterable) --> takewhile object\n\
1320\n\
1321Return successive entries from an iterable as long as the \n\
1322predicate evaluates to true for each entry.");
1323
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001324static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001325 PyVarObject_HEAD_INIT(NULL, 0)
1326 "itertools.takewhile", /* tp_name */
1327 sizeof(takewhileobject), /* tp_basicsize */
1328 0, /* tp_itemsize */
1329 /* methods */
1330 (destructor)takewhile_dealloc, /* tp_dealloc */
1331 0, /* tp_print */
1332 0, /* tp_getattr */
1333 0, /* tp_setattr */
1334 0, /* tp_reserved */
1335 0, /* tp_repr */
1336 0, /* tp_as_number */
1337 0, /* tp_as_sequence */
1338 0, /* tp_as_mapping */
1339 0, /* tp_hash */
1340 0, /* tp_call */
1341 0, /* tp_str */
1342 PyObject_GenericGetAttr, /* tp_getattro */
1343 0, /* tp_setattro */
1344 0, /* tp_as_buffer */
1345 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1346 Py_TPFLAGS_BASETYPE, /* tp_flags */
1347 takewhile_doc, /* tp_doc */
1348 (traverseproc)takewhile_traverse, /* tp_traverse */
1349 0, /* tp_clear */
1350 0, /* tp_richcompare */
1351 0, /* tp_weaklistoffset */
1352 PyObject_SelfIter, /* tp_iter */
1353 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001354 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001355 0, /* tp_members */
1356 0, /* tp_getset */
1357 0, /* tp_base */
1358 0, /* tp_dict */
1359 0, /* tp_descr_get */
1360 0, /* tp_descr_set */
1361 0, /* tp_dictoffset */
1362 0, /* tp_init */
1363 0, /* tp_alloc */
1364 takewhile_new, /* tp_new */
1365 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001366};
1367
1368
1369/* islice object ************************************************************/
1370
1371typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001372 PyObject_HEAD
1373 PyObject *it;
1374 Py_ssize_t next;
1375 Py_ssize_t stop;
1376 Py_ssize_t step;
1377 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001378} isliceobject;
1379
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001380static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001381
1382static PyObject *
1383islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1384{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001385 PyObject *seq;
1386 Py_ssize_t start=0, stop=-1, step=1;
1387 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1388 Py_ssize_t numargs;
1389 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001390
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001391 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1392 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001393
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001394 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1395 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001396
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001397 numargs = PyTuple_Size(args);
1398 if (numargs == 2) {
1399 if (a1 != Py_None) {
1400 stop = PyLong_AsSsize_t(a1);
1401 if (stop == -1) {
1402 if (PyErr_Occurred())
1403 PyErr_Clear();
1404 PyErr_SetString(PyExc_ValueError,
1405 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1406 return NULL;
1407 }
1408 }
1409 } else {
1410 if (a1 != Py_None)
1411 start = PyLong_AsSsize_t(a1);
1412 if (start == -1 && PyErr_Occurred())
1413 PyErr_Clear();
1414 if (a2 != Py_None) {
1415 stop = PyLong_AsSsize_t(a2);
1416 if (stop == -1) {
1417 if (PyErr_Occurred())
1418 PyErr_Clear();
1419 PyErr_SetString(PyExc_ValueError,
1420 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1421 return NULL;
1422 }
1423 }
1424 }
1425 if (start<0 || stop<-1) {
1426 PyErr_SetString(PyExc_ValueError,
1427 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1428 return NULL;
1429 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001430
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001431 if (a3 != NULL) {
1432 if (a3 != Py_None)
1433 step = PyLong_AsSsize_t(a3);
1434 if (step == -1 && PyErr_Occurred())
1435 PyErr_Clear();
1436 }
1437 if (step<1) {
1438 PyErr_SetString(PyExc_ValueError,
1439 "Step for islice() must be a positive integer or None.");
1440 return NULL;
1441 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001442
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001443 /* Get iterator. */
1444 it = PyObject_GetIter(seq);
1445 if (it == NULL)
1446 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001447
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001448 /* create isliceobject structure */
1449 lz = (isliceobject *)type->tp_alloc(type, 0);
1450 if (lz == NULL) {
1451 Py_DECREF(it);
1452 return NULL;
1453 }
1454 lz->it = it;
1455 lz->next = start;
1456 lz->stop = stop;
1457 lz->step = step;
1458 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001460 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001461}
1462
1463static void
1464islice_dealloc(isliceobject *lz)
1465{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001466 PyObject_GC_UnTrack(lz);
1467 Py_XDECREF(lz->it);
1468 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469}
1470
1471static int
1472islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1473{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001474 Py_VISIT(lz->it);
1475 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001476}
1477
1478static PyObject *
1479islice_next(isliceobject *lz)
1480{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001481 PyObject *item;
1482 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001483 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001484 Py_ssize_t oldnext;
1485 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001486
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001487 if (it == NULL)
1488 return NULL;
1489
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001490 iternext = *Py_TYPE(it)->tp_iternext;
1491 while (lz->cnt < lz->next) {
1492 item = iternext(it);
1493 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001494 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001495 Py_DECREF(item);
1496 lz->cnt++;
1497 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001498 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001499 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001500 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 lz->cnt++;
1504 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001505 /* The (size_t) cast below avoids the danger of undefined
1506 behaviour from signed integer overflow. */
1507 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001508 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1509 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001510 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001511
1512empty:
1513 Py_CLEAR(lz->it);
1514 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001515}
1516
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001517static PyObject *
1518islice_reduce(isliceobject *lz)
1519{
1520 /* When unpickled, generate a new object with the same bounds,
1521 * then 'setstate' with the next and count
1522 */
1523 PyObject *stop;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001524 if (lz->it == NULL) {
1525 PyObject *empty_list;
1526 PyObject *empty_it;
1527 empty_list = PyList_New(0);
1528 if (empty_list == NULL)
1529 return NULL;
1530 empty_it = PyObject_GetIter(empty_list);
1531 Py_DECREF(empty_list);
1532 if (empty_it == NULL)
1533 return NULL;
1534 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1535 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001536 if (lz->stop == -1) {
1537 stop = Py_None;
1538 Py_INCREF(stop);
1539 } else {
1540 stop = PyLong_FromSsize_t(lz->stop);
1541 if (stop == NULL)
1542 return NULL;
1543 }
1544 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1545 lz->it, lz->next, stop, lz->step,
1546 lz->cnt);
1547}
1548
1549static PyObject *
1550islice_setstate(isliceobject *lz, PyObject *state)
1551{
1552 Py_ssize_t cnt = PyLong_AsSsize_t(state);
1553 if (cnt == -1 && PyErr_Occurred())
1554 return NULL;
1555 lz->cnt = cnt;
1556 Py_RETURN_NONE;
1557}
1558
1559static PyMethodDef islice_methods[] = {
1560 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1561 reduce_doc},
1562 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1563 setstate_doc},
1564 {NULL, NULL} /* sentinel */
1565};
1566
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001567PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001568"islice(iterable, stop) --> islice object\n\
1569islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001570\n\
1571Return an iterator whose next() method returns selected values from an\n\
1572iterable. If start is specified, will skip all preceding elements;\n\
1573otherwise, start defaults to zero. Step defaults to one. If\n\
1574specified as another value, step determines how many values are \n\
1575skipped between successive calls. Works like a slice() on a list\n\
1576but returns an iterator.");
1577
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001578static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001579 PyVarObject_HEAD_INIT(NULL, 0)
1580 "itertools.islice", /* tp_name */
1581 sizeof(isliceobject), /* tp_basicsize */
1582 0, /* tp_itemsize */
1583 /* methods */
1584 (destructor)islice_dealloc, /* tp_dealloc */
1585 0, /* tp_print */
1586 0, /* tp_getattr */
1587 0, /* tp_setattr */
1588 0, /* tp_reserved */
1589 0, /* tp_repr */
1590 0, /* tp_as_number */
1591 0, /* tp_as_sequence */
1592 0, /* tp_as_mapping */
1593 0, /* tp_hash */
1594 0, /* tp_call */
1595 0, /* tp_str */
1596 PyObject_GenericGetAttr, /* tp_getattro */
1597 0, /* tp_setattro */
1598 0, /* tp_as_buffer */
1599 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1600 Py_TPFLAGS_BASETYPE, /* tp_flags */
1601 islice_doc, /* tp_doc */
1602 (traverseproc)islice_traverse, /* tp_traverse */
1603 0, /* tp_clear */
1604 0, /* tp_richcompare */
1605 0, /* tp_weaklistoffset */
1606 PyObject_SelfIter, /* tp_iter */
1607 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001608 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001609 0, /* tp_members */
1610 0, /* tp_getset */
1611 0, /* tp_base */
1612 0, /* tp_dict */
1613 0, /* tp_descr_get */
1614 0, /* tp_descr_set */
1615 0, /* tp_dictoffset */
1616 0, /* tp_init */
1617 0, /* tp_alloc */
1618 islice_new, /* tp_new */
1619 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001620};
1621
1622
1623/* starmap object ************************************************************/
1624
1625typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001626 PyObject_HEAD
1627 PyObject *func;
1628 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001629} starmapobject;
1630
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001631static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001632
1633static PyObject *
1634starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1635{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001636 PyObject *func, *seq;
1637 PyObject *it;
1638 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001639
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001640 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1641 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001642
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001643 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1644 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001645
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001646 /* Get iterator. */
1647 it = PyObject_GetIter(seq);
1648 if (it == NULL)
1649 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001650
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001651 /* create starmapobject structure */
1652 lz = (starmapobject *)type->tp_alloc(type, 0);
1653 if (lz == NULL) {
1654 Py_DECREF(it);
1655 return NULL;
1656 }
1657 Py_INCREF(func);
1658 lz->func = func;
1659 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001660
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001661 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001662}
1663
1664static void
1665starmap_dealloc(starmapobject *lz)
1666{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001667 PyObject_GC_UnTrack(lz);
1668 Py_XDECREF(lz->func);
1669 Py_XDECREF(lz->it);
1670 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001671}
1672
1673static int
1674starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1675{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001676 Py_VISIT(lz->it);
1677 Py_VISIT(lz->func);
1678 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001679}
1680
1681static PyObject *
1682starmap_next(starmapobject *lz)
1683{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001684 PyObject *args;
1685 PyObject *result;
1686 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001687
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001688 args = (*Py_TYPE(it)->tp_iternext)(it);
1689 if (args == NULL)
1690 return NULL;
1691 if (!PyTuple_CheckExact(args)) {
1692 PyObject *newargs = PySequence_Tuple(args);
1693 Py_DECREF(args);
1694 if (newargs == NULL)
1695 return NULL;
1696 args = newargs;
1697 }
1698 result = PyObject_Call(lz->func, args, NULL);
1699 Py_DECREF(args);
1700 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001701}
1702
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001703static PyObject *
1704starmap_reduce(starmapobject *lz)
1705{
1706 /* Just pickle the iterator */
1707 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1708}
1709
1710static PyMethodDef starmap_methods[] = {
1711 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1712 reduce_doc},
1713 {NULL, NULL} /* sentinel */
1714};
1715
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001716PyDoc_STRVAR(starmap_doc,
1717"starmap(function, sequence) --> starmap object\n\
1718\n\
1719Return an iterator whose values are returned from the function evaluated\n\
Serhiy Storchakad65c9492015-11-02 14:10:23 +02001720with an argument tuple taken from the given sequence.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001721
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001722static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001723 PyVarObject_HEAD_INIT(NULL, 0)
1724 "itertools.starmap", /* tp_name */
1725 sizeof(starmapobject), /* tp_basicsize */
1726 0, /* tp_itemsize */
1727 /* methods */
1728 (destructor)starmap_dealloc, /* tp_dealloc */
1729 0, /* tp_print */
1730 0, /* tp_getattr */
1731 0, /* tp_setattr */
1732 0, /* tp_reserved */
1733 0, /* tp_repr */
1734 0, /* tp_as_number */
1735 0, /* tp_as_sequence */
1736 0, /* tp_as_mapping */
1737 0, /* tp_hash */
1738 0, /* tp_call */
1739 0, /* tp_str */
1740 PyObject_GenericGetAttr, /* tp_getattro */
1741 0, /* tp_setattro */
1742 0, /* tp_as_buffer */
1743 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1744 Py_TPFLAGS_BASETYPE, /* tp_flags */
1745 starmap_doc, /* tp_doc */
1746 (traverseproc)starmap_traverse, /* tp_traverse */
1747 0, /* tp_clear */
1748 0, /* tp_richcompare */
1749 0, /* tp_weaklistoffset */
1750 PyObject_SelfIter, /* tp_iter */
1751 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001752 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001753 0, /* tp_members */
1754 0, /* tp_getset */
1755 0, /* tp_base */
1756 0, /* tp_dict */
1757 0, /* tp_descr_get */
1758 0, /* tp_descr_set */
1759 0, /* tp_dictoffset */
1760 0, /* tp_init */
1761 0, /* tp_alloc */
1762 starmap_new, /* tp_new */
1763 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001764};
1765
1766
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001767/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001768
1769typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001770 PyObject_HEAD
1771 PyObject *source; /* Iterator over input iterables */
1772 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001773} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001774
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001775static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001776
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001777static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001778chain_new_internal(PyTypeObject *type, PyObject *source)
1779{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001780 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001781
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001782 lz = (chainobject *)type->tp_alloc(type, 0);
1783 if (lz == NULL) {
1784 Py_DECREF(source);
1785 return NULL;
1786 }
1787
1788 lz->source = source;
1789 lz->active = NULL;
1790 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001791}
1792
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001793static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001794chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001795{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001796 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001797
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001798 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1799 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001800
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001801 source = PyObject_GetIter(args);
1802 if (source == NULL)
1803 return NULL;
1804
1805 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001806}
1807
1808static PyObject *
1809chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1810{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001811 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001812
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001813 source = PyObject_GetIter(arg);
1814 if (source == NULL)
1815 return NULL;
1816
1817 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001818}
1819
1820static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001821chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001822{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001823 PyObject_GC_UnTrack(lz);
1824 Py_XDECREF(lz->active);
1825 Py_XDECREF(lz->source);
1826 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001827}
1828
Raymond Hettinger2012f172003-02-07 05:32:58 +00001829static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001830chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001831{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001832 Py_VISIT(lz->source);
1833 Py_VISIT(lz->active);
1834 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001835}
1836
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001837static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001838chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001839{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001840 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001841
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001842 if (lz->source == NULL)
1843 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001844
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001845 if (lz->active == NULL) {
1846 PyObject *iterable = PyIter_Next(lz->source);
1847 if (iterable == NULL) {
1848 Py_CLEAR(lz->source);
1849 return NULL; /* no more input sources */
1850 }
1851 lz->active = PyObject_GetIter(iterable);
1852 Py_DECREF(iterable);
1853 if (lz->active == NULL) {
1854 Py_CLEAR(lz->source);
1855 return NULL; /* input not iterable */
1856 }
1857 }
1858 item = PyIter_Next(lz->active);
1859 if (item != NULL)
1860 return item;
1861 if (PyErr_Occurred()) {
1862 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1863 PyErr_Clear();
1864 else
1865 return NULL; /* input raised an exception */
1866 }
1867 Py_CLEAR(lz->active);
1868 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001869}
1870
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001871static PyObject *
1872chain_reduce(chainobject *lz)
1873{
1874 if (lz->source) {
1875 /* we can't pickle function objects (itertools.from_iterable) so
1876 * we must use setstate to replace the iterable. One day we
1877 * will fix pickling of functions
1878 */
1879 if (lz->active) {
1880 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1881 } else {
1882 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1883 }
1884 } else {
1885 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1886 }
1887 return NULL;
1888}
1889
1890static PyObject *
1891chain_setstate(chainobject *lz, PyObject *state)
1892{
1893 PyObject *source, *active=NULL;
1894 if (! PyArg_ParseTuple(state, "O|O", &source, &active))
1895 return NULL;
1896
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02001897 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001898 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02001899 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001900 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001901 Py_RETURN_NONE;
1902}
1903
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001904PyDoc_STRVAR(chain_doc,
1905"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001906\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001907Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001908first iterable until it is exhausted, then elements from the next\n\
1909iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001910
Christian Heimesf16baeb2008-02-29 14:57:44 +00001911PyDoc_STRVAR(chain_from_iterable_doc,
1912"chain.from_iterable(iterable) --> chain object\n\
1913\n\
Martin Pantereb995702016-07-28 01:11:04 +00001914Alternate chain() constructor taking a single iterable argument\n\
Christian Heimesf16baeb2008-02-29 14:57:44 +00001915that evaluates lazily.");
1916
1917static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001918 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1919 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001920 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1921 reduce_doc},
1922 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1923 setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001924 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001925};
1926
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001927static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001928 PyVarObject_HEAD_INIT(NULL, 0)
1929 "itertools.chain", /* tp_name */
1930 sizeof(chainobject), /* tp_basicsize */
1931 0, /* tp_itemsize */
1932 /* methods */
1933 (destructor)chain_dealloc, /* tp_dealloc */
1934 0, /* tp_print */
1935 0, /* tp_getattr */
1936 0, /* tp_setattr */
1937 0, /* tp_reserved */
1938 0, /* tp_repr */
1939 0, /* tp_as_number */
1940 0, /* tp_as_sequence */
1941 0, /* tp_as_mapping */
1942 0, /* tp_hash */
1943 0, /* tp_call */
1944 0, /* tp_str */
1945 PyObject_GenericGetAttr, /* tp_getattro */
1946 0, /* tp_setattro */
1947 0, /* tp_as_buffer */
1948 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1949 Py_TPFLAGS_BASETYPE, /* tp_flags */
1950 chain_doc, /* tp_doc */
1951 (traverseproc)chain_traverse, /* tp_traverse */
1952 0, /* tp_clear */
1953 0, /* tp_richcompare */
1954 0, /* tp_weaklistoffset */
1955 PyObject_SelfIter, /* tp_iter */
1956 (iternextfunc)chain_next, /* tp_iternext */
1957 chain_methods, /* tp_methods */
1958 0, /* tp_members */
1959 0, /* tp_getset */
1960 0, /* tp_base */
1961 0, /* tp_dict */
1962 0, /* tp_descr_get */
1963 0, /* tp_descr_set */
1964 0, /* tp_dictoffset */
1965 0, /* tp_init */
1966 0, /* tp_alloc */
1967 chain_new, /* tp_new */
1968 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001969};
1970
1971
Christian Heimesc3f30c42008-02-22 16:37:40 +00001972/* product object ************************************************************/
1973
1974typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001975 PyObject_HEAD
1976 PyObject *pools; /* tuple of pool tuples */
1977 Py_ssize_t *indices; /* one index per pool */
1978 PyObject *result; /* most recently returned result tuple */
1979 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001980} productobject;
1981
1982static PyTypeObject product_type;
1983
1984static PyObject *
1985product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1986{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001987 productobject *lz;
1988 Py_ssize_t nargs, npools, repeat=1;
1989 PyObject *pools = NULL;
1990 Py_ssize_t *indices = NULL;
1991 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001992
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001993 if (kwds != NULL) {
1994 char *kwlist[] = {"repeat", 0};
1995 PyObject *tmpargs = PyTuple_New(0);
1996 if (tmpargs == NULL)
1997 return NULL;
1998 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1999 Py_DECREF(tmpargs);
2000 return NULL;
2001 }
2002 Py_DECREF(tmpargs);
2003 if (repeat < 0) {
2004 PyErr_SetString(PyExc_ValueError,
2005 "repeat argument cannot be negative");
2006 return NULL;
2007 }
2008 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002009
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002010 assert(PyTuple_CheckExact(args));
2011 if (repeat == 0) {
2012 nargs = 0;
2013 } else {
2014 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002015 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002016 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2017 return NULL;
2018 }
2019 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002020 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002021
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002022 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002023 if (indices == NULL) {
2024 PyErr_NoMemory();
2025 goto error;
2026 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002027
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002028 pools = PyTuple_New(npools);
2029 if (pools == NULL)
2030 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002031
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002032 for (i=0; i < nargs ; ++i) {
2033 PyObject *item = PyTuple_GET_ITEM(args, i);
2034 PyObject *pool = PySequence_Tuple(item);
2035 if (pool == NULL)
2036 goto error;
2037 PyTuple_SET_ITEM(pools, i, pool);
2038 indices[i] = 0;
2039 }
2040 for ( ; i < npools; ++i) {
2041 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2042 Py_INCREF(pool);
2043 PyTuple_SET_ITEM(pools, i, pool);
2044 indices[i] = 0;
2045 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002046
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002047 /* create productobject structure */
2048 lz = (productobject *)type->tp_alloc(type, 0);
2049 if (lz == NULL)
2050 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002051
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002052 lz->pools = pools;
2053 lz->indices = indices;
2054 lz->result = NULL;
2055 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002056
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002057 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002058
2059error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002060 if (indices != NULL)
2061 PyMem_Free(indices);
2062 Py_XDECREF(pools);
2063 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002064}
2065
2066static void
2067product_dealloc(productobject *lz)
2068{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002069 PyObject_GC_UnTrack(lz);
2070 Py_XDECREF(lz->pools);
2071 Py_XDECREF(lz->result);
2072 if (lz->indices != NULL)
2073 PyMem_Free(lz->indices);
2074 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002075}
2076
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002077static PyObject *
2078product_sizeof(productobject *lz, void *unused)
2079{
2080 Py_ssize_t res;
2081
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002082 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002083 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2084 return PyLong_FromSsize_t(res);
2085}
2086
2087PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2088
Christian Heimesc3f30c42008-02-22 16:37:40 +00002089static int
2090product_traverse(productobject *lz, visitproc visit, void *arg)
2091{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002092 Py_VISIT(lz->pools);
2093 Py_VISIT(lz->result);
2094 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002095}
2096
2097static PyObject *
2098product_next(productobject *lz)
2099{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002100 PyObject *pool;
2101 PyObject *elem;
2102 PyObject *oldelem;
2103 PyObject *pools = lz->pools;
2104 PyObject *result = lz->result;
2105 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2106 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002107
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002108 if (lz->stopped)
2109 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002110
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002111 if (result == NULL) {
2112 /* On the first pass, return an initial tuple filled with the
2113 first element from each pool. */
2114 result = PyTuple_New(npools);
2115 if (result == NULL)
2116 goto empty;
2117 lz->result = result;
2118 for (i=0; i < npools; i++) {
2119 pool = PyTuple_GET_ITEM(pools, i);
2120 if (PyTuple_GET_SIZE(pool) == 0)
2121 goto empty;
2122 elem = PyTuple_GET_ITEM(pool, 0);
2123 Py_INCREF(elem);
2124 PyTuple_SET_ITEM(result, i, elem);
2125 }
2126 } else {
2127 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002128
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002129 /* Copy the previous result tuple or re-use it if available */
2130 if (Py_REFCNT(result) > 1) {
2131 PyObject *old_result = result;
2132 result = PyTuple_New(npools);
2133 if (result == NULL)
2134 goto empty;
2135 lz->result = result;
2136 for (i=0; i < npools; i++) {
2137 elem = PyTuple_GET_ITEM(old_result, i);
2138 Py_INCREF(elem);
2139 PyTuple_SET_ITEM(result, i, elem);
2140 }
2141 Py_DECREF(old_result);
2142 }
2143 /* Now, we've got the only copy so we can update it in-place */
2144 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002145
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002146 /* Update the pool indices right-to-left. Only advance to the
2147 next pool when the previous one rolls-over */
2148 for (i=npools-1 ; i >= 0 ; i--) {
2149 pool = PyTuple_GET_ITEM(pools, i);
2150 indices[i]++;
2151 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2152 /* Roll-over and advance to next pool */
2153 indices[i] = 0;
2154 elem = PyTuple_GET_ITEM(pool, 0);
2155 Py_INCREF(elem);
2156 oldelem = PyTuple_GET_ITEM(result, i);
2157 PyTuple_SET_ITEM(result, i, elem);
2158 Py_DECREF(oldelem);
2159 } else {
2160 /* No rollover. Just increment and stop here. */
2161 elem = PyTuple_GET_ITEM(pool, indices[i]);
2162 Py_INCREF(elem);
2163 oldelem = PyTuple_GET_ITEM(result, i);
2164 PyTuple_SET_ITEM(result, i, elem);
2165 Py_DECREF(oldelem);
2166 break;
2167 }
2168 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002169
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002170 /* If i is negative, then the indices have all rolled-over
2171 and we're done. */
2172 if (i < 0)
2173 goto empty;
2174 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002175
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002176 Py_INCREF(result);
2177 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002178
2179empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002180 lz->stopped = 1;
2181 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002182}
2183
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002184static PyObject *
2185product_reduce(productobject *lz)
2186{
2187 if (lz->stopped) {
2188 return Py_BuildValue("O(())", Py_TYPE(lz));
2189 } else if (lz->result == NULL) {
2190 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2191 } else {
2192 PyObject *indices;
2193 Py_ssize_t n, i;
2194
2195 /* we must pickle the indices use them for setstate, and
2196 * additionally indicate that the iterator has started
2197 */
2198 n = PyTuple_GET_SIZE(lz->pools);
2199 indices = PyTuple_New(n);
2200 if (indices == NULL)
2201 return NULL;
2202 for (i=0; i<n; i++){
2203 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2204 if (!index) {
2205 Py_DECREF(indices);
2206 return NULL;
2207 }
2208 PyTuple_SET_ITEM(indices, i, index);
2209 }
2210 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2211 }
2212}
2213
2214static PyObject *
2215product_setstate(productobject *lz, PyObject *state)
2216{
2217 PyObject *result;
2218 Py_ssize_t n, i;
2219
2220 n = PyTuple_GET_SIZE(lz->pools);
2221 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2222 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2223 return NULL;
2224 }
2225 for (i=0; i<n; i++)
2226 {
2227 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2228 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002229 PyObject* pool;
2230 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002231 if (index < 0 && PyErr_Occurred())
2232 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002233 pool = PyTuple_GET_ITEM(lz->pools, i);
2234 poolsize = PyTuple_GET_SIZE(pool);
2235 if (poolsize == 0) {
2236 lz->stopped = 1;
2237 Py_RETURN_NONE;
2238 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002239 /* clamp the index */
2240 if (index < 0)
2241 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002242 else if (index > poolsize-1)
2243 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002244 lz->indices[i] = index;
2245 }
2246
2247 result = PyTuple_New(n);
2248 if (!result)
2249 return NULL;
2250 for (i=0; i<n; i++) {
2251 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2252 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2253 Py_INCREF(element);
2254 PyTuple_SET_ITEM(result, i, element);
2255 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002256 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002257 Py_RETURN_NONE;
2258}
2259
2260static PyMethodDef product_methods[] = {
2261 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2262 reduce_doc},
2263 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2264 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002265 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2266 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002267 {NULL, NULL} /* sentinel */
2268};
2269
Christian Heimesc3f30c42008-02-22 16:37:40 +00002270PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002271"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002272\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002273Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002274For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2275The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2276cycle in a manner similar to an odometer (with the rightmost element changing\n\
2277on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002278To compute the product of an iterable with itself, specify the number\n\
2279of repetitions with the optional repeat keyword argument. For example,\n\
2280product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002281product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2282product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2283
2284static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002285 PyVarObject_HEAD_INIT(NULL, 0)
2286 "itertools.product", /* tp_name */
2287 sizeof(productobject), /* tp_basicsize */
2288 0, /* tp_itemsize */
2289 /* methods */
2290 (destructor)product_dealloc, /* tp_dealloc */
2291 0, /* tp_print */
2292 0, /* tp_getattr */
2293 0, /* tp_setattr */
2294 0, /* tp_reserved */
2295 0, /* tp_repr */
2296 0, /* tp_as_number */
2297 0, /* tp_as_sequence */
2298 0, /* tp_as_mapping */
2299 0, /* tp_hash */
2300 0, /* tp_call */
2301 0, /* tp_str */
2302 PyObject_GenericGetAttr, /* tp_getattro */
2303 0, /* tp_setattro */
2304 0, /* tp_as_buffer */
2305 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2306 Py_TPFLAGS_BASETYPE, /* tp_flags */
2307 product_doc, /* tp_doc */
2308 (traverseproc)product_traverse, /* tp_traverse */
2309 0, /* tp_clear */
2310 0, /* tp_richcompare */
2311 0, /* tp_weaklistoffset */
2312 PyObject_SelfIter, /* tp_iter */
2313 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002314 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002315 0, /* tp_members */
2316 0, /* tp_getset */
2317 0, /* tp_base */
2318 0, /* tp_dict */
2319 0, /* tp_descr_get */
2320 0, /* tp_descr_set */
2321 0, /* tp_dictoffset */
2322 0, /* tp_init */
2323 0, /* tp_alloc */
2324 product_new, /* tp_new */
2325 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002326};
2327
2328
Christian Heimes380f7f22008-02-28 11:19:05 +00002329/* combinations object ************************************************************/
2330
2331typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002332 PyObject_HEAD
2333 PyObject *pool; /* input converted to a tuple */
2334 Py_ssize_t *indices; /* one index per result element */
2335 PyObject *result; /* most recently returned result tuple */
2336 Py_ssize_t r; /* size of result tuple */
2337 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002338} combinationsobject;
2339
2340static PyTypeObject combinations_type;
2341
2342static PyObject *
2343combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2344{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002345 combinationsobject *co;
2346 Py_ssize_t n;
2347 Py_ssize_t r;
2348 PyObject *pool = NULL;
2349 PyObject *iterable = NULL;
2350 Py_ssize_t *indices = NULL;
2351 Py_ssize_t i;
2352 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002353
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002354 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2355 &iterable, &r))
2356 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002357
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002358 pool = PySequence_Tuple(iterable);
2359 if (pool == NULL)
2360 goto error;
2361 n = PyTuple_GET_SIZE(pool);
2362 if (r < 0) {
2363 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2364 goto error;
2365 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002366
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002367 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002368 if (indices == NULL) {
2369 PyErr_NoMemory();
2370 goto error;
2371 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002372
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002373 for (i=0 ; i<r ; i++)
2374 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002375
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002376 /* create combinationsobject structure */
2377 co = (combinationsobject *)type->tp_alloc(type, 0);
2378 if (co == NULL)
2379 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002380
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002381 co->pool = pool;
2382 co->indices = indices;
2383 co->result = NULL;
2384 co->r = r;
2385 co->stopped = r > n ? 1 : 0;
2386
2387 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002388
2389error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002390 if (indices != NULL)
2391 PyMem_Free(indices);
2392 Py_XDECREF(pool);
2393 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002394}
2395
2396static void
2397combinations_dealloc(combinationsobject *co)
2398{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002399 PyObject_GC_UnTrack(co);
2400 Py_XDECREF(co->pool);
2401 Py_XDECREF(co->result);
2402 if (co->indices != NULL)
2403 PyMem_Free(co->indices);
2404 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002405}
2406
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002407static PyObject *
2408combinations_sizeof(combinationsobject *co, void *unused)
2409{
2410 Py_ssize_t res;
2411
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002412 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002413 res += co->r * sizeof(Py_ssize_t);
2414 return PyLong_FromSsize_t(res);
2415}
2416
Christian Heimes380f7f22008-02-28 11:19:05 +00002417static int
2418combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2419{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002420 Py_VISIT(co->pool);
2421 Py_VISIT(co->result);
2422 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002423}
2424
2425static PyObject *
2426combinations_next(combinationsobject *co)
2427{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002428 PyObject *elem;
2429 PyObject *oldelem;
2430 PyObject *pool = co->pool;
2431 Py_ssize_t *indices = co->indices;
2432 PyObject *result = co->result;
2433 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2434 Py_ssize_t r = co->r;
2435 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002436
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002437 if (co->stopped)
2438 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002439
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002440 if (result == NULL) {
2441 /* On the first pass, initialize result tuple using the indices */
2442 result = PyTuple_New(r);
2443 if (result == NULL)
2444 goto empty;
2445 co->result = result;
2446 for (i=0; i<r ; i++) {
2447 index = indices[i];
2448 elem = PyTuple_GET_ITEM(pool, index);
2449 Py_INCREF(elem);
2450 PyTuple_SET_ITEM(result, i, elem);
2451 }
2452 } else {
2453 /* Copy the previous result tuple or re-use it if available */
2454 if (Py_REFCNT(result) > 1) {
2455 PyObject *old_result = result;
2456 result = PyTuple_New(r);
2457 if (result == NULL)
2458 goto empty;
2459 co->result = result;
2460 for (i=0; i<r ; i++) {
2461 elem = PyTuple_GET_ITEM(old_result, i);
2462 Py_INCREF(elem);
2463 PyTuple_SET_ITEM(result, i, elem);
2464 }
2465 Py_DECREF(old_result);
2466 }
2467 /* Now, we've got the only copy so we can update it in-place
2468 * CPython's empty tuple is a singleton and cached in
2469 * PyTuple's freelist.
2470 */
2471 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002472
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002473 /* Scan indices right-to-left until finding one that is not
2474 at its maximum (i + n - r). */
2475 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2476 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002477
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002478 /* If i is negative, then the indices are all at
2479 their maximum value and we're done. */
2480 if (i < 0)
2481 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002482
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002483 /* Increment the current index which we know is not at its
2484 maximum. Then move back to the right setting each index
2485 to its lowest possible value (one higher than the index
2486 to its left -- this maintains the sort order invariant). */
2487 indices[i]++;
2488 for (j=i+1 ; j<r ; j++)
2489 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002490
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002491 /* Update the result tuple for the new indices
2492 starting with i, the leftmost index that changed */
2493 for ( ; i<r ; i++) {
2494 index = indices[i];
2495 elem = PyTuple_GET_ITEM(pool, index);
2496 Py_INCREF(elem);
2497 oldelem = PyTuple_GET_ITEM(result, i);
2498 PyTuple_SET_ITEM(result, i, elem);
2499 Py_DECREF(oldelem);
2500 }
2501 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002502
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002503 Py_INCREF(result);
2504 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002505
2506empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002507 co->stopped = 1;
2508 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002509}
2510
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002511static PyObject *
2512combinations_reduce(combinationsobject *lz)
2513{
2514 if (lz->result == NULL) {
2515 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2516 } else if (lz->stopped) {
2517 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2518 } else {
2519 PyObject *indices;
2520 Py_ssize_t i;
2521
2522 /* we must pickle the indices and use them for setstate */
2523 indices = PyTuple_New(lz->r);
2524 if (!indices)
2525 return NULL;
2526 for (i=0; i<lz->r; i++)
2527 {
2528 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2529 if (!index) {
2530 Py_DECREF(indices);
2531 return NULL;
2532 }
2533 PyTuple_SET_ITEM(indices, i, index);
2534 }
2535
2536 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2537 }
2538}
2539
2540static PyObject *
2541combinations_setstate(combinationsobject *lz, PyObject *state)
2542{
2543 PyObject *result;
2544 Py_ssize_t i;
2545 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2546
2547 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2548 {
2549 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2550 return NULL;
2551 }
2552
2553 for (i=0; i<lz->r; i++)
2554 {
2555 Py_ssize_t max;
2556 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2557 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2558 if (index == -1 && PyErr_Occurred())
2559 return NULL; /* not an integer */
2560 max = i + n - lz->r;
2561 /* clamp the index (beware of negative max) */
2562 if (index > max)
2563 index = max;
2564 if (index < 0)
2565 index = 0;
2566 lz->indices[i] = index;
2567 }
2568
2569 result = PyTuple_New(lz->r);
2570 if (result == NULL)
2571 return NULL;
2572 for (i=0; i<lz->r; i++) {
2573 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2574 Py_INCREF(element);
2575 PyTuple_SET_ITEM(result, i, element);
2576 }
2577
Serhiy Storchaka48842712016-04-06 09:45:48 +03002578 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002579 Py_RETURN_NONE;
2580}
2581
2582static PyMethodDef combinations_methods[] = {
2583 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2584 reduce_doc},
2585 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2586 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002587 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2588 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002589 {NULL, NULL} /* sentinel */
2590};
2591
Christian Heimes380f7f22008-02-28 11:19:05 +00002592PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002593"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002594\n\
2595Return successive r-length combinations of elements in the iterable.\n\n\
2596combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2597
2598static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002599 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002600 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002601 sizeof(combinationsobject), /* tp_basicsize */
2602 0, /* tp_itemsize */
2603 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002604 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002605 0, /* tp_print */
2606 0, /* tp_getattr */
2607 0, /* tp_setattr */
2608 0, /* tp_reserved */
2609 0, /* tp_repr */
2610 0, /* tp_as_number */
2611 0, /* tp_as_sequence */
2612 0, /* tp_as_mapping */
2613 0, /* tp_hash */
2614 0, /* tp_call */
2615 0, /* tp_str */
2616 PyObject_GenericGetAttr, /* tp_getattro */
2617 0, /* tp_setattro */
2618 0, /* tp_as_buffer */
2619 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2620 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002621 combinations_doc, /* tp_doc */
2622 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002623 0, /* tp_clear */
2624 0, /* tp_richcompare */
2625 0, /* tp_weaklistoffset */
2626 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002627 (iternextfunc)combinations_next, /* tp_iternext */
2628 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002629 0, /* tp_members */
2630 0, /* tp_getset */
2631 0, /* tp_base */
2632 0, /* tp_dict */
2633 0, /* tp_descr_get */
2634 0, /* tp_descr_set */
2635 0, /* tp_dictoffset */
2636 0, /* tp_init */
2637 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002638 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002639 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002640};
2641
2642
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002643/* combinations with replacement object *******************************************/
2644
2645/* Equivalent to:
2646
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002647 def combinations_with_replacement(iterable, r):
2648 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2649 # number items returned: (n+r-1)! / r! / (n-1)!
2650 pool = tuple(iterable)
2651 n = len(pool)
2652 indices = [0] * r
2653 yield tuple(pool[i] for i in indices)
2654 while 1:
2655 for i in reversed(range(r)):
2656 if indices[i] != n - 1:
2657 break
2658 else:
2659 return
2660 indices[i:] = [indices[i] + 1] * (r - i)
2661 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002662
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002663 def combinations_with_replacement2(iterable, r):
2664 'Alternate version that filters from product()'
2665 pool = tuple(iterable)
2666 n = len(pool)
2667 for indices in product(range(n), repeat=r):
2668 if sorted(indices) == list(indices):
2669 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002670*/
2671typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002672 PyObject_HEAD
2673 PyObject *pool; /* input converted to a tuple */
2674 Py_ssize_t *indices; /* one index per result element */
2675 PyObject *result; /* most recently returned result tuple */
2676 Py_ssize_t r; /* size of result tuple */
2677 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002678} cwrobject;
2679
2680static PyTypeObject cwr_type;
2681
2682static PyObject *
2683cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2684{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002685 cwrobject *co;
2686 Py_ssize_t n;
2687 Py_ssize_t r;
2688 PyObject *pool = NULL;
2689 PyObject *iterable = NULL;
2690 Py_ssize_t *indices = NULL;
2691 Py_ssize_t i;
2692 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002693
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002694 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2695 &iterable, &r))
2696 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002697
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002698 pool = PySequence_Tuple(iterable);
2699 if (pool == NULL)
2700 goto error;
2701 n = PyTuple_GET_SIZE(pool);
2702 if (r < 0) {
2703 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2704 goto error;
2705 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002706
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002707 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002708 if (indices == NULL) {
2709 PyErr_NoMemory();
2710 goto error;
2711 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002712
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002713 for (i=0 ; i<r ; i++)
2714 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002715
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002716 /* create cwrobject structure */
2717 co = (cwrobject *)type->tp_alloc(type, 0);
2718 if (co == NULL)
2719 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002720
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002721 co->pool = pool;
2722 co->indices = indices;
2723 co->result = NULL;
2724 co->r = r;
2725 co->stopped = !n && r;
2726
2727 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002728
2729error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002730 if (indices != NULL)
2731 PyMem_Free(indices);
2732 Py_XDECREF(pool);
2733 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002734}
2735
2736static void
2737cwr_dealloc(cwrobject *co)
2738{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002739 PyObject_GC_UnTrack(co);
2740 Py_XDECREF(co->pool);
2741 Py_XDECREF(co->result);
2742 if (co->indices != NULL)
2743 PyMem_Free(co->indices);
2744 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002745}
2746
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002747static PyObject *
2748cwr_sizeof(cwrobject *co, void *unused)
2749{
2750 Py_ssize_t res;
2751
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002752 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002753 res += co->r * sizeof(Py_ssize_t);
2754 return PyLong_FromSsize_t(res);
2755}
2756
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002757static int
2758cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2759{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002760 Py_VISIT(co->pool);
2761 Py_VISIT(co->result);
2762 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002763}
2764
2765static PyObject *
2766cwr_next(cwrobject *co)
2767{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002768 PyObject *elem;
2769 PyObject *oldelem;
2770 PyObject *pool = co->pool;
2771 Py_ssize_t *indices = co->indices;
2772 PyObject *result = co->result;
2773 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2774 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002775 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002776
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002777 if (co->stopped)
2778 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002779
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002780 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002781 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002782 result = PyTuple_New(r);
2783 if (result == NULL)
2784 goto empty;
2785 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002786 if (n > 0) {
2787 elem = PyTuple_GET_ITEM(pool, 0);
2788 for (i=0; i<r ; i++) {
2789 assert(indices[i] == 0);
2790 Py_INCREF(elem);
2791 PyTuple_SET_ITEM(result, i, elem);
2792 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002793 }
2794 } else {
2795 /* Copy the previous result tuple or re-use it if available */
2796 if (Py_REFCNT(result) > 1) {
2797 PyObject *old_result = result;
2798 result = PyTuple_New(r);
2799 if (result == NULL)
2800 goto empty;
2801 co->result = result;
2802 for (i=0; i<r ; i++) {
2803 elem = PyTuple_GET_ITEM(old_result, i);
2804 Py_INCREF(elem);
2805 PyTuple_SET_ITEM(result, i, elem);
2806 }
2807 Py_DECREF(old_result);
2808 }
2809 /* Now, we've got the only copy so we can update it in-place CPython's
2810 empty tuple is a singleton and cached in PyTuple's freelist. */
2811 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002812
Tim Peters9edb1682013-09-03 11:49:31 -05002813 /* Scan indices right-to-left until finding one that is not
2814 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002815 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2816 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002817
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002818 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002819 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002820 if (i < 0)
2821 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002822
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002823 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002824 maximum. Then set all to the right to the same value. */
2825 index = indices[i] + 1;
2826 assert(index < n);
2827 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002828 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002829 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002830 Py_INCREF(elem);
2831 oldelem = PyTuple_GET_ITEM(result, i);
2832 PyTuple_SET_ITEM(result, i, elem);
2833 Py_DECREF(oldelem);
2834 }
2835 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002836
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002837 Py_INCREF(result);
2838 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002839
2840empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002841 co->stopped = 1;
2842 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002843}
2844
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002845static PyObject *
2846cwr_reduce(cwrobject *lz)
2847{
2848 if (lz->result == NULL) {
2849 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2850 } else if (lz->stopped) {
2851 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2852 } else {
2853 PyObject *indices;
2854 Py_ssize_t i;
2855
2856 /* we must pickle the indices and use them for setstate */
2857 indices = PyTuple_New(lz->r);
2858 if (!indices)
2859 return NULL;
2860 for (i=0; i<lz->r; i++)
2861 {
2862 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2863 if (!index) {
2864 Py_DECREF(indices);
2865 return NULL;
2866 }
2867 PyTuple_SET_ITEM(indices, i, index);
2868 }
2869
2870 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2871 }
2872}
2873
2874static PyObject *
2875cwr_setstate(cwrobject *lz, PyObject *state)
2876{
2877 PyObject *result;
2878 Py_ssize_t n, i;
2879
2880 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2881 {
2882 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2883 return NULL;
2884 }
2885
2886 n = PyTuple_GET_SIZE(lz->pool);
2887 for (i=0; i<lz->r; i++)
2888 {
2889 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2890 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2891 if (index < 0 && PyErr_Occurred())
2892 return NULL; /* not an integer */
2893 /* clamp the index */
2894 if (index < 0)
2895 index = 0;
2896 else if (index > n-1)
2897 index = n-1;
2898 lz->indices[i] = index;
2899 }
2900 result = PyTuple_New(lz->r);
2901 if (result == NULL)
2902 return NULL;
2903 for (i=0; i<lz->r; i++) {
2904 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2905 Py_INCREF(element);
2906 PyTuple_SET_ITEM(result, i, element);
2907 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002908 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002909 Py_RETURN_NONE;
2910}
2911
2912static PyMethodDef cwr_methods[] = {
2913 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2914 reduce_doc},
2915 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2916 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002917 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
2918 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002919 {NULL, NULL} /* sentinel */
2920};
2921
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002922PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002923"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002924\n\
2925Return successive r-length combinations of elements in the iterable\n\
2926allowing individual elements to have successive repeats.\n\
2927combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2928
2929static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002930 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002931 "itertools.combinations_with_replacement", /* tp_name */
2932 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002933 0, /* tp_itemsize */
2934 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002935 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002936 0, /* tp_print */
2937 0, /* tp_getattr */
2938 0, /* tp_setattr */
2939 0, /* tp_reserved */
2940 0, /* tp_repr */
2941 0, /* tp_as_number */
2942 0, /* tp_as_sequence */
2943 0, /* tp_as_mapping */
2944 0, /* tp_hash */
2945 0, /* tp_call */
2946 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002947 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002948 0, /* tp_setattro */
2949 0, /* tp_as_buffer */
2950 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002951 Py_TPFLAGS_BASETYPE, /* tp_flags */
2952 cwr_doc, /* tp_doc */
2953 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002954 0, /* tp_clear */
2955 0, /* tp_richcompare */
2956 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002957 PyObject_SelfIter, /* tp_iter */
2958 (iternextfunc)cwr_next, /* tp_iternext */
2959 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002960 0, /* tp_members */
2961 0, /* tp_getset */
2962 0, /* tp_base */
2963 0, /* tp_dict */
2964 0, /* tp_descr_get */
2965 0, /* tp_descr_set */
2966 0, /* tp_dictoffset */
2967 0, /* tp_init */
2968 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002969 cwr_new, /* tp_new */
2970 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002971};
2972
2973
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002974/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002975
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002976def permutations(iterable, r=None):
2977 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2978 pool = tuple(iterable)
2979 n = len(pool)
2980 r = n if r is None else r
2981 indices = range(n)
2982 cycles = range(n-r+1, n+1)[::-1]
2983 yield tuple(pool[i] for i in indices[:r])
2984 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03002985 for i in reversed(range(r)):
2986 cycles[i] -= 1
2987 if cycles[i] == 0:
2988 indices[i:] = indices[i+1:] + indices[i:i+1]
2989 cycles[i] = n - i
2990 else:
2991 j = cycles[i]
2992 indices[i], indices[-j] = indices[-j], indices[i]
2993 yield tuple(pool[i] for i in indices[:r])
2994 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002995 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03002996 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002997*/
2998
2999typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003000 PyObject_HEAD
3001 PyObject *pool; /* input converted to a tuple */
3002 Py_ssize_t *indices; /* one index per element in the pool */
3003 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3004 PyObject *result; /* most recently returned result tuple */
3005 Py_ssize_t r; /* size of result tuple */
3006 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003007} permutationsobject;
3008
3009static PyTypeObject permutations_type;
3010
3011static PyObject *
3012permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3013{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003014 permutationsobject *po;
3015 Py_ssize_t n;
3016 Py_ssize_t r;
3017 PyObject *robj = Py_None;
3018 PyObject *pool = NULL;
3019 PyObject *iterable = NULL;
3020 Py_ssize_t *indices = NULL;
3021 Py_ssize_t *cycles = NULL;
3022 Py_ssize_t i;
3023 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003024
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003025 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
3026 &iterable, &robj))
3027 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003028
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003029 pool = PySequence_Tuple(iterable);
3030 if (pool == NULL)
3031 goto error;
3032 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003033
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003034 r = n;
3035 if (robj != Py_None) {
3036 if (!PyLong_Check(robj)) {
3037 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3038 goto error;
3039 }
3040 r = PyLong_AsSsize_t(robj);
3041 if (r == -1 && PyErr_Occurred())
3042 goto error;
3043 }
3044 if (r < 0) {
3045 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3046 goto error;
3047 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003048
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003049 indices = PyMem_New(Py_ssize_t, n);
3050 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003051 if (indices == NULL || cycles == NULL) {
3052 PyErr_NoMemory();
3053 goto error;
3054 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003055
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003056 for (i=0 ; i<n ; i++)
3057 indices[i] = i;
3058 for (i=0 ; i<r ; i++)
3059 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003060
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003061 /* create permutationsobject structure */
3062 po = (permutationsobject *)type->tp_alloc(type, 0);
3063 if (po == NULL)
3064 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003065
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003066 po->pool = pool;
3067 po->indices = indices;
3068 po->cycles = cycles;
3069 po->result = NULL;
3070 po->r = r;
3071 po->stopped = r > n ? 1 : 0;
3072
3073 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003074
3075error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003076 if (indices != NULL)
3077 PyMem_Free(indices);
3078 if (cycles != NULL)
3079 PyMem_Free(cycles);
3080 Py_XDECREF(pool);
3081 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003082}
3083
3084static void
3085permutations_dealloc(permutationsobject *po)
3086{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003087 PyObject_GC_UnTrack(po);
3088 Py_XDECREF(po->pool);
3089 Py_XDECREF(po->result);
3090 PyMem_Free(po->indices);
3091 PyMem_Free(po->cycles);
3092 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003093}
3094
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003095static PyObject *
3096permutations_sizeof(permutationsobject *po, void *unused)
3097{
3098 Py_ssize_t res;
3099
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003100 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003101 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3102 res += po->r * sizeof(Py_ssize_t);
3103 return PyLong_FromSsize_t(res);
3104}
3105
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003106static int
3107permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3108{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003109 Py_VISIT(po->pool);
3110 Py_VISIT(po->result);
3111 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003112}
3113
3114static PyObject *
3115permutations_next(permutationsobject *po)
3116{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003117 PyObject *elem;
3118 PyObject *oldelem;
3119 PyObject *pool = po->pool;
3120 Py_ssize_t *indices = po->indices;
3121 Py_ssize_t *cycles = po->cycles;
3122 PyObject *result = po->result;
3123 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3124 Py_ssize_t r = po->r;
3125 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003126
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003127 if (po->stopped)
3128 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003129
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003130 if (result == NULL) {
3131 /* On the first pass, initialize result tuple using the indices */
3132 result = PyTuple_New(r);
3133 if (result == NULL)
3134 goto empty;
3135 po->result = result;
3136 for (i=0; i<r ; i++) {
3137 index = indices[i];
3138 elem = PyTuple_GET_ITEM(pool, index);
3139 Py_INCREF(elem);
3140 PyTuple_SET_ITEM(result, i, elem);
3141 }
3142 } else {
3143 if (n == 0)
3144 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003145
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003146 /* Copy the previous result tuple or re-use it if available */
3147 if (Py_REFCNT(result) > 1) {
3148 PyObject *old_result = result;
3149 result = PyTuple_New(r);
3150 if (result == NULL)
3151 goto empty;
3152 po->result = result;
3153 for (i=0; i<r ; i++) {
3154 elem = PyTuple_GET_ITEM(old_result, i);
3155 Py_INCREF(elem);
3156 PyTuple_SET_ITEM(result, i, elem);
3157 }
3158 Py_DECREF(old_result);
3159 }
3160 /* Now, we've got the only copy so we can update it in-place */
3161 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003162
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003163 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3164 for (i=r-1 ; i>=0 ; i--) {
3165 cycles[i] -= 1;
3166 if (cycles[i] == 0) {
3167 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3168 index = indices[i];
3169 for (j=i ; j<n-1 ; j++)
3170 indices[j] = indices[j+1];
3171 indices[n-1] = index;
3172 cycles[i] = n - i;
3173 } else {
3174 j = cycles[i];
3175 index = indices[i];
3176 indices[i] = indices[n-j];
3177 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003178
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003179 for (k=i; k<r ; k++) {
3180 /* start with i, the leftmost element that changed */
3181 /* yield tuple(pool[k] for k in indices[:r]) */
3182 index = indices[k];
3183 elem = PyTuple_GET_ITEM(pool, index);
3184 Py_INCREF(elem);
3185 oldelem = PyTuple_GET_ITEM(result, k);
3186 PyTuple_SET_ITEM(result, k, elem);
3187 Py_DECREF(oldelem);
3188 }
3189 break;
3190 }
3191 }
3192 /* If i is negative, then the cycles have all
3193 rolled-over and we're done. */
3194 if (i < 0)
3195 goto empty;
3196 }
3197 Py_INCREF(result);
3198 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003199
3200empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003201 po->stopped = 1;
3202 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003203}
3204
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003205static PyObject *
3206permutations_reduce(permutationsobject *po)
3207{
3208 if (po->result == NULL) {
3209 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3210 } else if (po->stopped) {
3211 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3212 } else {
3213 PyObject *indices=NULL, *cycles=NULL;
3214 Py_ssize_t n, i;
3215
3216 /* we must pickle the indices and cycles and use them for setstate */
3217 n = PyTuple_GET_SIZE(po->pool);
3218 indices = PyTuple_New(n);
3219 if (indices == NULL)
3220 goto err;
3221 for (i=0; i<n; i++){
3222 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3223 if (!index)
3224 goto err;
3225 PyTuple_SET_ITEM(indices, i, index);
3226 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003227
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003228 cycles = PyTuple_New(po->r);
3229 if (cycles == NULL)
3230 goto err;
3231 for (i=0; i<po->r; i++)
3232 {
3233 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3234 if (!index)
3235 goto err;
3236 PyTuple_SET_ITEM(cycles, i, index);
3237 }
3238 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3239 po->pool, po->r,
3240 indices, cycles);
3241 err:
3242 Py_XDECREF(indices);
3243 Py_XDECREF(cycles);
3244 return NULL;
3245 }
3246}
3247
3248static PyObject *
3249permutations_setstate(permutationsobject *po, PyObject *state)
3250{
3251 PyObject *indices, *cycles, *result;
3252 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003253
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003254 if (!PyArg_ParseTuple(state, "O!O!",
3255 &PyTuple_Type, &indices,
3256 &PyTuple_Type, &cycles))
3257 return NULL;
3258
3259 n = PyTuple_GET_SIZE(po->pool);
3260 if (PyTuple_GET_SIZE(indices) != n ||
3261 PyTuple_GET_SIZE(cycles) != po->r)
3262 {
3263 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3264 return NULL;
3265 }
3266
3267 for (i=0; i<n; i++)
3268 {
3269 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3270 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3271 if (index < 0 && PyErr_Occurred())
3272 return NULL; /* not an integer */
3273 /* clamp the index */
3274 if (index < 0)
3275 index = 0;
3276 else if (index > n-1)
3277 index = n-1;
3278 po->indices[i] = index;
3279 }
3280
3281 for (i=0; i<po->r; i++)
3282 {
3283 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3284 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3285 if (index < 0 && PyErr_Occurred())
3286 return NULL; /* not an integer */
3287 if (index < 1)
3288 index = 1;
3289 else if (index > n-i)
3290 index = n-i;
3291 po->cycles[i] = index;
3292 }
3293 result = PyTuple_New(po->r);
3294 if (result == NULL)
3295 return NULL;
3296 for (i=0; i<po->r; i++) {
3297 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3298 Py_INCREF(element);
3299 PyTuple_SET_ITEM(result, i, element);
3300 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003301 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003302 Py_RETURN_NONE;
3303}
3304
3305static PyMethodDef permuations_methods[] = {
3306 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3307 reduce_doc},
3308 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3309 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003310 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3311 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003312 {NULL, NULL} /* sentinel */
3313};
3314
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003315PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003316"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003317\n\
3318Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003319permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003320
3321static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003322 PyVarObject_HEAD_INIT(NULL, 0)
3323 "itertools.permutations", /* tp_name */
3324 sizeof(permutationsobject), /* tp_basicsize */
3325 0, /* tp_itemsize */
3326 /* methods */
3327 (destructor)permutations_dealloc, /* tp_dealloc */
3328 0, /* tp_print */
3329 0, /* tp_getattr */
3330 0, /* tp_setattr */
3331 0, /* tp_reserved */
3332 0, /* tp_repr */
3333 0, /* tp_as_number */
3334 0, /* tp_as_sequence */
3335 0, /* tp_as_mapping */
3336 0, /* tp_hash */
3337 0, /* tp_call */
3338 0, /* tp_str */
3339 PyObject_GenericGetAttr, /* tp_getattro */
3340 0, /* tp_setattro */
3341 0, /* tp_as_buffer */
3342 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3343 Py_TPFLAGS_BASETYPE, /* tp_flags */
3344 permutations_doc, /* tp_doc */
3345 (traverseproc)permutations_traverse, /* tp_traverse */
3346 0, /* tp_clear */
3347 0, /* tp_richcompare */
3348 0, /* tp_weaklistoffset */
3349 PyObject_SelfIter, /* tp_iter */
3350 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003351 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003352 0, /* tp_members */
3353 0, /* tp_getset */
3354 0, /* tp_base */
3355 0, /* tp_dict */
3356 0, /* tp_descr_get */
3357 0, /* tp_descr_set */
3358 0, /* tp_dictoffset */
3359 0, /* tp_init */
3360 0, /* tp_alloc */
3361 permutations_new, /* tp_new */
3362 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003363};
3364
Raymond Hettinger482ba772010-12-01 22:48:00 +00003365/* accumulate object ************************************************************/
3366
3367typedef struct {
3368 PyObject_HEAD
3369 PyObject *total;
3370 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003371 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003372} accumulateobject;
3373
3374static PyTypeObject accumulate_type;
3375
3376static PyObject *
3377accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3378{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003379 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003380 PyObject *iterable;
3381 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003382 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003383 accumulateobject *lz;
3384
Raymond Hettinger5d446132011-03-27 18:52:10 -07003385 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3386 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003387 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003388
3389 /* Get iterator. */
3390 it = PyObject_GetIter(iterable);
3391 if (it == NULL)
3392 return NULL;
3393
Raymond Hettinger482ba772010-12-01 22:48:00 +00003394 /* create accumulateobject structure */
3395 lz = (accumulateobject *)type->tp_alloc(type, 0);
3396 if (lz == NULL) {
3397 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003398 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003399 }
3400
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003401 if (binop != Py_None) {
3402 Py_XINCREF(binop);
3403 lz->binop = binop;
3404 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003405 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003406 lz->it = it;
3407 return (PyObject *)lz;
3408}
3409
3410static void
3411accumulate_dealloc(accumulateobject *lz)
3412{
3413 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003414 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003415 Py_XDECREF(lz->total);
3416 Py_XDECREF(lz->it);
3417 Py_TYPE(lz)->tp_free(lz);
3418}
3419
3420static int
3421accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3422{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003423 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003424 Py_VISIT(lz->it);
3425 Py_VISIT(lz->total);
3426 return 0;
3427}
3428
3429static PyObject *
3430accumulate_next(accumulateobject *lz)
3431{
3432 PyObject *val, *oldtotal, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003433
Raymond Hettinger482ba772010-12-01 22:48:00 +00003434 val = PyIter_Next(lz->it);
3435 if (val == NULL)
3436 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003437
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003438 if (lz->total == NULL) {
3439 Py_INCREF(val);
3440 lz->total = val;
3441 return lz->total;
3442 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003443
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003444 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003445 newtotal = PyNumber_Add(lz->total, val);
3446 else
3447 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003448 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003449 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003450 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003451
3452 oldtotal = lz->total;
3453 lz->total = newtotal;
3454 Py_DECREF(oldtotal);
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003455
Raymond Hettinger482ba772010-12-01 22:48:00 +00003456 Py_INCREF(newtotal);
3457 return newtotal;
3458}
3459
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003460static PyObject *
3461accumulate_reduce(accumulateobject *lz)
3462{
Serhiy Storchakad5516252016-03-06 14:00:45 +02003463 if (lz->total == Py_None) {
3464 PyObject *it;
3465
3466 if (PyType_Ready(&chain_type) < 0)
3467 return NULL;
3468 if (PyType_Ready(&islice_type) < 0)
3469 return NULL;
3470 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3471 lz->total, lz->it);
3472 if (it == NULL)
3473 return NULL;
3474 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3475 it, lz->binop ? lz->binop : Py_None);
3476 if (it == NULL)
3477 return NULL;
3478 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3479 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003480 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3481 lz->it, lz->binop?lz->binop:Py_None,
3482 lz->total?lz->total:Py_None);
3483 }
3484
3485static PyObject *
3486accumulate_setstate(accumulateobject *lz, PyObject *state)
3487{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003488 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003489 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003490 Py_RETURN_NONE;
3491}
3492
3493static PyMethodDef accumulate_methods[] = {
3494 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3495 reduce_doc},
3496 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3497 setstate_doc},
3498 {NULL, NULL} /* sentinel */
3499};
3500
Raymond Hettinger482ba772010-12-01 22:48:00 +00003501PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003502"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003503\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003504Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003505
3506static PyTypeObject accumulate_type = {
3507 PyVarObject_HEAD_INIT(NULL, 0)
3508 "itertools.accumulate", /* tp_name */
3509 sizeof(accumulateobject), /* tp_basicsize */
3510 0, /* tp_itemsize */
3511 /* methods */
3512 (destructor)accumulate_dealloc, /* tp_dealloc */
3513 0, /* tp_print */
3514 0, /* tp_getattr */
3515 0, /* tp_setattr */
3516 0, /* tp_reserved */
3517 0, /* tp_repr */
3518 0, /* tp_as_number */
3519 0, /* tp_as_sequence */
3520 0, /* tp_as_mapping */
3521 0, /* tp_hash */
3522 0, /* tp_call */
3523 0, /* tp_str */
3524 PyObject_GenericGetAttr, /* tp_getattro */
3525 0, /* tp_setattro */
3526 0, /* tp_as_buffer */
3527 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3528 Py_TPFLAGS_BASETYPE, /* tp_flags */
3529 accumulate_doc, /* tp_doc */
3530 (traverseproc)accumulate_traverse, /* tp_traverse */
3531 0, /* tp_clear */
3532 0, /* tp_richcompare */
3533 0, /* tp_weaklistoffset */
3534 PyObject_SelfIter, /* tp_iter */
3535 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003536 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003537 0, /* tp_members */
3538 0, /* tp_getset */
3539 0, /* tp_base */
3540 0, /* tp_dict */
3541 0, /* tp_descr_get */
3542 0, /* tp_descr_set */
3543 0, /* tp_dictoffset */
3544 0, /* tp_init */
3545 0, /* tp_alloc */
3546 accumulate_new, /* tp_new */
3547 PyObject_GC_Del, /* tp_free */
3548};
3549
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003550
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003551/* compress object ************************************************************/
3552
3553/* Equivalent to:
3554
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003555 def compress(data, selectors):
3556 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3557 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003558*/
3559
3560typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003561 PyObject_HEAD
3562 PyObject *data;
3563 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003564} compressobject;
3565
3566static PyTypeObject compress_type;
3567
3568static PyObject *
3569compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3570{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003571 PyObject *seq1, *seq2;
3572 PyObject *data=NULL, *selectors=NULL;
3573 compressobject *lz;
3574 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003575
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003576 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3577 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003578
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003579 data = PyObject_GetIter(seq1);
3580 if (data == NULL)
3581 goto fail;
3582 selectors = PyObject_GetIter(seq2);
3583 if (selectors == NULL)
3584 goto fail;
3585
3586 /* create compressobject structure */
3587 lz = (compressobject *)type->tp_alloc(type, 0);
3588 if (lz == NULL)
3589 goto fail;
3590 lz->data = data;
3591 lz->selectors = selectors;
3592 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003593
3594fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003595 Py_XDECREF(data);
3596 Py_XDECREF(selectors);
3597 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003598}
3599
3600static void
3601compress_dealloc(compressobject *lz)
3602{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003603 PyObject_GC_UnTrack(lz);
3604 Py_XDECREF(lz->data);
3605 Py_XDECREF(lz->selectors);
3606 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003607}
3608
3609static int
3610compress_traverse(compressobject *lz, visitproc visit, void *arg)
3611{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003612 Py_VISIT(lz->data);
3613 Py_VISIT(lz->selectors);
3614 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003615}
3616
3617static PyObject *
3618compress_next(compressobject *lz)
3619{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003620 PyObject *data = lz->data, *selectors = lz->selectors;
3621 PyObject *datum, *selector;
3622 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3623 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3624 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003625
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003626 while (1) {
3627 /* Steps: get datum, get selector, evaluate selector.
3628 Order is important (to match the pure python version
3629 in terms of which input gets a chance to raise an
3630 exception first).
3631 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003632
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003633 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003634 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003635 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003636
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003637 selector = selectornext(selectors);
3638 if (selector == NULL) {
3639 Py_DECREF(datum);
3640 return NULL;
3641 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003642
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003643 ok = PyObject_IsTrue(selector);
3644 Py_DECREF(selector);
3645 if (ok == 1)
3646 return datum;
3647 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003648 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003649 return NULL;
3650 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003651}
3652
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003653static PyObject *
3654compress_reduce(compressobject *lz)
3655{
3656 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3657 lz->data, lz->selectors);
3658 }
3659
3660static PyMethodDef compress_methods[] = {
3661 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3662 reduce_doc},
3663 {NULL, NULL} /* sentinel */
3664};
3665
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003666PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003667"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003668\n\
3669Return data elements corresponding to true selector elements.\n\
3670Forms a shorter iterator from selected data elements using the\n\
3671selectors to choose the data elements.");
3672
3673static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003674 PyVarObject_HEAD_INIT(NULL, 0)
3675 "itertools.compress", /* tp_name */
3676 sizeof(compressobject), /* tp_basicsize */
3677 0, /* tp_itemsize */
3678 /* methods */
3679 (destructor)compress_dealloc, /* tp_dealloc */
3680 0, /* tp_print */
3681 0, /* tp_getattr */
3682 0, /* tp_setattr */
3683 0, /* tp_reserved */
3684 0, /* tp_repr */
3685 0, /* tp_as_number */
3686 0, /* tp_as_sequence */
3687 0, /* tp_as_mapping */
3688 0, /* tp_hash */
3689 0, /* tp_call */
3690 0, /* tp_str */
3691 PyObject_GenericGetAttr, /* tp_getattro */
3692 0, /* tp_setattro */
3693 0, /* tp_as_buffer */
3694 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3695 Py_TPFLAGS_BASETYPE, /* tp_flags */
3696 compress_doc, /* tp_doc */
3697 (traverseproc)compress_traverse, /* tp_traverse */
3698 0, /* tp_clear */
3699 0, /* tp_richcompare */
3700 0, /* tp_weaklistoffset */
3701 PyObject_SelfIter, /* tp_iter */
3702 (iternextfunc)compress_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003703 compress_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003704 0, /* tp_members */
3705 0, /* tp_getset */
3706 0, /* tp_base */
3707 0, /* tp_dict */
3708 0, /* tp_descr_get */
3709 0, /* tp_descr_set */
3710 0, /* tp_dictoffset */
3711 0, /* tp_init */
3712 0, /* tp_alloc */
3713 compress_new, /* tp_new */
3714 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003715};
3716
3717
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003718/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003719
3720typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003721 PyObject_HEAD
3722 PyObject *func;
3723 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003724} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003725
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003726static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003727
3728static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003729filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003730{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003731 PyObject *func, *seq;
3732 PyObject *it;
3733 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003734
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003735 if (type == &filterfalse_type &&
3736 !_PyArg_NoKeywords("filterfalse()", kwds))
3737 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003738
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003739 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3740 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003741
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003742 /* Get iterator. */
3743 it = PyObject_GetIter(seq);
3744 if (it == NULL)
3745 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003746
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003747 /* create filterfalseobject structure */
3748 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3749 if (lz == NULL) {
3750 Py_DECREF(it);
3751 return NULL;
3752 }
3753 Py_INCREF(func);
3754 lz->func = func;
3755 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003756
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003757 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003758}
3759
3760static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003761filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003762{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003763 PyObject_GC_UnTrack(lz);
3764 Py_XDECREF(lz->func);
3765 Py_XDECREF(lz->it);
3766 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003767}
3768
3769static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003770filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003771{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003772 Py_VISIT(lz->it);
3773 Py_VISIT(lz->func);
3774 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003775}
3776
3777static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003778filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003779{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003780 PyObject *item;
3781 PyObject *it = lz->it;
3782 long ok;
3783 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003784
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003785 iternext = *Py_TYPE(it)->tp_iternext;
3786 for (;;) {
3787 item = iternext(it);
3788 if (item == NULL)
3789 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003790
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003791 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3792 ok = PyObject_IsTrue(item);
3793 } else {
3794 PyObject *good;
3795 good = PyObject_CallFunctionObjArgs(lz->func,
3796 item, NULL);
3797 if (good == NULL) {
3798 Py_DECREF(item);
3799 return NULL;
3800 }
3801 ok = PyObject_IsTrue(good);
3802 Py_DECREF(good);
3803 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003804 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003805 return item;
3806 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003807 if (ok < 0)
3808 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003809 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003810}
3811
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003812static PyObject *
3813filterfalse_reduce(filterfalseobject *lz)
3814{
3815 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3816 lz->func, lz->it);
3817 }
3818
3819static PyMethodDef filterfalse_methods[] = {
3820 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3821 reduce_doc},
3822 {NULL, NULL} /* sentinel */
3823};
3824
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003825PyDoc_STRVAR(filterfalse_doc,
3826"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003827\n\
3828Return those items of sequence for which function(item) is false.\n\
3829If function is None, return the items that are false.");
3830
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003831static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003832 PyVarObject_HEAD_INIT(NULL, 0)
3833 "itertools.filterfalse", /* tp_name */
3834 sizeof(filterfalseobject), /* tp_basicsize */
3835 0, /* tp_itemsize */
3836 /* methods */
3837 (destructor)filterfalse_dealloc, /* tp_dealloc */
3838 0, /* tp_print */
3839 0, /* tp_getattr */
3840 0, /* tp_setattr */
3841 0, /* tp_reserved */
3842 0, /* tp_repr */
3843 0, /* tp_as_number */
3844 0, /* tp_as_sequence */
3845 0, /* tp_as_mapping */
3846 0, /* tp_hash */
3847 0, /* tp_call */
3848 0, /* tp_str */
3849 PyObject_GenericGetAttr, /* tp_getattro */
3850 0, /* tp_setattro */
3851 0, /* tp_as_buffer */
3852 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3853 Py_TPFLAGS_BASETYPE, /* tp_flags */
3854 filterfalse_doc, /* tp_doc */
3855 (traverseproc)filterfalse_traverse, /* tp_traverse */
3856 0, /* tp_clear */
3857 0, /* tp_richcompare */
3858 0, /* tp_weaklistoffset */
3859 PyObject_SelfIter, /* tp_iter */
3860 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003861 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003862 0, /* tp_members */
3863 0, /* tp_getset */
3864 0, /* tp_base */
3865 0, /* tp_dict */
3866 0, /* tp_descr_get */
3867 0, /* tp_descr_set */
3868 0, /* tp_dictoffset */
3869 0, /* tp_init */
3870 0, /* tp_alloc */
3871 filterfalse_new, /* tp_new */
3872 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003873};
3874
3875
3876/* count object ************************************************************/
3877
3878typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003879 PyObject_HEAD
3880 Py_ssize_t cnt;
3881 PyObject *long_cnt;
3882 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003883} countobject;
3884
Raymond Hettinger30729212009-02-12 06:28:27 +00003885/* Counting logic and invariants:
3886
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003887fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003888
Serhiy Storchaka483405b2015-02-17 10:14:30 +02003889 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003890 Advances with: cnt += 1
3891 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003892
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003893slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003894
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003895 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3896 All counting is done with python objects (no overflows or underflows).
3897 Advances with: long_cnt += long_step
3898 Step may be zero -- effectively a slow version of repeat(cnt).
3899 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003900*/
3901
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003902static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003903
3904static PyObject *
3905count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3906{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003907 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003908 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003909 Py_ssize_t cnt = 0;
3910 PyObject *long_cnt = NULL;
3911 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003912 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003913 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003914
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003915 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3916 kwlist, &long_cnt, &long_step))
3917 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003918
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003919 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3920 (long_step != NULL && !PyNumber_Check(long_step))) {
3921 PyErr_SetString(PyExc_TypeError, "a number is required");
3922 return NULL;
3923 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003924
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003925 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
3926 (long_step == NULL || PyLong_Check(long_step));
3927
3928 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003929 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003930 if (fast_mode) {
3931 assert(PyLong_Check(long_cnt));
3932 cnt = PyLong_AsSsize_t(long_cnt);
3933 if (cnt == -1 && PyErr_Occurred()) {
3934 PyErr_Clear();
3935 fast_mode = 0;
3936 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003937 }
3938 Py_INCREF(long_cnt);
3939 } else {
3940 cnt = 0;
3941 long_cnt = PyLong_FromLong(0);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003942 if (long_cnt == NULL) {
3943 return NULL;
3944 }
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 not specified, step defaults to 1 */
3948 if (long_step == NULL) {
3949 long_step = PyLong_FromLong(1);
3950 if (long_step == NULL) {
3951 Py_DECREF(long_cnt);
3952 return NULL;
3953 }
3954 } else
3955 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003956
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003957 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003958
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003959 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003960 if (fast_mode) {
3961 assert(PyLong_Check(long_step));
3962 step = PyLong_AsLong(long_step);
3963 if (step != 1) {
3964 fast_mode = 0;
3965 if (step == -1 && PyErr_Occurred())
3966 PyErr_Clear();
3967 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003968 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003969
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003970 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003971 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003972 else
3973 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003974
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003975 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
3976 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
3977 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003978 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003979
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003980 /* create countobject structure */
3981 lz = (countobject *)type->tp_alloc(type, 0);
3982 if (lz == NULL) {
3983 Py_XDECREF(long_cnt);
3984 return NULL;
3985 }
3986 lz->cnt = cnt;
3987 lz->long_cnt = long_cnt;
3988 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003989
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003990 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003991}
3992
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003993static void
3994count_dealloc(countobject *lz)
3995{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003996 PyObject_GC_UnTrack(lz);
3997 Py_XDECREF(lz->long_cnt);
3998 Py_XDECREF(lz->long_step);
3999 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004000}
4001
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004002static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004003count_traverse(countobject *lz, visitproc visit, void *arg)
4004{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004005 Py_VISIT(lz->long_cnt);
4006 Py_VISIT(lz->long_step);
4007 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004008}
4009
4010static PyObject *
4011count_nextlong(countobject *lz)
4012{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004013 PyObject *long_cnt;
4014 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004015
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004016 long_cnt = lz->long_cnt;
4017 if (long_cnt == NULL) {
4018 /* Switch to slow_mode */
4019 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4020 if (long_cnt == NULL)
4021 return NULL;
4022 }
4023 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004024
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004025 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4026 if (stepped_up == NULL)
4027 return NULL;
4028 lz->long_cnt = stepped_up;
4029 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004030}
4031
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004032static PyObject *
4033count_next(countobject *lz)
4034{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004035 if (lz->cnt == PY_SSIZE_T_MAX)
4036 return count_nextlong(lz);
4037 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004038}
4039
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004040static PyObject *
4041count_repr(countobject *lz)
4042{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004043 if (lz->cnt != PY_SSIZE_T_MAX)
4044 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004045
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004046 if (PyLong_Check(lz->long_step)) {
4047 long step = PyLong_AsLong(lz->long_step);
4048 if (step == -1 && PyErr_Occurred()) {
4049 PyErr_Clear();
4050 }
4051 if (step == 1) {
4052 /* Don't display step when it is an integer equal to 1 */
4053 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
4054 }
4055 }
4056 return PyUnicode_FromFormat("count(%R, %R)",
4057 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004058}
4059
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004060static PyObject *
4061count_reduce(countobject *lz)
4062{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004063 if (lz->cnt == PY_SSIZE_T_MAX)
4064 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4065 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004066}
4067
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004068static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004069 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004070 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004071 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004072};
4073
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004074PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00004075 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004076\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004077Return a count object whose .__next__() method returns consecutive values.\n\
4078Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00004079 def count(firstval=0, step=1):\n\
Eli Benderskyc554f722013-09-03 06:37:19 -07004080 x = firstval\n\
4081 while 1:\n\
4082 yield x\n\
4083 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004084
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004085static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004086 PyVarObject_HEAD_INIT(NULL, 0)
4087 "itertools.count", /* tp_name */
4088 sizeof(countobject), /* tp_basicsize */
4089 0, /* tp_itemsize */
4090 /* methods */
4091 (destructor)count_dealloc, /* tp_dealloc */
4092 0, /* tp_print */
4093 0, /* tp_getattr */
4094 0, /* tp_setattr */
4095 0, /* tp_reserved */
4096 (reprfunc)count_repr, /* tp_repr */
4097 0, /* tp_as_number */
4098 0, /* tp_as_sequence */
4099 0, /* tp_as_mapping */
4100 0, /* tp_hash */
4101 0, /* tp_call */
4102 0, /* tp_str */
4103 PyObject_GenericGetAttr, /* tp_getattro */
4104 0, /* tp_setattro */
4105 0, /* tp_as_buffer */
4106 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4107 Py_TPFLAGS_BASETYPE, /* tp_flags */
4108 count_doc, /* tp_doc */
4109 (traverseproc)count_traverse, /* tp_traverse */
4110 0, /* tp_clear */
4111 0, /* tp_richcompare */
4112 0, /* tp_weaklistoffset */
4113 PyObject_SelfIter, /* tp_iter */
4114 (iternextfunc)count_next, /* tp_iternext */
4115 count_methods, /* tp_methods */
4116 0, /* tp_members */
4117 0, /* tp_getset */
4118 0, /* tp_base */
4119 0, /* tp_dict */
4120 0, /* tp_descr_get */
4121 0, /* tp_descr_set */
4122 0, /* tp_dictoffset */
4123 0, /* tp_init */
4124 0, /* tp_alloc */
4125 count_new, /* tp_new */
4126 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004127};
4128
4129
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004130/* repeat object ************************************************************/
4131
4132typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004133 PyObject_HEAD
4134 PyObject *element;
4135 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004136} repeatobject;
4137
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004138static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004139
4140static PyObject *
4141repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4142{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004143 repeatobject *ro;
4144 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004145 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004146 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004147
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004148 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4149 &element, &cnt))
4150 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004151
Raymond Hettinger97d35552014-06-24 21:36:58 -07004152 if (kwds != NULL)
4153 n_kwds = PyDict_Size(kwds);
4154 /* Does user supply times argument? */
4155 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004156 cnt = 0;
4157
4158 ro = (repeatobject *)type->tp_alloc(type, 0);
4159 if (ro == NULL)
4160 return NULL;
4161 Py_INCREF(element);
4162 ro->element = element;
4163 ro->cnt = cnt;
4164 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004165}
4166
4167static void
4168repeat_dealloc(repeatobject *ro)
4169{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004170 PyObject_GC_UnTrack(ro);
4171 Py_XDECREF(ro->element);
4172 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004173}
4174
4175static int
4176repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4177{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004178 Py_VISIT(ro->element);
4179 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004180}
4181
4182static PyObject *
4183repeat_next(repeatobject *ro)
4184{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004185 if (ro->cnt == 0)
4186 return NULL;
4187 if (ro->cnt > 0)
4188 ro->cnt--;
4189 Py_INCREF(ro->element);
4190 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004191}
4192
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004193static PyObject *
4194repeat_repr(repeatobject *ro)
4195{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004196 if (ro->cnt == -1)
4197 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4198 else
4199 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4200}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004201
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004202static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004203repeat_len(repeatobject *ro)
4204{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004205 if (ro->cnt == -1) {
4206 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4207 return NULL;
4208 }
4209 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004210}
4211
Armin Rigof5b3e362006-02-11 21:32:43 +00004212PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004213
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004214static PyObject *
4215repeat_reduce(repeatobject *ro)
4216{
4217 /* unpickle this so that a new repeat iterator is constructed with an
4218 * object, then call __setstate__ on it to set cnt
4219 */
4220 if (ro->cnt >= 0)
4221 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4222 else
4223 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4224}
4225
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004226static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004227 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004228 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004229 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004230};
4231
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004232PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004233"repeat(object [,times]) -> create an iterator which returns the object\n\
4234for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004235endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004236
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004237static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004238 PyVarObject_HEAD_INIT(NULL, 0)
4239 "itertools.repeat", /* tp_name */
4240 sizeof(repeatobject), /* tp_basicsize */
4241 0, /* tp_itemsize */
4242 /* methods */
4243 (destructor)repeat_dealloc, /* tp_dealloc */
4244 0, /* tp_print */
4245 0, /* tp_getattr */
4246 0, /* tp_setattr */
4247 0, /* tp_reserved */
4248 (reprfunc)repeat_repr, /* tp_repr */
4249 0, /* tp_as_number */
4250 0, /* tp_as_sequence */
4251 0, /* tp_as_mapping */
4252 0, /* tp_hash */
4253 0, /* tp_call */
4254 0, /* tp_str */
4255 PyObject_GenericGetAttr, /* tp_getattro */
4256 0, /* tp_setattro */
4257 0, /* tp_as_buffer */
4258 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4259 Py_TPFLAGS_BASETYPE, /* tp_flags */
4260 repeat_doc, /* tp_doc */
4261 (traverseproc)repeat_traverse, /* tp_traverse */
4262 0, /* tp_clear */
4263 0, /* tp_richcompare */
4264 0, /* tp_weaklistoffset */
4265 PyObject_SelfIter, /* tp_iter */
4266 (iternextfunc)repeat_next, /* tp_iternext */
4267 repeat_methods, /* tp_methods */
4268 0, /* tp_members */
4269 0, /* tp_getset */
4270 0, /* tp_base */
4271 0, /* tp_dict */
4272 0, /* tp_descr_get */
4273 0, /* tp_descr_set */
4274 0, /* tp_dictoffset */
4275 0, /* tp_init */
4276 0, /* tp_alloc */
4277 repeat_new, /* tp_new */
4278 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004279};
4280
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004281/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004282
4283#include "Python.h"
4284
4285typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004286 PyObject_HEAD
4287 Py_ssize_t tuplesize;
4288 Py_ssize_t numactive;
4289 PyObject *ittuple; /* tuple of iterators */
4290 PyObject *result;
4291 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004292} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004293
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004294static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004295
4296static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004297zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004298{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004299 ziplongestobject *lz;
4300 Py_ssize_t i;
4301 PyObject *ittuple; /* tuple of iterators */
4302 PyObject *result;
4303 PyObject *fillvalue = Py_None;
4304 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004305
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004306 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
4307 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4308 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
4309 PyErr_SetString(PyExc_TypeError,
4310 "zip_longest() got an unexpected keyword argument");
4311 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004312 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004313 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004314
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004315 /* args must be a tuple */
4316 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00004317
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004318 /* obtain iterators */
4319 ittuple = PyTuple_New(tuplesize);
4320 if (ittuple == NULL)
4321 return NULL;
4322 for (i=0; i < tuplesize; ++i) {
4323 PyObject *item = PyTuple_GET_ITEM(args, i);
4324 PyObject *it = PyObject_GetIter(item);
4325 if (it == NULL) {
4326 if (PyErr_ExceptionMatches(PyExc_TypeError))
4327 PyErr_Format(PyExc_TypeError,
4328 "zip_longest argument #%zd must support iteration",
4329 i+1);
4330 Py_DECREF(ittuple);
4331 return NULL;
4332 }
4333 PyTuple_SET_ITEM(ittuple, i, it);
4334 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004335
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004336 /* create a result holder */
4337 result = PyTuple_New(tuplesize);
4338 if (result == NULL) {
4339 Py_DECREF(ittuple);
4340 return NULL;
4341 }
4342 for (i=0 ; i < tuplesize ; i++) {
4343 Py_INCREF(Py_None);
4344 PyTuple_SET_ITEM(result, i, Py_None);
4345 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004346
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004347 /* create ziplongestobject structure */
4348 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4349 if (lz == NULL) {
4350 Py_DECREF(ittuple);
4351 Py_DECREF(result);
4352 return NULL;
4353 }
4354 lz->ittuple = ittuple;
4355 lz->tuplesize = tuplesize;
4356 lz->numactive = tuplesize;
4357 lz->result = result;
4358 Py_INCREF(fillvalue);
4359 lz->fillvalue = fillvalue;
4360 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004361}
4362
4363static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004364zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004365{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004366 PyObject_GC_UnTrack(lz);
4367 Py_XDECREF(lz->ittuple);
4368 Py_XDECREF(lz->result);
4369 Py_XDECREF(lz->fillvalue);
4370 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004371}
4372
4373static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004374zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004375{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004376 Py_VISIT(lz->ittuple);
4377 Py_VISIT(lz->result);
4378 Py_VISIT(lz->fillvalue);
4379 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004380}
4381
4382static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004383zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004384{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004385 Py_ssize_t i;
4386 Py_ssize_t tuplesize = lz->tuplesize;
4387 PyObject *result = lz->result;
4388 PyObject *it;
4389 PyObject *item;
4390 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004391
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004392 if (tuplesize == 0)
4393 return NULL;
4394 if (lz->numactive == 0)
4395 return NULL;
4396 if (Py_REFCNT(result) == 1) {
4397 Py_INCREF(result);
4398 for (i=0 ; i < tuplesize ; i++) {
4399 it = PyTuple_GET_ITEM(lz->ittuple, i);
4400 if (it == NULL) {
4401 Py_INCREF(lz->fillvalue);
4402 item = lz->fillvalue;
4403 } else {
4404 item = PyIter_Next(it);
4405 if (item == NULL) {
4406 lz->numactive -= 1;
4407 if (lz->numactive == 0 || PyErr_Occurred()) {
4408 lz->numactive = 0;
4409 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004410 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004411 } else {
4412 Py_INCREF(lz->fillvalue);
4413 item = lz->fillvalue;
4414 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4415 Py_DECREF(it);
4416 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004417 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004418 }
4419 olditem = PyTuple_GET_ITEM(result, i);
4420 PyTuple_SET_ITEM(result, i, item);
4421 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004422 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004423 } else {
4424 result = PyTuple_New(tuplesize);
4425 if (result == NULL)
4426 return NULL;
4427 for (i=0 ; i < tuplesize ; i++) {
4428 it = PyTuple_GET_ITEM(lz->ittuple, i);
4429 if (it == NULL) {
4430 Py_INCREF(lz->fillvalue);
4431 item = lz->fillvalue;
4432 } else {
4433 item = PyIter_Next(it);
4434 if (item == NULL) {
4435 lz->numactive -= 1;
4436 if (lz->numactive == 0 || PyErr_Occurred()) {
4437 lz->numactive = 0;
4438 Py_DECREF(result);
4439 return NULL;
4440 } else {
4441 Py_INCREF(lz->fillvalue);
4442 item = lz->fillvalue;
4443 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4444 Py_DECREF(it);
4445 }
4446 }
4447 }
4448 PyTuple_SET_ITEM(result, i, item);
4449 }
4450 }
4451 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004452}
4453
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004454static PyObject *
4455zip_longest_reduce(ziplongestobject *lz)
4456{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004457
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004458 /* Create a new tuple with empty sequences where appropriate to pickle.
4459 * Then use setstate to set the fillvalue
4460 */
4461 int i;
4462 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
4463 if (args == NULL)
4464 return NULL;
4465 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4466 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4467 if (elem == NULL) {
4468 elem = PyTuple_New(0);
4469 if (elem == NULL) {
4470 Py_DECREF(args);
4471 return NULL;
4472 }
4473 } else
4474 Py_INCREF(elem);
4475 PyTuple_SET_ITEM(args, i, elem);
4476 }
4477 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4478}
4479
4480static PyObject *
4481zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4482{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004483 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004484 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004485 Py_RETURN_NONE;
4486}
4487
4488static PyMethodDef zip_longest_methods[] = {
4489 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4490 reduce_doc},
4491 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4492 setstate_doc},
4493 {NULL, NULL} /* sentinel */
4494};
4495
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004496PyDoc_STRVAR(zip_longest_doc,
4497"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004498\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004499Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004500the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004501method continues until the longest iterable in the argument sequence\n\
4502is exhausted and then it raises StopIteration. When the shorter iterables\n\
4503are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4504defaults to None or can be specified by a keyword argument.\n\
4505");
4506
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004507static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004508 PyVarObject_HEAD_INIT(NULL, 0)
4509 "itertools.zip_longest", /* tp_name */
4510 sizeof(ziplongestobject), /* tp_basicsize */
4511 0, /* tp_itemsize */
4512 /* methods */
4513 (destructor)zip_longest_dealloc, /* tp_dealloc */
4514 0, /* tp_print */
4515 0, /* tp_getattr */
4516 0, /* tp_setattr */
4517 0, /* tp_reserved */
4518 0, /* tp_repr */
4519 0, /* tp_as_number */
4520 0, /* tp_as_sequence */
4521 0, /* tp_as_mapping */
4522 0, /* tp_hash */
4523 0, /* tp_call */
4524 0, /* tp_str */
4525 PyObject_GenericGetAttr, /* tp_getattro */
4526 0, /* tp_setattro */
4527 0, /* tp_as_buffer */
4528 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4529 Py_TPFLAGS_BASETYPE, /* tp_flags */
4530 zip_longest_doc, /* tp_doc */
4531 (traverseproc)zip_longest_traverse, /* tp_traverse */
4532 0, /* tp_clear */
4533 0, /* tp_richcompare */
4534 0, /* tp_weaklistoffset */
4535 PyObject_SelfIter, /* tp_iter */
4536 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004537 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004538 0, /* tp_members */
4539 0, /* tp_getset */
4540 0, /* tp_base */
4541 0, /* tp_dict */
4542 0, /* tp_descr_get */
4543 0, /* tp_descr_set */
4544 0, /* tp_dictoffset */
4545 0, /* tp_init */
4546 0, /* tp_alloc */
4547 zip_longest_new, /* tp_new */
4548 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004549};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004550
4551/* module level code ********************************************************/
4552
4553PyDoc_STRVAR(module_doc,
4554"Functional tools for creating and using iterators.\n\
4555\n\
4556Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004557count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004558cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004559repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004560\n\
4561Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004562accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004563chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger8df58f72013-09-09 01:29:40 -05004564chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004565compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4566dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4567groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004568filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004569islice(seq, [start,] stop [, step]) --> elements from\n\
4570 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004571starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004572tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004573takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004574zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004575\n\
4576Combinatoric generators:\n\
4577product(p, q, ... [repeat=1]) --> cartesian product\n\
4578permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004579combinations(p, r)\n\
4580combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004581");
4582
4583
Raymond Hettingerad983e72003-11-12 14:32:26 +00004584static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004585 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4586 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004587};
4588
Martin v. Löwis1a214512008-06-11 05:26:20 +00004589
4590static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004591 PyModuleDef_HEAD_INIT,
4592 "itertools",
4593 module_doc,
4594 -1,
4595 module_methods,
4596 NULL,
4597 NULL,
4598 NULL,
4599 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004600};
4601
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004602PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004603PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004604{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004605 int i;
4606 PyObject *m;
4607 char *name;
4608 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004609 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004610 &combinations_type,
4611 &cwr_type,
4612 &cycle_type,
4613 &dropwhile_type,
4614 &takewhile_type,
4615 &islice_type,
4616 &starmap_type,
4617 &chain_type,
4618 &compress_type,
4619 &filterfalse_type,
4620 &count_type,
4621 &ziplongest_type,
4622 &permutations_type,
4623 &product_type,
4624 &repeat_type,
4625 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004626 &_grouper_type,
4627 &tee_type,
4628 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004629 NULL
4630 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004631
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004632 Py_TYPE(&teedataobject_type) = &PyType_Type;
4633 m = PyModule_Create(&itertoolsmodule);
4634 if (m == NULL)
4635 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004636
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004637 for (i=0 ; typelist[i] != NULL ; i++) {
4638 if (PyType_Ready(typelist[i]) < 0)
4639 return NULL;
4640 name = strchr(typelist[i]->tp_name, '.');
4641 assert (name != NULL);
4642 Py_INCREF(typelist[i]);
4643 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4644 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004645
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004646 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004647}