blob: be0f4982526f62a836abc1d52c9b3d19897762dc [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;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300158 if (!PyTuple_Check(state)) {
159 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000160 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300161 }
162 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey)) {
163 return NULL;
164 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200165 Py_INCREF(currkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300166 Py_XSETREF(lz->currkey, currkey);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200167 Py_INCREF(currvalue);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300168 Py_XSETREF(lz->currvalue, currvalue);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200169 Py_INCREF(tgtkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300170 Py_XSETREF(lz->tgtkey, tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000171 Py_RETURN_NONE;
172}
173
174PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
175
176static PyMethodDef groupby_methods[] = {
177 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
178 reduce_doc},
179 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
180 setstate_doc},
181 {NULL, NULL} /* sentinel */
182};
183
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000184PyDoc_STRVAR(groupby_doc,
185"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
186(key, sub-iterator) grouped by each value of key(value).\n");
187
188static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000189 PyVarObject_HEAD_INIT(NULL, 0)
190 "itertools.groupby", /* tp_name */
191 sizeof(groupbyobject), /* tp_basicsize */
192 0, /* tp_itemsize */
193 /* methods */
194 (destructor)groupby_dealloc, /* tp_dealloc */
195 0, /* tp_print */
196 0, /* tp_getattr */
197 0, /* tp_setattr */
198 0, /* tp_reserved */
199 0, /* tp_repr */
200 0, /* tp_as_number */
201 0, /* tp_as_sequence */
202 0, /* tp_as_mapping */
203 0, /* tp_hash */
204 0, /* tp_call */
205 0, /* tp_str */
206 PyObject_GenericGetAttr, /* tp_getattro */
207 0, /* tp_setattro */
208 0, /* tp_as_buffer */
209 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
210 Py_TPFLAGS_BASETYPE, /* tp_flags */
211 groupby_doc, /* tp_doc */
212 (traverseproc)groupby_traverse, /* tp_traverse */
213 0, /* tp_clear */
214 0, /* tp_richcompare */
215 0, /* tp_weaklistoffset */
216 PyObject_SelfIter, /* tp_iter */
217 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000218 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000219 0, /* tp_members */
220 0, /* tp_getset */
221 0, /* tp_base */
222 0, /* tp_dict */
223 0, /* tp_descr_get */
224 0, /* tp_descr_set */
225 0, /* tp_dictoffset */
226 0, /* tp_init */
227 0, /* tp_alloc */
228 groupby_new, /* tp_new */
229 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000230};
231
232
233/* _grouper object (internal) ************************************************/
234
235typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000236 PyObject_HEAD
237 PyObject *parent;
238 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000239} _grouperobject;
240
241static PyTypeObject _grouper_type;
242
243static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000244_grouper_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
245{
246 PyObject *parent, *tgtkey;
247
248 if (!PyArg_ParseTuple(args, "O!O", &groupby_type, &parent, &tgtkey))
249 return NULL;
250
251 return _grouper_create((groupbyobject*) parent, tgtkey);
252}
253
254static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000255_grouper_create(groupbyobject *parent, PyObject *tgtkey)
256{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000257 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000259 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
260 if (igo == NULL)
261 return NULL;
262 igo->parent = (PyObject *)parent;
263 Py_INCREF(parent);
264 igo->tgtkey = tgtkey;
265 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000266
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000267 PyObject_GC_Track(igo);
268 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000269}
270
271static void
272_grouper_dealloc(_grouperobject *igo)
273{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000274 PyObject_GC_UnTrack(igo);
275 Py_DECREF(igo->parent);
276 Py_DECREF(igo->tgtkey);
277 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000278}
279
280static int
281_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
282{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000283 Py_VISIT(igo->parent);
284 Py_VISIT(igo->tgtkey);
285 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000286}
287
288static PyObject *
289_grouper_next(_grouperobject *igo)
290{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000291 groupbyobject *gbo = (groupbyobject *)igo->parent;
292 PyObject *newvalue, *newkey, *r;
293 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000294
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000295 if (gbo->currvalue == NULL) {
296 newvalue = PyIter_Next(gbo->it);
297 if (newvalue == NULL)
298 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000299
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000300 if (gbo->keyfunc == Py_None) {
301 newkey = newvalue;
302 Py_INCREF(newvalue);
303 } else {
304 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
305 newvalue, NULL);
306 if (newkey == NULL) {
307 Py_DECREF(newvalue);
308 return NULL;
309 }
310 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000311
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000312 assert(gbo->currkey == NULL);
313 gbo->currkey = newkey;
314 gbo->currvalue = newvalue;
315 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000316
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000317 assert(gbo->currkey != NULL);
318 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
319 if (rcmp <= 0)
320 /* got any error or current group is end */
321 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000322
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000323 r = gbo->currvalue;
324 gbo->currvalue = NULL;
325 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000326
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000327 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000328}
329
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000330static PyObject *
331_grouper_reduce(_grouperobject *lz)
332{
333 return Py_BuildValue("O(OO)", Py_TYPE(lz),
334 lz->parent, lz->tgtkey);
335}
336
337static PyMethodDef _grouper_methods[] = {
338 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
339 reduce_doc},
340 {NULL, NULL} /* sentinel */
341};
342
343
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000344static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000345 PyVarObject_HEAD_INIT(NULL, 0)
346 "itertools._grouper", /* tp_name */
347 sizeof(_grouperobject), /* tp_basicsize */
348 0, /* tp_itemsize */
349 /* methods */
350 (destructor)_grouper_dealloc, /* tp_dealloc */
351 0, /* tp_print */
352 0, /* tp_getattr */
353 0, /* tp_setattr */
354 0, /* tp_reserved */
355 0, /* tp_repr */
356 0, /* tp_as_number */
357 0, /* tp_as_sequence */
358 0, /* tp_as_mapping */
359 0, /* tp_hash */
360 0, /* tp_call */
361 0, /* tp_str */
362 PyObject_GenericGetAttr, /* tp_getattro */
363 0, /* tp_setattro */
364 0, /* tp_as_buffer */
365 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
366 0, /* tp_doc */
367 (traverseproc)_grouper_traverse,/* tp_traverse */
368 0, /* tp_clear */
369 0, /* tp_richcompare */
370 0, /* tp_weaklistoffset */
371 PyObject_SelfIter, /* tp_iter */
372 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000373 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000374 0, /* tp_members */
375 0, /* tp_getset */
376 0, /* tp_base */
377 0, /* tp_dict */
378 0, /* tp_descr_get */
379 0, /* tp_descr_set */
380 0, /* tp_dictoffset */
381 0, /* tp_init */
382 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000383 _grouper_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000384 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000385};
386
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000387
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000388
Raymond Hettingerad983e72003-11-12 14:32:26 +0000389/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000390
Raymond Hettingerad983e72003-11-12 14:32:26 +0000391/* The teedataobject pre-allocates space for LINKCELLS number of objects.
392 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000393 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000394 the other structure members including PyHEAD overhead. The larger the
395 value, the less memory overhead per object and the less time spent
396 allocating/deallocating new links. The smaller the number, the less
397 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000398*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000399#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000400
401typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000402 PyObject_HEAD
403 PyObject *it;
Antoine Pitroub4a46cb2013-09-20 22:19:22 +0200404 int numread; /* 0 <= numread <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000405 PyObject *nextlink;
406 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000407} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000408
409typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000410 PyObject_HEAD
411 teedataobject *dataobj;
Antoine Pitroub4a46cb2013-09-20 22:19:22 +0200412 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000413 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000414} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415
Raymond Hettingerad983e72003-11-12 14:32:26 +0000416static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417
418static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000419teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000420{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000421 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000422
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000423 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
424 if (tdo == NULL)
425 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000426
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000427 tdo->numread = 0;
428 tdo->nextlink = NULL;
429 Py_INCREF(it);
430 tdo->it = it;
431 PyObject_GC_Track(tdo);
432 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000433}
434
435static PyObject *
436teedataobject_jumplink(teedataobject *tdo)
437{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000438 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000439 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000440 Py_XINCREF(tdo->nextlink);
441 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000442}
443
444static PyObject *
445teedataobject_getitem(teedataobject *tdo, int i)
446{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000447 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000448
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000449 assert(i < LINKCELLS);
450 if (i < tdo->numread)
451 value = tdo->values[i];
452 else {
453 /* this is the lead iterator, so fetch more data */
454 assert(i == tdo->numread);
455 value = PyIter_Next(tdo->it);
456 if (value == NULL)
457 return NULL;
458 tdo->numread++;
459 tdo->values[i] = value;
460 }
461 Py_INCREF(value);
462 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000463}
464
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000465static int
466teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
467{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000468 int i;
469 Py_VISIT(tdo->it);
470 for (i = 0; i < tdo->numread; i++)
471 Py_VISIT(tdo->values[i]);
472 Py_VISIT(tdo->nextlink);
473 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000474}
475
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200476static void
477teedataobject_safe_decref(PyObject *obj)
478{
479 while (obj && Py_TYPE(obj) == &teedataobject_type &&
480 Py_REFCNT(obj) == 1) {
481 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
482 ((teedataobject *)obj)->nextlink = NULL;
483 Py_DECREF(obj);
484 obj = nextlink;
485 }
486 Py_XDECREF(obj);
487}
488
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000489static int
490teedataobject_clear(teedataobject *tdo)
491{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000492 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200493 PyObject *tmp;
494
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000495 Py_CLEAR(tdo->it);
496 for (i=0 ; i<tdo->numread ; i++)
497 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200498 tmp = tdo->nextlink;
499 tdo->nextlink = NULL;
500 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000501 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000502}
503
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000504static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000505teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000506{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000507 PyObject_GC_UnTrack(tdo);
508 teedataobject_clear(tdo);
509 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000510}
511
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000512static PyObject *
513teedataobject_reduce(teedataobject *tdo)
514{
515 int i;
516 /* create a temporary list of already iterated values */
517 PyObject *values = PyList_New(tdo->numread);
518 if (!values)
519 return NULL;
520 for (i=0 ; i<tdo->numread ; i++) {
521 Py_INCREF(tdo->values[i]);
522 PyList_SET_ITEM(values, i, tdo->values[i]);
523 }
524 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
525 values,
526 tdo->nextlink ? tdo->nextlink : Py_None);
527}
528
529static PyTypeObject teedataobject_type;
530
531static PyObject *
532teedataobject_new(PyTypeObject *type, PyObject *args, PyObject *kw)
533{
534 teedataobject *tdo;
535 PyObject *it, *values, *next;
536 Py_ssize_t i, len;
537
538 assert(type == &teedataobject_type);
539 if (!PyArg_ParseTuple(args, "OO!O", &it, &PyList_Type, &values, &next))
540 return NULL;
541
542 tdo = (teedataobject *)teedataobject_newinternal(it);
543 if (!tdo)
544 return NULL;
545
546 len = PyList_GET_SIZE(values);
547 if (len > LINKCELLS)
548 goto err;
549 for (i=0; i<len; i++) {
550 tdo->values[i] = PyList_GET_ITEM(values, i);
551 Py_INCREF(tdo->values[i]);
552 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200553 /* len <= LINKCELLS < INT_MAX */
554 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000555
556 if (len == LINKCELLS) {
557 if (next != Py_None) {
558 if (Py_TYPE(next) != &teedataobject_type)
559 goto err;
560 assert(tdo->nextlink == NULL);
561 Py_INCREF(next);
562 tdo->nextlink = next;
563 }
564 } else {
565 if (next != Py_None)
566 goto err; /* shouldn't have a next if we are not full */
567 }
568 return (PyObject*)tdo;
569
570err:
571 Py_XDECREF(tdo);
572 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
573 return NULL;
574}
575
576static PyMethodDef teedataobject_methods[] = {
577 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
578 reduce_doc},
579 {NULL, NULL} /* sentinel */
580};
581
Raymond Hettingerad983e72003-11-12 14:32:26 +0000582PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000583
Raymond Hettingerad983e72003-11-12 14:32:26 +0000584static PyTypeObject teedataobject_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000585 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000586 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000587 sizeof(teedataobject), /* tp_basicsize */
588 0, /* tp_itemsize */
589 /* methods */
590 (destructor)teedataobject_dealloc, /* tp_dealloc */
591 0, /* tp_print */
592 0, /* tp_getattr */
593 0, /* tp_setattr */
594 0, /* tp_reserved */
595 0, /* tp_repr */
596 0, /* tp_as_number */
597 0, /* tp_as_sequence */
598 0, /* tp_as_mapping */
599 0, /* tp_hash */
600 0, /* tp_call */
601 0, /* tp_str */
602 PyObject_GenericGetAttr, /* tp_getattro */
603 0, /* tp_setattro */
604 0, /* tp_as_buffer */
605 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
606 teedataobject_doc, /* tp_doc */
607 (traverseproc)teedataobject_traverse, /* tp_traverse */
608 (inquiry)teedataobject_clear, /* tp_clear */
609 0, /* tp_richcompare */
610 0, /* tp_weaklistoffset */
611 0, /* tp_iter */
612 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000613 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000614 0, /* tp_members */
615 0, /* tp_getset */
616 0, /* tp_base */
617 0, /* tp_dict */
618 0, /* tp_descr_get */
619 0, /* tp_descr_set */
620 0, /* tp_dictoffset */
621 0, /* tp_init */
622 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000623 teedataobject_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000624 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000625};
626
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000627
628static PyTypeObject tee_type;
629
630static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000631tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000632{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000633 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000634
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000635 if (to->index >= LINKCELLS) {
636 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200637 if (link == NULL)
638 return NULL;
Serhiy Storchaka57a01d32016-04-10 18:05:40 +0300639 Py_SETREF(to->dataobj, (teedataobject *)link);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000640 to->index = 0;
641 }
642 value = teedataobject_getitem(to->dataobj, to->index);
643 if (value == NULL)
644 return NULL;
645 to->index++;
646 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000647}
648
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000649static int
650tee_traverse(teeobject *to, visitproc visit, void *arg)
651{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000652 Py_VISIT((PyObject *)to->dataobj);
653 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000654}
655
Raymond Hettingerad983e72003-11-12 14:32:26 +0000656static PyObject *
657tee_copy(teeobject *to)
658{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000659 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000660
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000661 newto = PyObject_GC_New(teeobject, &tee_type);
662 if (newto == NULL)
663 return NULL;
664 Py_INCREF(to->dataobj);
665 newto->dataobj = to->dataobj;
666 newto->index = to->index;
667 newto->weakreflist = NULL;
668 PyObject_GC_Track(newto);
669 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000670}
671
672PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
673
674static PyObject *
675tee_fromiterable(PyObject *iterable)
676{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000677 teeobject *to;
678 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000679
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000680 it = PyObject_GetIter(iterable);
681 if (it == NULL)
682 return NULL;
683 if (PyObject_TypeCheck(it, &tee_type)) {
684 to = (teeobject *)tee_copy((teeobject *)it);
685 goto done;
686 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000687
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000688 to = PyObject_GC_New(teeobject, &tee_type);
689 if (to == NULL)
690 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000691 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000692 if (!to->dataobj) {
693 PyObject_GC_Del(to);
694 to = NULL;
695 goto done;
696 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000697
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000698 to->index = 0;
699 to->weakreflist = NULL;
700 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000701done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000702 Py_XDECREF(it);
703 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000704}
705
706static PyObject *
707tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
708{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000709 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000710
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000711 if (!PyArg_UnpackTuple(args, "_tee", 1, 1, &iterable))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000712 return NULL;
713 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000714}
715
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000716static int
717tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000718{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000719 if (to->weakreflist != NULL)
720 PyObject_ClearWeakRefs((PyObject *) to);
721 Py_CLEAR(to->dataobj);
722 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000723}
724
725static void
726tee_dealloc(teeobject *to)
727{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000728 PyObject_GC_UnTrack(to);
729 tee_clear(to);
730 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000731}
732
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000733static PyObject *
734tee_reduce(teeobject *to)
735{
736 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
737}
738
739static PyObject *
740tee_setstate(teeobject *to, PyObject *state)
741{
742 teedataobject *tdo;
743 int index;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300744 if (!PyTuple_Check(state)) {
745 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000746 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300747 }
748 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
749 return NULL;
750 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000751 if (index < 0 || index > LINKCELLS) {
752 PyErr_SetString(PyExc_ValueError, "Index out of range");
753 return NULL;
754 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200755 Py_INCREF(tdo);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300756 Py_XSETREF(to->dataobj, tdo);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000757 to->index = index;
758 Py_RETURN_NONE;
759}
760
Raymond Hettingerad983e72003-11-12 14:32:26 +0000761PyDoc_STRVAR(teeobject_doc,
762"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000763
Raymond Hettingerad983e72003-11-12 14:32:26 +0000764static PyMethodDef tee_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000765 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000766 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
767 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000768 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000769};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000770
771static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000772 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000773 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000774 sizeof(teeobject), /* tp_basicsize */
775 0, /* tp_itemsize */
776 /* methods */
777 (destructor)tee_dealloc, /* tp_dealloc */
778 0, /* tp_print */
779 0, /* tp_getattr */
780 0, /* tp_setattr */
781 0, /* tp_reserved */
782 0, /* tp_repr */
783 0, /* tp_as_number */
784 0, /* tp_as_sequence */
785 0, /* tp_as_mapping */
786 0, /* tp_hash */
787 0, /* tp_call */
788 0, /* tp_str */
789 0, /* tp_getattro */
790 0, /* tp_setattro */
791 0, /* tp_as_buffer */
792 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
793 teeobject_doc, /* tp_doc */
794 (traverseproc)tee_traverse, /* tp_traverse */
795 (inquiry)tee_clear, /* tp_clear */
796 0, /* tp_richcompare */
797 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
798 PyObject_SelfIter, /* tp_iter */
799 (iternextfunc)tee_next, /* tp_iternext */
800 tee_methods, /* tp_methods */
801 0, /* tp_members */
802 0, /* tp_getset */
803 0, /* tp_base */
804 0, /* tp_dict */
805 0, /* tp_descr_get */
806 0, /* tp_descr_set */
807 0, /* tp_dictoffset */
808 0, /* tp_init */
809 0, /* tp_alloc */
810 tee_new, /* tp_new */
811 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000812};
813
Raymond Hettingerad983e72003-11-12 14:32:26 +0000814static PyObject *
815tee(PyObject *self, PyObject *args)
816{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000817 Py_ssize_t i, n=2;
818 PyObject *it, *iterable, *copyable, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200819 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000820
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000821 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
822 return NULL;
823 if (n < 0) {
824 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
825 return NULL;
826 }
827 result = PyTuple_New(n);
828 if (result == NULL)
829 return NULL;
830 if (n == 0)
831 return result;
832 it = PyObject_GetIter(iterable);
833 if (it == NULL) {
834 Py_DECREF(result);
835 return NULL;
836 }
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200837 if (!_PyObject_HasAttrId(it, &PyId___copy__)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000838 copyable = tee_fromiterable(it);
839 Py_DECREF(it);
840 if (copyable == NULL) {
841 Py_DECREF(result);
842 return NULL;
843 }
844 } else
845 copyable = it;
846 PyTuple_SET_ITEM(result, 0, copyable);
847 for (i=1 ; i<n ; i++) {
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200848
849 copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000850 if (copyable == NULL) {
851 Py_DECREF(result);
852 return NULL;
853 }
854 PyTuple_SET_ITEM(result, i, copyable);
855 }
856 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000857}
858
859PyDoc_STRVAR(tee_doc,
860"tee(iterable, n=2) --> tuple of n independent iterators.");
861
862
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000863/* cycle object **********************************************************/
864
865typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000866 PyObject_HEAD
867 PyObject *it;
868 PyObject *saved;
869 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000870} cycleobject;
871
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000872static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000873
874static PyObject *
875cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
876{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000877 PyObject *it;
878 PyObject *iterable;
879 PyObject *saved;
880 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000881
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000882 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
883 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000884
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000885 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
886 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000887
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000888 /* Get iterator. */
889 it = PyObject_GetIter(iterable);
890 if (it == NULL)
891 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000892
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000893 saved = PyList_New(0);
894 if (saved == NULL) {
895 Py_DECREF(it);
896 return NULL;
897 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000898
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000899 /* create cycleobject structure */
900 lz = (cycleobject *)type->tp_alloc(type, 0);
901 if (lz == NULL) {
902 Py_DECREF(it);
903 Py_DECREF(saved);
904 return NULL;
905 }
906 lz->it = it;
907 lz->saved = saved;
908 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000909
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000910 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000911}
912
913static void
914cycle_dealloc(cycleobject *lz)
915{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000916 PyObject_GC_UnTrack(lz);
917 Py_XDECREF(lz->saved);
918 Py_XDECREF(lz->it);
919 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000920}
921
922static int
923cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
924{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000925 Py_VISIT(lz->it);
926 Py_VISIT(lz->saved);
927 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000928}
929
930static PyObject *
931cycle_next(cycleobject *lz)
932{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000933 PyObject *item;
934 PyObject *it;
935 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000936
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000937 while (1) {
938 item = PyIter_Next(lz->it);
939 if (item != NULL) {
940 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
941 Py_DECREF(item);
942 return NULL;
943 }
944 return item;
945 }
946 if (PyErr_Occurred()) {
947 if (PyErr_ExceptionMatches(PyExc_StopIteration))
948 PyErr_Clear();
949 else
950 return NULL;
951 }
952 if (PyList_Size(lz->saved) == 0)
953 return NULL;
954 it = PyObject_GetIter(lz->saved);
955 if (it == NULL)
956 return NULL;
957 tmp = lz->it;
958 lz->it = it;
959 lz->firstpass = 1;
960 Py_DECREF(tmp);
961 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000962}
963
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000964static PyObject *
965cycle_reduce(cycleobject *lz)
966{
967 /* Create a new cycle with the iterator tuple, then set
968 * the saved state on it.
969 */
Serhiy Storchakad269b5e2013-01-25 13:38:56 +0200970 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz),
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000971 lz->it, lz->saved, lz->firstpass);
972 }
973
974static PyObject *
975cycle_setstate(cycleobject *lz, PyObject *state)
976{
977 PyObject *saved=NULL;
978 int firstpass;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300979 if (!PyTuple_Check(state)) {
980 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000981 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300982 }
983 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
984 return NULL;
985 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200986 Py_XINCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300987 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000988 lz->firstpass = firstpass != 0;
989 Py_RETURN_NONE;
990}
991
992static PyMethodDef cycle_methods[] = {
993 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
994 reduce_doc},
995 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
996 setstate_doc},
997 {NULL, NULL} /* sentinel */
998};
999
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001000PyDoc_STRVAR(cycle_doc,
1001"cycle(iterable) --> cycle object\n\
1002\n\
1003Return elements from the iterable until it is exhausted.\n\
1004Then repeat the sequence indefinitely.");
1005
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001006static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001007 PyVarObject_HEAD_INIT(NULL, 0)
1008 "itertools.cycle", /* tp_name */
1009 sizeof(cycleobject), /* tp_basicsize */
1010 0, /* tp_itemsize */
1011 /* methods */
1012 (destructor)cycle_dealloc, /* tp_dealloc */
1013 0, /* tp_print */
1014 0, /* tp_getattr */
1015 0, /* tp_setattr */
1016 0, /* tp_reserved */
1017 0, /* tp_repr */
1018 0, /* tp_as_number */
1019 0, /* tp_as_sequence */
1020 0, /* tp_as_mapping */
1021 0, /* tp_hash */
1022 0, /* tp_call */
1023 0, /* tp_str */
1024 PyObject_GenericGetAttr, /* tp_getattro */
1025 0, /* tp_setattro */
1026 0, /* tp_as_buffer */
1027 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1028 Py_TPFLAGS_BASETYPE, /* tp_flags */
1029 cycle_doc, /* tp_doc */
1030 (traverseproc)cycle_traverse, /* tp_traverse */
1031 0, /* tp_clear */
1032 0, /* tp_richcompare */
1033 0, /* tp_weaklistoffset */
1034 PyObject_SelfIter, /* tp_iter */
1035 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001036 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001037 0, /* tp_members */
1038 0, /* tp_getset */
1039 0, /* tp_base */
1040 0, /* tp_dict */
1041 0, /* tp_descr_get */
1042 0, /* tp_descr_set */
1043 0, /* tp_dictoffset */
1044 0, /* tp_init */
1045 0, /* tp_alloc */
1046 cycle_new, /* tp_new */
1047 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001048};
1049
1050
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001051/* dropwhile object **********************************************************/
1052
1053typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001054 PyObject_HEAD
1055 PyObject *func;
1056 PyObject *it;
1057 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001058} dropwhileobject;
1059
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001060static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001061
1062static PyObject *
1063dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1064{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001065 PyObject *func, *seq;
1066 PyObject *it;
1067 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001068
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001069 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
1070 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001071
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001072 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
1073 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001074
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001075 /* Get iterator. */
1076 it = PyObject_GetIter(seq);
1077 if (it == NULL)
1078 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001079
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001080 /* create dropwhileobject structure */
1081 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1082 if (lz == NULL) {
1083 Py_DECREF(it);
1084 return NULL;
1085 }
1086 Py_INCREF(func);
1087 lz->func = func;
1088 lz->it = it;
1089 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001090
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001091 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001092}
1093
1094static void
1095dropwhile_dealloc(dropwhileobject *lz)
1096{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001097 PyObject_GC_UnTrack(lz);
1098 Py_XDECREF(lz->func);
1099 Py_XDECREF(lz->it);
1100 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001101}
1102
1103static int
1104dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1105{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001106 Py_VISIT(lz->it);
1107 Py_VISIT(lz->func);
1108 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001109}
1110
1111static PyObject *
1112dropwhile_next(dropwhileobject *lz)
1113{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001114 PyObject *item, *good;
1115 PyObject *it = lz->it;
1116 long ok;
1117 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001118
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001119 iternext = *Py_TYPE(it)->tp_iternext;
1120 for (;;) {
1121 item = iternext(it);
1122 if (item == NULL)
1123 return NULL;
1124 if (lz->start == 1)
1125 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001126
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001127 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1128 if (good == NULL) {
1129 Py_DECREF(item);
1130 return NULL;
1131 }
1132 ok = PyObject_IsTrue(good);
1133 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001134 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001135 lz->start = 1;
1136 return item;
1137 }
1138 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001139 if (ok < 0)
1140 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001141 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001142}
1143
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001144static PyObject *
1145dropwhile_reduce(dropwhileobject *lz)
1146{
1147 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1148 lz->func, lz->it, lz->start);
1149}
1150
1151static PyObject *
1152dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1153{
1154 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001155 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001156 return NULL;
1157 lz->start = start;
1158 Py_RETURN_NONE;
1159}
1160
1161static PyMethodDef dropwhile_methods[] = {
1162 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1163 reduce_doc},
1164 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1165 setstate_doc},
1166 {NULL, NULL} /* sentinel */
1167};
1168
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001169PyDoc_STRVAR(dropwhile_doc,
1170"dropwhile(predicate, iterable) --> dropwhile object\n\
1171\n\
1172Drop items from the iterable while predicate(item) is true.\n\
1173Afterwards, return every element until the iterable is exhausted.");
1174
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001175static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001176 PyVarObject_HEAD_INIT(NULL, 0)
1177 "itertools.dropwhile", /* tp_name */
1178 sizeof(dropwhileobject), /* tp_basicsize */
1179 0, /* tp_itemsize */
1180 /* methods */
1181 (destructor)dropwhile_dealloc, /* tp_dealloc */
1182 0, /* tp_print */
1183 0, /* tp_getattr */
1184 0, /* tp_setattr */
1185 0, /* tp_reserved */
1186 0, /* tp_repr */
1187 0, /* tp_as_number */
1188 0, /* tp_as_sequence */
1189 0, /* tp_as_mapping */
1190 0, /* tp_hash */
1191 0, /* tp_call */
1192 0, /* tp_str */
1193 PyObject_GenericGetAttr, /* tp_getattro */
1194 0, /* tp_setattro */
1195 0, /* tp_as_buffer */
1196 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1197 Py_TPFLAGS_BASETYPE, /* tp_flags */
1198 dropwhile_doc, /* tp_doc */
1199 (traverseproc)dropwhile_traverse, /* tp_traverse */
1200 0, /* tp_clear */
1201 0, /* tp_richcompare */
1202 0, /* tp_weaklistoffset */
1203 PyObject_SelfIter, /* tp_iter */
1204 (iternextfunc)dropwhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001205 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001206 0, /* tp_members */
1207 0, /* tp_getset */
1208 0, /* tp_base */
1209 0, /* tp_dict */
1210 0, /* tp_descr_get */
1211 0, /* tp_descr_set */
1212 0, /* tp_dictoffset */
1213 0, /* tp_init */
1214 0, /* tp_alloc */
1215 dropwhile_new, /* tp_new */
1216 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001217};
1218
1219
1220/* takewhile object **********************************************************/
1221
1222typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001223 PyObject_HEAD
1224 PyObject *func;
1225 PyObject *it;
1226 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001227} takewhileobject;
1228
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001229static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001230
1231static PyObject *
1232takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1233{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001234 PyObject *func, *seq;
1235 PyObject *it;
1236 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001237
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001238 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1239 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001240
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001241 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1242 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001243
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001244 /* Get iterator. */
1245 it = PyObject_GetIter(seq);
1246 if (it == NULL)
1247 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001248
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001249 /* create takewhileobject structure */
1250 lz = (takewhileobject *)type->tp_alloc(type, 0);
1251 if (lz == NULL) {
1252 Py_DECREF(it);
1253 return NULL;
1254 }
1255 Py_INCREF(func);
1256 lz->func = func;
1257 lz->it = it;
1258 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001259
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001260 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001261}
1262
1263static void
1264takewhile_dealloc(takewhileobject *lz)
1265{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001266 PyObject_GC_UnTrack(lz);
1267 Py_XDECREF(lz->func);
1268 Py_XDECREF(lz->it);
1269 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001270}
1271
1272static int
1273takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1274{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001275 Py_VISIT(lz->it);
1276 Py_VISIT(lz->func);
1277 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001278}
1279
1280static PyObject *
1281takewhile_next(takewhileobject *lz)
1282{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001283 PyObject *item, *good;
1284 PyObject *it = lz->it;
1285 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001286
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001287 if (lz->stop == 1)
1288 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001289
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001290 item = (*Py_TYPE(it)->tp_iternext)(it);
1291 if (item == NULL)
1292 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001293
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001294 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1295 if (good == NULL) {
1296 Py_DECREF(item);
1297 return NULL;
1298 }
1299 ok = PyObject_IsTrue(good);
1300 Py_DECREF(good);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001301 if (ok == 1)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001302 return item;
1303 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001304 if (ok == 0)
1305 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001306 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001307}
1308
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001309static PyObject *
1310takewhile_reduce(takewhileobject *lz)
1311{
1312 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1313 lz->func, lz->it, lz->stop);
1314}
1315
1316static PyObject *
1317takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1318{
1319 int stop = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001320 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001321 return NULL;
1322 lz->stop = stop;
1323 Py_RETURN_NONE;
1324}
1325
1326static PyMethodDef takewhile_reduce_methods[] = {
1327 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1328 reduce_doc},
1329 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1330 setstate_doc},
1331 {NULL, NULL} /* sentinel */
1332};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001333PyDoc_STRVAR(takewhile_doc,
1334"takewhile(predicate, iterable) --> takewhile object\n\
1335\n\
1336Return successive entries from an iterable as long as the \n\
1337predicate evaluates to true for each entry.");
1338
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001339static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001340 PyVarObject_HEAD_INIT(NULL, 0)
1341 "itertools.takewhile", /* tp_name */
1342 sizeof(takewhileobject), /* tp_basicsize */
1343 0, /* tp_itemsize */
1344 /* methods */
1345 (destructor)takewhile_dealloc, /* tp_dealloc */
1346 0, /* tp_print */
1347 0, /* tp_getattr */
1348 0, /* tp_setattr */
1349 0, /* tp_reserved */
1350 0, /* tp_repr */
1351 0, /* tp_as_number */
1352 0, /* tp_as_sequence */
1353 0, /* tp_as_mapping */
1354 0, /* tp_hash */
1355 0, /* tp_call */
1356 0, /* tp_str */
1357 PyObject_GenericGetAttr, /* tp_getattro */
1358 0, /* tp_setattro */
1359 0, /* tp_as_buffer */
1360 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1361 Py_TPFLAGS_BASETYPE, /* tp_flags */
1362 takewhile_doc, /* tp_doc */
1363 (traverseproc)takewhile_traverse, /* tp_traverse */
1364 0, /* tp_clear */
1365 0, /* tp_richcompare */
1366 0, /* tp_weaklistoffset */
1367 PyObject_SelfIter, /* tp_iter */
1368 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001369 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001370 0, /* tp_members */
1371 0, /* tp_getset */
1372 0, /* tp_base */
1373 0, /* tp_dict */
1374 0, /* tp_descr_get */
1375 0, /* tp_descr_set */
1376 0, /* tp_dictoffset */
1377 0, /* tp_init */
1378 0, /* tp_alloc */
1379 takewhile_new, /* tp_new */
1380 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001381};
1382
1383
1384/* islice object ************************************************************/
1385
1386typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001387 PyObject_HEAD
1388 PyObject *it;
1389 Py_ssize_t next;
1390 Py_ssize_t stop;
1391 Py_ssize_t step;
1392 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001393} isliceobject;
1394
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001395static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001396
1397static PyObject *
1398islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1399{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001400 PyObject *seq;
1401 Py_ssize_t start=0, stop=-1, step=1;
1402 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1403 Py_ssize_t numargs;
1404 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001405
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001406 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1407 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001408
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001409 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1410 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001411
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001412 numargs = PyTuple_Size(args);
1413 if (numargs == 2) {
1414 if (a1 != Py_None) {
1415 stop = PyLong_AsSsize_t(a1);
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 } else {
1425 if (a1 != Py_None)
1426 start = PyLong_AsSsize_t(a1);
1427 if (start == -1 && PyErr_Occurred())
1428 PyErr_Clear();
1429 if (a2 != Py_None) {
1430 stop = PyLong_AsSsize_t(a2);
1431 if (stop == -1) {
1432 if (PyErr_Occurred())
1433 PyErr_Clear();
1434 PyErr_SetString(PyExc_ValueError,
1435 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1436 return NULL;
1437 }
1438 }
1439 }
1440 if (start<0 || stop<-1) {
1441 PyErr_SetString(PyExc_ValueError,
1442 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1443 return NULL;
1444 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001445
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001446 if (a3 != NULL) {
1447 if (a3 != Py_None)
1448 step = PyLong_AsSsize_t(a3);
1449 if (step == -1 && PyErr_Occurred())
1450 PyErr_Clear();
1451 }
1452 if (step<1) {
1453 PyErr_SetString(PyExc_ValueError,
1454 "Step for islice() must be a positive integer or None.");
1455 return NULL;
1456 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001458 /* Get iterator. */
1459 it = PyObject_GetIter(seq);
1460 if (it == NULL)
1461 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001462
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001463 /* create isliceobject structure */
1464 lz = (isliceobject *)type->tp_alloc(type, 0);
1465 if (lz == NULL) {
1466 Py_DECREF(it);
1467 return NULL;
1468 }
1469 lz->it = it;
1470 lz->next = start;
1471 lz->stop = stop;
1472 lz->step = step;
1473 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001474
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001475 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001476}
1477
1478static void
1479islice_dealloc(isliceobject *lz)
1480{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001481 PyObject_GC_UnTrack(lz);
1482 Py_XDECREF(lz->it);
1483 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001484}
1485
1486static int
1487islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1488{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001489 Py_VISIT(lz->it);
1490 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001491}
1492
1493static PyObject *
1494islice_next(isliceobject *lz)
1495{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001496 PyObject *item;
1497 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001498 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001499 Py_ssize_t oldnext;
1500 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001501
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001502 if (it == NULL)
1503 return NULL;
1504
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001505 iternext = *Py_TYPE(it)->tp_iternext;
1506 while (lz->cnt < lz->next) {
1507 item = iternext(it);
1508 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001509 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001510 Py_DECREF(item);
1511 lz->cnt++;
1512 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001513 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001514 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001515 item = iternext(it);
1516 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001517 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001518 lz->cnt++;
1519 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001520 /* The (size_t) cast below avoids the danger of undefined
1521 behaviour from signed integer overflow. */
1522 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001523 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1524 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001525 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001526
1527empty:
1528 Py_CLEAR(lz->it);
1529 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001530}
1531
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001532static PyObject *
1533islice_reduce(isliceobject *lz)
1534{
1535 /* When unpickled, generate a new object with the same bounds,
1536 * then 'setstate' with the next and count
1537 */
1538 PyObject *stop;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001539 if (lz->it == NULL) {
1540 PyObject *empty_list;
1541 PyObject *empty_it;
1542 empty_list = PyList_New(0);
1543 if (empty_list == NULL)
1544 return NULL;
1545 empty_it = PyObject_GetIter(empty_list);
1546 Py_DECREF(empty_list);
1547 if (empty_it == NULL)
1548 return NULL;
1549 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1550 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001551 if (lz->stop == -1) {
1552 stop = Py_None;
1553 Py_INCREF(stop);
1554 } else {
1555 stop = PyLong_FromSsize_t(lz->stop);
1556 if (stop == NULL)
1557 return NULL;
1558 }
1559 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1560 lz->it, lz->next, stop, lz->step,
1561 lz->cnt);
1562}
1563
1564static PyObject *
1565islice_setstate(isliceobject *lz, PyObject *state)
1566{
1567 Py_ssize_t cnt = PyLong_AsSsize_t(state);
1568 if (cnt == -1 && PyErr_Occurred())
1569 return NULL;
1570 lz->cnt = cnt;
1571 Py_RETURN_NONE;
1572}
1573
1574static PyMethodDef islice_methods[] = {
1575 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1576 reduce_doc},
1577 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1578 setstate_doc},
1579 {NULL, NULL} /* sentinel */
1580};
1581
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001582PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001583"islice(iterable, stop) --> islice object\n\
1584islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001585\n\
1586Return an iterator whose next() method returns selected values from an\n\
1587iterable. If start is specified, will skip all preceding elements;\n\
1588otherwise, start defaults to zero. Step defaults to one. If\n\
1589specified as another value, step determines how many values are \n\
1590skipped between successive calls. Works like a slice() on a list\n\
1591but returns an iterator.");
1592
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001593static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001594 PyVarObject_HEAD_INIT(NULL, 0)
1595 "itertools.islice", /* tp_name */
1596 sizeof(isliceobject), /* tp_basicsize */
1597 0, /* tp_itemsize */
1598 /* methods */
1599 (destructor)islice_dealloc, /* tp_dealloc */
1600 0, /* tp_print */
1601 0, /* tp_getattr */
1602 0, /* tp_setattr */
1603 0, /* tp_reserved */
1604 0, /* tp_repr */
1605 0, /* tp_as_number */
1606 0, /* tp_as_sequence */
1607 0, /* tp_as_mapping */
1608 0, /* tp_hash */
1609 0, /* tp_call */
1610 0, /* tp_str */
1611 PyObject_GenericGetAttr, /* tp_getattro */
1612 0, /* tp_setattro */
1613 0, /* tp_as_buffer */
1614 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1615 Py_TPFLAGS_BASETYPE, /* tp_flags */
1616 islice_doc, /* tp_doc */
1617 (traverseproc)islice_traverse, /* tp_traverse */
1618 0, /* tp_clear */
1619 0, /* tp_richcompare */
1620 0, /* tp_weaklistoffset */
1621 PyObject_SelfIter, /* tp_iter */
1622 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001623 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001624 0, /* tp_members */
1625 0, /* tp_getset */
1626 0, /* tp_base */
1627 0, /* tp_dict */
1628 0, /* tp_descr_get */
1629 0, /* tp_descr_set */
1630 0, /* tp_dictoffset */
1631 0, /* tp_init */
1632 0, /* tp_alloc */
1633 islice_new, /* tp_new */
1634 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001635};
1636
1637
1638/* starmap object ************************************************************/
1639
1640typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001641 PyObject_HEAD
1642 PyObject *func;
1643 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001644} starmapobject;
1645
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001646static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001647
1648static PyObject *
1649starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1650{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001651 PyObject *func, *seq;
1652 PyObject *it;
1653 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001654
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001655 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1656 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001657
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001658 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1659 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001660
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001661 /* Get iterator. */
1662 it = PyObject_GetIter(seq);
1663 if (it == NULL)
1664 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001665
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001666 /* create starmapobject structure */
1667 lz = (starmapobject *)type->tp_alloc(type, 0);
1668 if (lz == NULL) {
1669 Py_DECREF(it);
1670 return NULL;
1671 }
1672 Py_INCREF(func);
1673 lz->func = func;
1674 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001675
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001676 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001677}
1678
1679static void
1680starmap_dealloc(starmapobject *lz)
1681{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001682 PyObject_GC_UnTrack(lz);
1683 Py_XDECREF(lz->func);
1684 Py_XDECREF(lz->it);
1685 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001686}
1687
1688static int
1689starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1690{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001691 Py_VISIT(lz->it);
1692 Py_VISIT(lz->func);
1693 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001694}
1695
1696static PyObject *
1697starmap_next(starmapobject *lz)
1698{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001699 PyObject *args;
1700 PyObject *result;
1701 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001702
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001703 args = (*Py_TYPE(it)->tp_iternext)(it);
1704 if (args == NULL)
1705 return NULL;
1706 if (!PyTuple_CheckExact(args)) {
1707 PyObject *newargs = PySequence_Tuple(args);
1708 Py_DECREF(args);
1709 if (newargs == NULL)
1710 return NULL;
1711 args = newargs;
1712 }
1713 result = PyObject_Call(lz->func, args, NULL);
1714 Py_DECREF(args);
1715 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001716}
1717
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001718static PyObject *
1719starmap_reduce(starmapobject *lz)
1720{
1721 /* Just pickle the iterator */
1722 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1723}
1724
1725static PyMethodDef starmap_methods[] = {
1726 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1727 reduce_doc},
1728 {NULL, NULL} /* sentinel */
1729};
1730
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001731PyDoc_STRVAR(starmap_doc,
1732"starmap(function, sequence) --> starmap object\n\
1733\n\
1734Return an iterator whose values are returned from the function evaluated\n\
Serhiy Storchakad65c9492015-11-02 14:10:23 +02001735with an argument tuple taken from the given sequence.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001736
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001737static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001738 PyVarObject_HEAD_INIT(NULL, 0)
1739 "itertools.starmap", /* tp_name */
1740 sizeof(starmapobject), /* tp_basicsize */
1741 0, /* tp_itemsize */
1742 /* methods */
1743 (destructor)starmap_dealloc, /* tp_dealloc */
1744 0, /* tp_print */
1745 0, /* tp_getattr */
1746 0, /* tp_setattr */
1747 0, /* tp_reserved */
1748 0, /* tp_repr */
1749 0, /* tp_as_number */
1750 0, /* tp_as_sequence */
1751 0, /* tp_as_mapping */
1752 0, /* tp_hash */
1753 0, /* tp_call */
1754 0, /* tp_str */
1755 PyObject_GenericGetAttr, /* tp_getattro */
1756 0, /* tp_setattro */
1757 0, /* tp_as_buffer */
1758 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1759 Py_TPFLAGS_BASETYPE, /* tp_flags */
1760 starmap_doc, /* tp_doc */
1761 (traverseproc)starmap_traverse, /* tp_traverse */
1762 0, /* tp_clear */
1763 0, /* tp_richcompare */
1764 0, /* tp_weaklistoffset */
1765 PyObject_SelfIter, /* tp_iter */
1766 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001767 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001768 0, /* tp_members */
1769 0, /* tp_getset */
1770 0, /* tp_base */
1771 0, /* tp_dict */
1772 0, /* tp_descr_get */
1773 0, /* tp_descr_set */
1774 0, /* tp_dictoffset */
1775 0, /* tp_init */
1776 0, /* tp_alloc */
1777 starmap_new, /* tp_new */
1778 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001779};
1780
1781
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001782/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001783
1784typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001785 PyObject_HEAD
1786 PyObject *source; /* Iterator over input iterables */
1787 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001788} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001789
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001790static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001791
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001792static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001793chain_new_internal(PyTypeObject *type, PyObject *source)
1794{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001795 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001796
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001797 lz = (chainobject *)type->tp_alloc(type, 0);
1798 if (lz == NULL) {
1799 Py_DECREF(source);
1800 return NULL;
1801 }
1802
1803 lz->source = source;
1804 lz->active = NULL;
1805 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001806}
1807
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001808static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001809chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001810{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001811 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001812
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001813 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1814 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001815
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001816 source = PyObject_GetIter(args);
1817 if (source == NULL)
1818 return NULL;
1819
1820 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001821}
1822
1823static PyObject *
1824chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1825{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001826 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001827
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001828 source = PyObject_GetIter(arg);
1829 if (source == NULL)
1830 return NULL;
1831
1832 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001833}
1834
1835static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001836chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001837{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001838 PyObject_GC_UnTrack(lz);
1839 Py_XDECREF(lz->active);
1840 Py_XDECREF(lz->source);
1841 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001842}
1843
Raymond Hettinger2012f172003-02-07 05:32:58 +00001844static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001845chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001846{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001847 Py_VISIT(lz->source);
1848 Py_VISIT(lz->active);
1849 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001850}
1851
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001852static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001853chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001854{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001855 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001856
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001857 if (lz->source == NULL)
1858 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001859
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001860 if (lz->active == NULL) {
1861 PyObject *iterable = PyIter_Next(lz->source);
1862 if (iterable == NULL) {
1863 Py_CLEAR(lz->source);
1864 return NULL; /* no more input sources */
1865 }
1866 lz->active = PyObject_GetIter(iterable);
1867 Py_DECREF(iterable);
1868 if (lz->active == NULL) {
1869 Py_CLEAR(lz->source);
1870 return NULL; /* input not iterable */
1871 }
1872 }
1873 item = PyIter_Next(lz->active);
1874 if (item != NULL)
1875 return item;
1876 if (PyErr_Occurred()) {
1877 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1878 PyErr_Clear();
1879 else
1880 return NULL; /* input raised an exception */
1881 }
1882 Py_CLEAR(lz->active);
1883 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001884}
1885
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001886static PyObject *
1887chain_reduce(chainobject *lz)
1888{
1889 if (lz->source) {
1890 /* we can't pickle function objects (itertools.from_iterable) so
1891 * we must use setstate to replace the iterable. One day we
1892 * will fix pickling of functions
1893 */
1894 if (lz->active) {
1895 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1896 } else {
1897 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1898 }
1899 } else {
1900 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1901 }
1902 return NULL;
1903}
1904
1905static PyObject *
1906chain_setstate(chainobject *lz, PyObject *state)
1907{
1908 PyObject *source, *active=NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001909
1910 if (!PyTuple_Check(state)) {
1911 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001912 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001913 }
1914 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
1915 return NULL;
1916 }
1917 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
1918 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
1919 return NULL;
1920 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001921
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02001922 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001923 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02001924 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001925 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001926 Py_RETURN_NONE;
1927}
1928
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001929PyDoc_STRVAR(chain_doc,
1930"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001931\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001932Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001933first iterable until it is exhausted, then elements from the next\n\
1934iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001935
Christian Heimesf16baeb2008-02-29 14:57:44 +00001936PyDoc_STRVAR(chain_from_iterable_doc,
1937"chain.from_iterable(iterable) --> chain object\n\
1938\n\
Martin Pantereb995702016-07-28 01:11:04 +00001939Alternate chain() constructor taking a single iterable argument\n\
Christian Heimesf16baeb2008-02-29 14:57:44 +00001940that evaluates lazily.");
1941
1942static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001943 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1944 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001945 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1946 reduce_doc},
1947 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1948 setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001949 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001950};
1951
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001952static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001953 PyVarObject_HEAD_INIT(NULL, 0)
1954 "itertools.chain", /* tp_name */
1955 sizeof(chainobject), /* tp_basicsize */
1956 0, /* tp_itemsize */
1957 /* methods */
1958 (destructor)chain_dealloc, /* tp_dealloc */
1959 0, /* tp_print */
1960 0, /* tp_getattr */
1961 0, /* tp_setattr */
1962 0, /* tp_reserved */
1963 0, /* tp_repr */
1964 0, /* tp_as_number */
1965 0, /* tp_as_sequence */
1966 0, /* tp_as_mapping */
1967 0, /* tp_hash */
1968 0, /* tp_call */
1969 0, /* tp_str */
1970 PyObject_GenericGetAttr, /* tp_getattro */
1971 0, /* tp_setattro */
1972 0, /* tp_as_buffer */
1973 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1974 Py_TPFLAGS_BASETYPE, /* tp_flags */
1975 chain_doc, /* tp_doc */
1976 (traverseproc)chain_traverse, /* tp_traverse */
1977 0, /* tp_clear */
1978 0, /* tp_richcompare */
1979 0, /* tp_weaklistoffset */
1980 PyObject_SelfIter, /* tp_iter */
1981 (iternextfunc)chain_next, /* tp_iternext */
1982 chain_methods, /* tp_methods */
1983 0, /* tp_members */
1984 0, /* tp_getset */
1985 0, /* tp_base */
1986 0, /* tp_dict */
1987 0, /* tp_descr_get */
1988 0, /* tp_descr_set */
1989 0, /* tp_dictoffset */
1990 0, /* tp_init */
1991 0, /* tp_alloc */
1992 chain_new, /* tp_new */
1993 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001994};
1995
1996
Christian Heimesc3f30c42008-02-22 16:37:40 +00001997/* product object ************************************************************/
1998
1999typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002000 PyObject_HEAD
2001 PyObject *pools; /* tuple of pool tuples */
2002 Py_ssize_t *indices; /* one index per pool */
2003 PyObject *result; /* most recently returned result tuple */
2004 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002005} productobject;
2006
2007static PyTypeObject product_type;
2008
2009static PyObject *
2010product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2011{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002012 productobject *lz;
2013 Py_ssize_t nargs, npools, repeat=1;
2014 PyObject *pools = NULL;
2015 Py_ssize_t *indices = NULL;
2016 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002017
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002018 if (kwds != NULL) {
2019 char *kwlist[] = {"repeat", 0};
2020 PyObject *tmpargs = PyTuple_New(0);
2021 if (tmpargs == NULL)
2022 return NULL;
2023 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
2024 Py_DECREF(tmpargs);
2025 return NULL;
2026 }
2027 Py_DECREF(tmpargs);
2028 if (repeat < 0) {
2029 PyErr_SetString(PyExc_ValueError,
2030 "repeat argument cannot be negative");
2031 return NULL;
2032 }
2033 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002034
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002035 assert(PyTuple_CheckExact(args));
2036 if (repeat == 0) {
2037 nargs = 0;
2038 } else {
2039 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002040 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002041 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2042 return NULL;
2043 }
2044 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002045 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002046
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002047 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002048 if (indices == NULL) {
2049 PyErr_NoMemory();
2050 goto error;
2051 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002052
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002053 pools = PyTuple_New(npools);
2054 if (pools == NULL)
2055 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002056
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002057 for (i=0; i < nargs ; ++i) {
2058 PyObject *item = PyTuple_GET_ITEM(args, i);
2059 PyObject *pool = PySequence_Tuple(item);
2060 if (pool == NULL)
2061 goto error;
2062 PyTuple_SET_ITEM(pools, i, pool);
2063 indices[i] = 0;
2064 }
2065 for ( ; i < npools; ++i) {
2066 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2067 Py_INCREF(pool);
2068 PyTuple_SET_ITEM(pools, i, pool);
2069 indices[i] = 0;
2070 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002071
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002072 /* create productobject structure */
2073 lz = (productobject *)type->tp_alloc(type, 0);
2074 if (lz == NULL)
2075 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002076
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002077 lz->pools = pools;
2078 lz->indices = indices;
2079 lz->result = NULL;
2080 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002081
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002082 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002083
2084error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002085 if (indices != NULL)
2086 PyMem_Free(indices);
2087 Py_XDECREF(pools);
2088 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002089}
2090
2091static void
2092product_dealloc(productobject *lz)
2093{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002094 PyObject_GC_UnTrack(lz);
2095 Py_XDECREF(lz->pools);
2096 Py_XDECREF(lz->result);
2097 if (lz->indices != NULL)
2098 PyMem_Free(lz->indices);
2099 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002100}
2101
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002102static PyObject *
2103product_sizeof(productobject *lz, void *unused)
2104{
2105 Py_ssize_t res;
2106
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002107 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002108 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2109 return PyLong_FromSsize_t(res);
2110}
2111
2112PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2113
Christian Heimesc3f30c42008-02-22 16:37:40 +00002114static int
2115product_traverse(productobject *lz, visitproc visit, void *arg)
2116{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002117 Py_VISIT(lz->pools);
2118 Py_VISIT(lz->result);
2119 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002120}
2121
2122static PyObject *
2123product_next(productobject *lz)
2124{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002125 PyObject *pool;
2126 PyObject *elem;
2127 PyObject *oldelem;
2128 PyObject *pools = lz->pools;
2129 PyObject *result = lz->result;
2130 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2131 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002132
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002133 if (lz->stopped)
2134 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002135
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002136 if (result == NULL) {
2137 /* On the first pass, return an initial tuple filled with the
2138 first element from each pool. */
2139 result = PyTuple_New(npools);
2140 if (result == NULL)
2141 goto empty;
2142 lz->result = result;
2143 for (i=0; i < npools; i++) {
2144 pool = PyTuple_GET_ITEM(pools, i);
2145 if (PyTuple_GET_SIZE(pool) == 0)
2146 goto empty;
2147 elem = PyTuple_GET_ITEM(pool, 0);
2148 Py_INCREF(elem);
2149 PyTuple_SET_ITEM(result, i, elem);
2150 }
2151 } else {
2152 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002153
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002154 /* Copy the previous result tuple or re-use it if available */
2155 if (Py_REFCNT(result) > 1) {
2156 PyObject *old_result = result;
2157 result = PyTuple_New(npools);
2158 if (result == NULL)
2159 goto empty;
2160 lz->result = result;
2161 for (i=0; i < npools; i++) {
2162 elem = PyTuple_GET_ITEM(old_result, i);
2163 Py_INCREF(elem);
2164 PyTuple_SET_ITEM(result, i, elem);
2165 }
2166 Py_DECREF(old_result);
2167 }
2168 /* Now, we've got the only copy so we can update it in-place */
2169 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002170
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002171 /* Update the pool indices right-to-left. Only advance to the
2172 next pool when the previous one rolls-over */
2173 for (i=npools-1 ; i >= 0 ; i--) {
2174 pool = PyTuple_GET_ITEM(pools, i);
2175 indices[i]++;
2176 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2177 /* Roll-over and advance to next pool */
2178 indices[i] = 0;
2179 elem = PyTuple_GET_ITEM(pool, 0);
2180 Py_INCREF(elem);
2181 oldelem = PyTuple_GET_ITEM(result, i);
2182 PyTuple_SET_ITEM(result, i, elem);
2183 Py_DECREF(oldelem);
2184 } else {
2185 /* No rollover. Just increment and stop here. */
2186 elem = PyTuple_GET_ITEM(pool, indices[i]);
2187 Py_INCREF(elem);
2188 oldelem = PyTuple_GET_ITEM(result, i);
2189 PyTuple_SET_ITEM(result, i, elem);
2190 Py_DECREF(oldelem);
2191 break;
2192 }
2193 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002194
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002195 /* If i is negative, then the indices have all rolled-over
2196 and we're done. */
2197 if (i < 0)
2198 goto empty;
2199 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002200
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002201 Py_INCREF(result);
2202 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002203
2204empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002205 lz->stopped = 1;
2206 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002207}
2208
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002209static PyObject *
2210product_reduce(productobject *lz)
2211{
2212 if (lz->stopped) {
2213 return Py_BuildValue("O(())", Py_TYPE(lz));
2214 } else if (lz->result == NULL) {
2215 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2216 } else {
2217 PyObject *indices;
2218 Py_ssize_t n, i;
2219
2220 /* we must pickle the indices use them for setstate, and
2221 * additionally indicate that the iterator has started
2222 */
2223 n = PyTuple_GET_SIZE(lz->pools);
2224 indices = PyTuple_New(n);
2225 if (indices == NULL)
2226 return NULL;
2227 for (i=0; i<n; i++){
2228 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2229 if (!index) {
2230 Py_DECREF(indices);
2231 return NULL;
2232 }
2233 PyTuple_SET_ITEM(indices, i, index);
2234 }
2235 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2236 }
2237}
2238
2239static PyObject *
2240product_setstate(productobject *lz, PyObject *state)
2241{
2242 PyObject *result;
2243 Py_ssize_t n, i;
2244
2245 n = PyTuple_GET_SIZE(lz->pools);
2246 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2247 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2248 return NULL;
2249 }
2250 for (i=0; i<n; i++)
2251 {
2252 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2253 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002254 PyObject* pool;
2255 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002256 if (index < 0 && PyErr_Occurred())
2257 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002258 pool = PyTuple_GET_ITEM(lz->pools, i);
2259 poolsize = PyTuple_GET_SIZE(pool);
2260 if (poolsize == 0) {
2261 lz->stopped = 1;
2262 Py_RETURN_NONE;
2263 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002264 /* clamp the index */
2265 if (index < 0)
2266 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002267 else if (index > poolsize-1)
2268 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002269 lz->indices[i] = index;
2270 }
2271
2272 result = PyTuple_New(n);
2273 if (!result)
2274 return NULL;
2275 for (i=0; i<n; i++) {
2276 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2277 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2278 Py_INCREF(element);
2279 PyTuple_SET_ITEM(result, i, element);
2280 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002281 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002282 Py_RETURN_NONE;
2283}
2284
2285static PyMethodDef product_methods[] = {
2286 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2287 reduce_doc},
2288 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2289 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002290 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2291 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002292 {NULL, NULL} /* sentinel */
2293};
2294
Christian Heimesc3f30c42008-02-22 16:37:40 +00002295PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002296"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002297\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002298Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002299For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2300The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2301cycle in a manner similar to an odometer (with the rightmost element changing\n\
2302on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002303To compute the product of an iterable with itself, specify the number\n\
2304of repetitions with the optional repeat keyword argument. For example,\n\
2305product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002306product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2307product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2308
2309static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002310 PyVarObject_HEAD_INIT(NULL, 0)
2311 "itertools.product", /* tp_name */
2312 sizeof(productobject), /* tp_basicsize */
2313 0, /* tp_itemsize */
2314 /* methods */
2315 (destructor)product_dealloc, /* tp_dealloc */
2316 0, /* tp_print */
2317 0, /* tp_getattr */
2318 0, /* tp_setattr */
2319 0, /* tp_reserved */
2320 0, /* tp_repr */
2321 0, /* tp_as_number */
2322 0, /* tp_as_sequence */
2323 0, /* tp_as_mapping */
2324 0, /* tp_hash */
2325 0, /* tp_call */
2326 0, /* tp_str */
2327 PyObject_GenericGetAttr, /* tp_getattro */
2328 0, /* tp_setattro */
2329 0, /* tp_as_buffer */
2330 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2331 Py_TPFLAGS_BASETYPE, /* tp_flags */
2332 product_doc, /* tp_doc */
2333 (traverseproc)product_traverse, /* tp_traverse */
2334 0, /* tp_clear */
2335 0, /* tp_richcompare */
2336 0, /* tp_weaklistoffset */
2337 PyObject_SelfIter, /* tp_iter */
2338 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002339 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002340 0, /* tp_members */
2341 0, /* tp_getset */
2342 0, /* tp_base */
2343 0, /* tp_dict */
2344 0, /* tp_descr_get */
2345 0, /* tp_descr_set */
2346 0, /* tp_dictoffset */
2347 0, /* tp_init */
2348 0, /* tp_alloc */
2349 product_new, /* tp_new */
2350 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002351};
2352
2353
Christian Heimes380f7f22008-02-28 11:19:05 +00002354/* combinations object ************************************************************/
2355
2356typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002357 PyObject_HEAD
2358 PyObject *pool; /* input converted to a tuple */
2359 Py_ssize_t *indices; /* one index per result element */
2360 PyObject *result; /* most recently returned result tuple */
2361 Py_ssize_t r; /* size of result tuple */
2362 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002363} combinationsobject;
2364
2365static PyTypeObject combinations_type;
2366
2367static PyObject *
2368combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2369{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002370 combinationsobject *co;
2371 Py_ssize_t n;
2372 Py_ssize_t r;
2373 PyObject *pool = NULL;
2374 PyObject *iterable = NULL;
2375 Py_ssize_t *indices = NULL;
2376 Py_ssize_t i;
2377 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002378
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002379 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2380 &iterable, &r))
2381 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002382
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002383 pool = PySequence_Tuple(iterable);
2384 if (pool == NULL)
2385 goto error;
2386 n = PyTuple_GET_SIZE(pool);
2387 if (r < 0) {
2388 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2389 goto error;
2390 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002391
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002392 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002393 if (indices == NULL) {
2394 PyErr_NoMemory();
2395 goto error;
2396 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002397
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002398 for (i=0 ; i<r ; i++)
2399 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002400
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002401 /* create combinationsobject structure */
2402 co = (combinationsobject *)type->tp_alloc(type, 0);
2403 if (co == NULL)
2404 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002405
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002406 co->pool = pool;
2407 co->indices = indices;
2408 co->result = NULL;
2409 co->r = r;
2410 co->stopped = r > n ? 1 : 0;
2411
2412 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002413
2414error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002415 if (indices != NULL)
2416 PyMem_Free(indices);
2417 Py_XDECREF(pool);
2418 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002419}
2420
2421static void
2422combinations_dealloc(combinationsobject *co)
2423{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002424 PyObject_GC_UnTrack(co);
2425 Py_XDECREF(co->pool);
2426 Py_XDECREF(co->result);
2427 if (co->indices != NULL)
2428 PyMem_Free(co->indices);
2429 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002430}
2431
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002432static PyObject *
2433combinations_sizeof(combinationsobject *co, void *unused)
2434{
2435 Py_ssize_t res;
2436
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002437 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002438 res += co->r * sizeof(Py_ssize_t);
2439 return PyLong_FromSsize_t(res);
2440}
2441
Christian Heimes380f7f22008-02-28 11:19:05 +00002442static int
2443combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2444{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002445 Py_VISIT(co->pool);
2446 Py_VISIT(co->result);
2447 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002448}
2449
2450static PyObject *
2451combinations_next(combinationsobject *co)
2452{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002453 PyObject *elem;
2454 PyObject *oldelem;
2455 PyObject *pool = co->pool;
2456 Py_ssize_t *indices = co->indices;
2457 PyObject *result = co->result;
2458 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2459 Py_ssize_t r = co->r;
2460 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002461
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002462 if (co->stopped)
2463 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002464
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002465 if (result == NULL) {
2466 /* On the first pass, initialize result tuple using the indices */
2467 result = PyTuple_New(r);
2468 if (result == NULL)
2469 goto empty;
2470 co->result = result;
2471 for (i=0; i<r ; i++) {
2472 index = indices[i];
2473 elem = PyTuple_GET_ITEM(pool, index);
2474 Py_INCREF(elem);
2475 PyTuple_SET_ITEM(result, i, elem);
2476 }
2477 } else {
2478 /* Copy the previous result tuple or re-use it if available */
2479 if (Py_REFCNT(result) > 1) {
2480 PyObject *old_result = result;
2481 result = PyTuple_New(r);
2482 if (result == NULL)
2483 goto empty;
2484 co->result = result;
2485 for (i=0; i<r ; i++) {
2486 elem = PyTuple_GET_ITEM(old_result, i);
2487 Py_INCREF(elem);
2488 PyTuple_SET_ITEM(result, i, elem);
2489 }
2490 Py_DECREF(old_result);
2491 }
2492 /* Now, we've got the only copy so we can update it in-place
2493 * CPython's empty tuple is a singleton and cached in
2494 * PyTuple's freelist.
2495 */
2496 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002497
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002498 /* Scan indices right-to-left until finding one that is not
2499 at its maximum (i + n - r). */
2500 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2501 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002502
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002503 /* If i is negative, then the indices are all at
2504 their maximum value and we're done. */
2505 if (i < 0)
2506 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002507
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002508 /* Increment the current index which we know is not at its
2509 maximum. Then move back to the right setting each index
2510 to its lowest possible value (one higher than the index
2511 to its left -- this maintains the sort order invariant). */
2512 indices[i]++;
2513 for (j=i+1 ; j<r ; j++)
2514 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002515
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002516 /* Update the result tuple for the new indices
2517 starting with i, the leftmost index that changed */
2518 for ( ; i<r ; i++) {
2519 index = indices[i];
2520 elem = PyTuple_GET_ITEM(pool, index);
2521 Py_INCREF(elem);
2522 oldelem = PyTuple_GET_ITEM(result, i);
2523 PyTuple_SET_ITEM(result, i, elem);
2524 Py_DECREF(oldelem);
2525 }
2526 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002527
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002528 Py_INCREF(result);
2529 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002530
2531empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002532 co->stopped = 1;
2533 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002534}
2535
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002536static PyObject *
2537combinations_reduce(combinationsobject *lz)
2538{
2539 if (lz->result == NULL) {
2540 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2541 } else if (lz->stopped) {
2542 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2543 } else {
2544 PyObject *indices;
2545 Py_ssize_t i;
2546
2547 /* we must pickle the indices and use them for setstate */
2548 indices = PyTuple_New(lz->r);
2549 if (!indices)
2550 return NULL;
2551 for (i=0; i<lz->r; i++)
2552 {
2553 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2554 if (!index) {
2555 Py_DECREF(indices);
2556 return NULL;
2557 }
2558 PyTuple_SET_ITEM(indices, i, index);
2559 }
2560
2561 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2562 }
2563}
2564
2565static PyObject *
2566combinations_setstate(combinationsobject *lz, PyObject *state)
2567{
2568 PyObject *result;
2569 Py_ssize_t i;
2570 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2571
2572 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2573 {
2574 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2575 return NULL;
2576 }
2577
2578 for (i=0; i<lz->r; i++)
2579 {
2580 Py_ssize_t max;
2581 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2582 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2583 if (index == -1 && PyErr_Occurred())
2584 return NULL; /* not an integer */
2585 max = i + n - lz->r;
2586 /* clamp the index (beware of negative max) */
2587 if (index > max)
2588 index = max;
2589 if (index < 0)
2590 index = 0;
2591 lz->indices[i] = index;
2592 }
2593
2594 result = PyTuple_New(lz->r);
2595 if (result == NULL)
2596 return NULL;
2597 for (i=0; i<lz->r; i++) {
2598 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2599 Py_INCREF(element);
2600 PyTuple_SET_ITEM(result, i, element);
2601 }
2602
Serhiy Storchaka48842712016-04-06 09:45:48 +03002603 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002604 Py_RETURN_NONE;
2605}
2606
2607static PyMethodDef combinations_methods[] = {
2608 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2609 reduce_doc},
2610 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2611 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002612 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2613 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002614 {NULL, NULL} /* sentinel */
2615};
2616
Christian Heimes380f7f22008-02-28 11:19:05 +00002617PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002618"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002619\n\
2620Return successive r-length combinations of elements in the iterable.\n\n\
2621combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2622
2623static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002624 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002625 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002626 sizeof(combinationsobject), /* tp_basicsize */
2627 0, /* tp_itemsize */
2628 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002629 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002630 0, /* tp_print */
2631 0, /* tp_getattr */
2632 0, /* tp_setattr */
2633 0, /* tp_reserved */
2634 0, /* tp_repr */
2635 0, /* tp_as_number */
2636 0, /* tp_as_sequence */
2637 0, /* tp_as_mapping */
2638 0, /* tp_hash */
2639 0, /* tp_call */
2640 0, /* tp_str */
2641 PyObject_GenericGetAttr, /* tp_getattro */
2642 0, /* tp_setattro */
2643 0, /* tp_as_buffer */
2644 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2645 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002646 combinations_doc, /* tp_doc */
2647 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002648 0, /* tp_clear */
2649 0, /* tp_richcompare */
2650 0, /* tp_weaklistoffset */
2651 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002652 (iternextfunc)combinations_next, /* tp_iternext */
2653 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002654 0, /* tp_members */
2655 0, /* tp_getset */
2656 0, /* tp_base */
2657 0, /* tp_dict */
2658 0, /* tp_descr_get */
2659 0, /* tp_descr_set */
2660 0, /* tp_dictoffset */
2661 0, /* tp_init */
2662 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002663 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002664 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002665};
2666
2667
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002668/* combinations with replacement object *******************************************/
2669
2670/* Equivalent to:
2671
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002672 def combinations_with_replacement(iterable, r):
2673 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2674 # number items returned: (n+r-1)! / r! / (n-1)!
2675 pool = tuple(iterable)
2676 n = len(pool)
2677 indices = [0] * r
2678 yield tuple(pool[i] for i in indices)
2679 while 1:
2680 for i in reversed(range(r)):
2681 if indices[i] != n - 1:
2682 break
2683 else:
2684 return
2685 indices[i:] = [indices[i] + 1] * (r - i)
2686 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002687
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002688 def combinations_with_replacement2(iterable, r):
2689 'Alternate version that filters from product()'
2690 pool = tuple(iterable)
2691 n = len(pool)
2692 for indices in product(range(n), repeat=r):
2693 if sorted(indices) == list(indices):
2694 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002695*/
2696typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002697 PyObject_HEAD
2698 PyObject *pool; /* input converted to a tuple */
2699 Py_ssize_t *indices; /* one index per result element */
2700 PyObject *result; /* most recently returned result tuple */
2701 Py_ssize_t r; /* size of result tuple */
2702 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002703} cwrobject;
2704
2705static PyTypeObject cwr_type;
2706
2707static PyObject *
2708cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2709{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002710 cwrobject *co;
2711 Py_ssize_t n;
2712 Py_ssize_t r;
2713 PyObject *pool = NULL;
2714 PyObject *iterable = NULL;
2715 Py_ssize_t *indices = NULL;
2716 Py_ssize_t i;
2717 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002718
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002719 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2720 &iterable, &r))
2721 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002722
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002723 pool = PySequence_Tuple(iterable);
2724 if (pool == NULL)
2725 goto error;
2726 n = PyTuple_GET_SIZE(pool);
2727 if (r < 0) {
2728 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2729 goto error;
2730 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002731
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002732 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002733 if (indices == NULL) {
2734 PyErr_NoMemory();
2735 goto error;
2736 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002737
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002738 for (i=0 ; i<r ; i++)
2739 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002740
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002741 /* create cwrobject structure */
2742 co = (cwrobject *)type->tp_alloc(type, 0);
2743 if (co == NULL)
2744 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002745
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002746 co->pool = pool;
2747 co->indices = indices;
2748 co->result = NULL;
2749 co->r = r;
2750 co->stopped = !n && r;
2751
2752 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002753
2754error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002755 if (indices != NULL)
2756 PyMem_Free(indices);
2757 Py_XDECREF(pool);
2758 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002759}
2760
2761static void
2762cwr_dealloc(cwrobject *co)
2763{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002764 PyObject_GC_UnTrack(co);
2765 Py_XDECREF(co->pool);
2766 Py_XDECREF(co->result);
2767 if (co->indices != NULL)
2768 PyMem_Free(co->indices);
2769 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002770}
2771
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002772static PyObject *
2773cwr_sizeof(cwrobject *co, void *unused)
2774{
2775 Py_ssize_t res;
2776
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002777 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002778 res += co->r * sizeof(Py_ssize_t);
2779 return PyLong_FromSsize_t(res);
2780}
2781
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002782static int
2783cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2784{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002785 Py_VISIT(co->pool);
2786 Py_VISIT(co->result);
2787 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002788}
2789
2790static PyObject *
2791cwr_next(cwrobject *co)
2792{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002793 PyObject *elem;
2794 PyObject *oldelem;
2795 PyObject *pool = co->pool;
2796 Py_ssize_t *indices = co->indices;
2797 PyObject *result = co->result;
2798 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2799 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002800 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002801
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002802 if (co->stopped)
2803 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002804
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002805 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002806 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002807 result = PyTuple_New(r);
2808 if (result == NULL)
2809 goto empty;
2810 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002811 if (n > 0) {
2812 elem = PyTuple_GET_ITEM(pool, 0);
2813 for (i=0; i<r ; i++) {
2814 assert(indices[i] == 0);
2815 Py_INCREF(elem);
2816 PyTuple_SET_ITEM(result, i, elem);
2817 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002818 }
2819 } else {
2820 /* Copy the previous result tuple or re-use it if available */
2821 if (Py_REFCNT(result) > 1) {
2822 PyObject *old_result = result;
2823 result = PyTuple_New(r);
2824 if (result == NULL)
2825 goto empty;
2826 co->result = result;
2827 for (i=0; i<r ; i++) {
2828 elem = PyTuple_GET_ITEM(old_result, i);
2829 Py_INCREF(elem);
2830 PyTuple_SET_ITEM(result, i, elem);
2831 }
2832 Py_DECREF(old_result);
2833 }
2834 /* Now, we've got the only copy so we can update it in-place CPython's
2835 empty tuple is a singleton and cached in PyTuple's freelist. */
2836 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002837
Tim Peters9edb1682013-09-03 11:49:31 -05002838 /* Scan indices right-to-left until finding one that is not
2839 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002840 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2841 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002842
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002843 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002844 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002845 if (i < 0)
2846 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002847
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002848 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002849 maximum. Then set all to the right to the same value. */
2850 index = indices[i] + 1;
2851 assert(index < n);
2852 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002853 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002854 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002855 Py_INCREF(elem);
2856 oldelem = PyTuple_GET_ITEM(result, i);
2857 PyTuple_SET_ITEM(result, i, elem);
2858 Py_DECREF(oldelem);
2859 }
2860 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002861
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002862 Py_INCREF(result);
2863 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002864
2865empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002866 co->stopped = 1;
2867 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002868}
2869
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002870static PyObject *
2871cwr_reduce(cwrobject *lz)
2872{
2873 if (lz->result == NULL) {
2874 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2875 } else if (lz->stopped) {
2876 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2877 } else {
2878 PyObject *indices;
2879 Py_ssize_t i;
2880
2881 /* we must pickle the indices and use them for setstate */
2882 indices = PyTuple_New(lz->r);
2883 if (!indices)
2884 return NULL;
2885 for (i=0; i<lz->r; i++)
2886 {
2887 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2888 if (!index) {
2889 Py_DECREF(indices);
2890 return NULL;
2891 }
2892 PyTuple_SET_ITEM(indices, i, index);
2893 }
2894
2895 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2896 }
2897}
2898
2899static PyObject *
2900cwr_setstate(cwrobject *lz, PyObject *state)
2901{
2902 PyObject *result;
2903 Py_ssize_t n, i;
2904
2905 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2906 {
2907 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2908 return NULL;
2909 }
2910
2911 n = PyTuple_GET_SIZE(lz->pool);
2912 for (i=0; i<lz->r; i++)
2913 {
2914 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2915 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2916 if (index < 0 && PyErr_Occurred())
2917 return NULL; /* not an integer */
2918 /* clamp the index */
2919 if (index < 0)
2920 index = 0;
2921 else if (index > n-1)
2922 index = n-1;
2923 lz->indices[i] = index;
2924 }
2925 result = PyTuple_New(lz->r);
2926 if (result == NULL)
2927 return NULL;
2928 for (i=0; i<lz->r; i++) {
2929 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2930 Py_INCREF(element);
2931 PyTuple_SET_ITEM(result, i, element);
2932 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002933 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002934 Py_RETURN_NONE;
2935}
2936
2937static PyMethodDef cwr_methods[] = {
2938 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2939 reduce_doc},
2940 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2941 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002942 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
2943 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002944 {NULL, NULL} /* sentinel */
2945};
2946
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002947PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002948"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002949\n\
2950Return successive r-length combinations of elements in the iterable\n\
2951allowing individual elements to have successive repeats.\n\
2952combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2953
2954static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002955 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002956 "itertools.combinations_with_replacement", /* tp_name */
2957 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002958 0, /* tp_itemsize */
2959 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002960 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002961 0, /* tp_print */
2962 0, /* tp_getattr */
2963 0, /* tp_setattr */
2964 0, /* tp_reserved */
2965 0, /* tp_repr */
2966 0, /* tp_as_number */
2967 0, /* tp_as_sequence */
2968 0, /* tp_as_mapping */
2969 0, /* tp_hash */
2970 0, /* tp_call */
2971 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002972 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002973 0, /* tp_setattro */
2974 0, /* tp_as_buffer */
2975 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002976 Py_TPFLAGS_BASETYPE, /* tp_flags */
2977 cwr_doc, /* tp_doc */
2978 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002979 0, /* tp_clear */
2980 0, /* tp_richcompare */
2981 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002982 PyObject_SelfIter, /* tp_iter */
2983 (iternextfunc)cwr_next, /* tp_iternext */
2984 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002985 0, /* tp_members */
2986 0, /* tp_getset */
2987 0, /* tp_base */
2988 0, /* tp_dict */
2989 0, /* tp_descr_get */
2990 0, /* tp_descr_set */
2991 0, /* tp_dictoffset */
2992 0, /* tp_init */
2993 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002994 cwr_new, /* tp_new */
2995 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002996};
2997
2998
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002999/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003000
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003001def permutations(iterable, r=None):
3002 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
3003 pool = tuple(iterable)
3004 n = len(pool)
3005 r = n if r is None else r
3006 indices = range(n)
3007 cycles = range(n-r+1, n+1)[::-1]
3008 yield tuple(pool[i] for i in indices[:r])
3009 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003010 for i in reversed(range(r)):
3011 cycles[i] -= 1
3012 if cycles[i] == 0:
3013 indices[i:] = indices[i+1:] + indices[i:i+1]
3014 cycles[i] = n - i
3015 else:
3016 j = cycles[i]
3017 indices[i], indices[-j] = indices[-j], indices[i]
3018 yield tuple(pool[i] for i in indices[:r])
3019 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003020 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003021 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003022*/
3023
3024typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003025 PyObject_HEAD
3026 PyObject *pool; /* input converted to a tuple */
3027 Py_ssize_t *indices; /* one index per element in the pool */
3028 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3029 PyObject *result; /* most recently returned result tuple */
3030 Py_ssize_t r; /* size of result tuple */
3031 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003032} permutationsobject;
3033
3034static PyTypeObject permutations_type;
3035
3036static PyObject *
3037permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3038{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003039 permutationsobject *po;
3040 Py_ssize_t n;
3041 Py_ssize_t r;
3042 PyObject *robj = Py_None;
3043 PyObject *pool = NULL;
3044 PyObject *iterable = NULL;
3045 Py_ssize_t *indices = NULL;
3046 Py_ssize_t *cycles = NULL;
3047 Py_ssize_t i;
3048 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003049
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003050 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
3051 &iterable, &robj))
3052 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003053
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003054 pool = PySequence_Tuple(iterable);
3055 if (pool == NULL)
3056 goto error;
3057 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003058
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003059 r = n;
3060 if (robj != Py_None) {
3061 if (!PyLong_Check(robj)) {
3062 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3063 goto error;
3064 }
3065 r = PyLong_AsSsize_t(robj);
3066 if (r == -1 && PyErr_Occurred())
3067 goto error;
3068 }
3069 if (r < 0) {
3070 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3071 goto error;
3072 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003073
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003074 indices = PyMem_New(Py_ssize_t, n);
3075 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003076 if (indices == NULL || cycles == NULL) {
3077 PyErr_NoMemory();
3078 goto error;
3079 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003080
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003081 for (i=0 ; i<n ; i++)
3082 indices[i] = i;
3083 for (i=0 ; i<r ; i++)
3084 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003085
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003086 /* create permutationsobject structure */
3087 po = (permutationsobject *)type->tp_alloc(type, 0);
3088 if (po == NULL)
3089 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003090
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003091 po->pool = pool;
3092 po->indices = indices;
3093 po->cycles = cycles;
3094 po->result = NULL;
3095 po->r = r;
3096 po->stopped = r > n ? 1 : 0;
3097
3098 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003099
3100error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003101 if (indices != NULL)
3102 PyMem_Free(indices);
3103 if (cycles != NULL)
3104 PyMem_Free(cycles);
3105 Py_XDECREF(pool);
3106 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003107}
3108
3109static void
3110permutations_dealloc(permutationsobject *po)
3111{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003112 PyObject_GC_UnTrack(po);
3113 Py_XDECREF(po->pool);
3114 Py_XDECREF(po->result);
3115 PyMem_Free(po->indices);
3116 PyMem_Free(po->cycles);
3117 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003118}
3119
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003120static PyObject *
3121permutations_sizeof(permutationsobject *po, void *unused)
3122{
3123 Py_ssize_t res;
3124
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003125 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003126 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3127 res += po->r * sizeof(Py_ssize_t);
3128 return PyLong_FromSsize_t(res);
3129}
3130
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003131static int
3132permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3133{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003134 Py_VISIT(po->pool);
3135 Py_VISIT(po->result);
3136 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003137}
3138
3139static PyObject *
3140permutations_next(permutationsobject *po)
3141{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003142 PyObject *elem;
3143 PyObject *oldelem;
3144 PyObject *pool = po->pool;
3145 Py_ssize_t *indices = po->indices;
3146 Py_ssize_t *cycles = po->cycles;
3147 PyObject *result = po->result;
3148 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3149 Py_ssize_t r = po->r;
3150 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003151
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003152 if (po->stopped)
3153 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003154
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003155 if (result == NULL) {
3156 /* On the first pass, initialize result tuple using the indices */
3157 result = PyTuple_New(r);
3158 if (result == NULL)
3159 goto empty;
3160 po->result = result;
3161 for (i=0; i<r ; i++) {
3162 index = indices[i];
3163 elem = PyTuple_GET_ITEM(pool, index);
3164 Py_INCREF(elem);
3165 PyTuple_SET_ITEM(result, i, elem);
3166 }
3167 } else {
3168 if (n == 0)
3169 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003170
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003171 /* Copy the previous result tuple or re-use it if available */
3172 if (Py_REFCNT(result) > 1) {
3173 PyObject *old_result = result;
3174 result = PyTuple_New(r);
3175 if (result == NULL)
3176 goto empty;
3177 po->result = result;
3178 for (i=0; i<r ; i++) {
3179 elem = PyTuple_GET_ITEM(old_result, i);
3180 Py_INCREF(elem);
3181 PyTuple_SET_ITEM(result, i, elem);
3182 }
3183 Py_DECREF(old_result);
3184 }
3185 /* Now, we've got the only copy so we can update it in-place */
3186 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003187
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003188 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3189 for (i=r-1 ; i>=0 ; i--) {
3190 cycles[i] -= 1;
3191 if (cycles[i] == 0) {
3192 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3193 index = indices[i];
3194 for (j=i ; j<n-1 ; j++)
3195 indices[j] = indices[j+1];
3196 indices[n-1] = index;
3197 cycles[i] = n - i;
3198 } else {
3199 j = cycles[i];
3200 index = indices[i];
3201 indices[i] = indices[n-j];
3202 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003203
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003204 for (k=i; k<r ; k++) {
3205 /* start with i, the leftmost element that changed */
3206 /* yield tuple(pool[k] for k in indices[:r]) */
3207 index = indices[k];
3208 elem = PyTuple_GET_ITEM(pool, index);
3209 Py_INCREF(elem);
3210 oldelem = PyTuple_GET_ITEM(result, k);
3211 PyTuple_SET_ITEM(result, k, elem);
3212 Py_DECREF(oldelem);
3213 }
3214 break;
3215 }
3216 }
3217 /* If i is negative, then the cycles have all
3218 rolled-over and we're done. */
3219 if (i < 0)
3220 goto empty;
3221 }
3222 Py_INCREF(result);
3223 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003224
3225empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003226 po->stopped = 1;
3227 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003228}
3229
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003230static PyObject *
3231permutations_reduce(permutationsobject *po)
3232{
3233 if (po->result == NULL) {
3234 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3235 } else if (po->stopped) {
3236 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3237 } else {
3238 PyObject *indices=NULL, *cycles=NULL;
3239 Py_ssize_t n, i;
3240
3241 /* we must pickle the indices and cycles and use them for setstate */
3242 n = PyTuple_GET_SIZE(po->pool);
3243 indices = PyTuple_New(n);
3244 if (indices == NULL)
3245 goto err;
3246 for (i=0; i<n; i++){
3247 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3248 if (!index)
3249 goto err;
3250 PyTuple_SET_ITEM(indices, i, index);
3251 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003252
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003253 cycles = PyTuple_New(po->r);
3254 if (cycles == NULL)
3255 goto err;
3256 for (i=0; i<po->r; i++)
3257 {
3258 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3259 if (!index)
3260 goto err;
3261 PyTuple_SET_ITEM(cycles, i, index);
3262 }
3263 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3264 po->pool, po->r,
3265 indices, cycles);
3266 err:
3267 Py_XDECREF(indices);
3268 Py_XDECREF(cycles);
3269 return NULL;
3270 }
3271}
3272
3273static PyObject *
3274permutations_setstate(permutationsobject *po, PyObject *state)
3275{
3276 PyObject *indices, *cycles, *result;
3277 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003278
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003279 if (!PyTuple_Check(state)) {
3280 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3281 return NULL;
3282 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003283 if (!PyArg_ParseTuple(state, "O!O!",
3284 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003285 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003286 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003287 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003288
3289 n = PyTuple_GET_SIZE(po->pool);
3290 if (PyTuple_GET_SIZE(indices) != n ||
3291 PyTuple_GET_SIZE(cycles) != po->r)
3292 {
3293 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3294 return NULL;
3295 }
3296
3297 for (i=0; i<n; i++)
3298 {
3299 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3300 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3301 if (index < 0 && PyErr_Occurred())
3302 return NULL; /* not an integer */
3303 /* clamp the index */
3304 if (index < 0)
3305 index = 0;
3306 else if (index > n-1)
3307 index = n-1;
3308 po->indices[i] = index;
3309 }
3310
3311 for (i=0; i<po->r; i++)
3312 {
3313 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3314 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3315 if (index < 0 && PyErr_Occurred())
3316 return NULL; /* not an integer */
3317 if (index < 1)
3318 index = 1;
3319 else if (index > n-i)
3320 index = n-i;
3321 po->cycles[i] = index;
3322 }
3323 result = PyTuple_New(po->r);
3324 if (result == NULL)
3325 return NULL;
3326 for (i=0; i<po->r; i++) {
3327 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3328 Py_INCREF(element);
3329 PyTuple_SET_ITEM(result, i, element);
3330 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003331 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003332 Py_RETURN_NONE;
3333}
3334
3335static PyMethodDef permuations_methods[] = {
3336 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3337 reduce_doc},
3338 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3339 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003340 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3341 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003342 {NULL, NULL} /* sentinel */
3343};
3344
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003345PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003346"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003347\n\
3348Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003349permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003350
3351static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003352 PyVarObject_HEAD_INIT(NULL, 0)
3353 "itertools.permutations", /* tp_name */
3354 sizeof(permutationsobject), /* tp_basicsize */
3355 0, /* tp_itemsize */
3356 /* methods */
3357 (destructor)permutations_dealloc, /* tp_dealloc */
3358 0, /* tp_print */
3359 0, /* tp_getattr */
3360 0, /* tp_setattr */
3361 0, /* tp_reserved */
3362 0, /* tp_repr */
3363 0, /* tp_as_number */
3364 0, /* tp_as_sequence */
3365 0, /* tp_as_mapping */
3366 0, /* tp_hash */
3367 0, /* tp_call */
3368 0, /* tp_str */
3369 PyObject_GenericGetAttr, /* tp_getattro */
3370 0, /* tp_setattro */
3371 0, /* tp_as_buffer */
3372 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3373 Py_TPFLAGS_BASETYPE, /* tp_flags */
3374 permutations_doc, /* tp_doc */
3375 (traverseproc)permutations_traverse, /* tp_traverse */
3376 0, /* tp_clear */
3377 0, /* tp_richcompare */
3378 0, /* tp_weaklistoffset */
3379 PyObject_SelfIter, /* tp_iter */
3380 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003381 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003382 0, /* tp_members */
3383 0, /* tp_getset */
3384 0, /* tp_base */
3385 0, /* tp_dict */
3386 0, /* tp_descr_get */
3387 0, /* tp_descr_set */
3388 0, /* tp_dictoffset */
3389 0, /* tp_init */
3390 0, /* tp_alloc */
3391 permutations_new, /* tp_new */
3392 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003393};
3394
Raymond Hettinger482ba772010-12-01 22:48:00 +00003395/* accumulate object ************************************************************/
3396
3397typedef struct {
3398 PyObject_HEAD
3399 PyObject *total;
3400 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003401 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003402} accumulateobject;
3403
3404static PyTypeObject accumulate_type;
3405
3406static PyObject *
3407accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3408{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003409 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003410 PyObject *iterable;
3411 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003412 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003413 accumulateobject *lz;
3414
Raymond Hettinger5d446132011-03-27 18:52:10 -07003415 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3416 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003417 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003418
3419 /* Get iterator. */
3420 it = PyObject_GetIter(iterable);
3421 if (it == NULL)
3422 return NULL;
3423
Raymond Hettinger482ba772010-12-01 22:48:00 +00003424 /* create accumulateobject structure */
3425 lz = (accumulateobject *)type->tp_alloc(type, 0);
3426 if (lz == NULL) {
3427 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003428 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003429 }
3430
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003431 if (binop != Py_None) {
3432 Py_XINCREF(binop);
3433 lz->binop = binop;
3434 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003435 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003436 lz->it = it;
3437 return (PyObject *)lz;
3438}
3439
3440static void
3441accumulate_dealloc(accumulateobject *lz)
3442{
3443 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003444 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003445 Py_XDECREF(lz->total);
3446 Py_XDECREF(lz->it);
3447 Py_TYPE(lz)->tp_free(lz);
3448}
3449
3450static int
3451accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3452{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003453 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003454 Py_VISIT(lz->it);
3455 Py_VISIT(lz->total);
3456 return 0;
3457}
3458
3459static PyObject *
3460accumulate_next(accumulateobject *lz)
3461{
3462 PyObject *val, *oldtotal, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003463
Raymond Hettinger482ba772010-12-01 22:48:00 +00003464 val = PyIter_Next(lz->it);
3465 if (val == NULL)
3466 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003467
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003468 if (lz->total == NULL) {
3469 Py_INCREF(val);
3470 lz->total = val;
3471 return lz->total;
3472 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003473
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003474 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003475 newtotal = PyNumber_Add(lz->total, val);
3476 else
3477 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003478 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003479 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003480 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003481
3482 oldtotal = lz->total;
3483 lz->total = newtotal;
3484 Py_DECREF(oldtotal);
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003485
Raymond Hettinger482ba772010-12-01 22:48:00 +00003486 Py_INCREF(newtotal);
3487 return newtotal;
3488}
3489
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003490static PyObject *
3491accumulate_reduce(accumulateobject *lz)
3492{
Serhiy Storchakad5516252016-03-06 14:00:45 +02003493 if (lz->total == Py_None) {
3494 PyObject *it;
3495
3496 if (PyType_Ready(&chain_type) < 0)
3497 return NULL;
3498 if (PyType_Ready(&islice_type) < 0)
3499 return NULL;
3500 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3501 lz->total, lz->it);
3502 if (it == NULL)
3503 return NULL;
3504 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3505 it, lz->binop ? lz->binop : Py_None);
3506 if (it == NULL)
3507 return NULL;
3508 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3509 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003510 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3511 lz->it, lz->binop?lz->binop:Py_None,
3512 lz->total?lz->total:Py_None);
3513 }
3514
3515static PyObject *
3516accumulate_setstate(accumulateobject *lz, PyObject *state)
3517{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003518 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003519 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003520 Py_RETURN_NONE;
3521}
3522
3523static PyMethodDef accumulate_methods[] = {
3524 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3525 reduce_doc},
3526 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3527 setstate_doc},
3528 {NULL, NULL} /* sentinel */
3529};
3530
Raymond Hettinger482ba772010-12-01 22:48:00 +00003531PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003532"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003533\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003534Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003535
3536static PyTypeObject accumulate_type = {
3537 PyVarObject_HEAD_INIT(NULL, 0)
3538 "itertools.accumulate", /* tp_name */
3539 sizeof(accumulateobject), /* tp_basicsize */
3540 0, /* tp_itemsize */
3541 /* methods */
3542 (destructor)accumulate_dealloc, /* tp_dealloc */
3543 0, /* tp_print */
3544 0, /* tp_getattr */
3545 0, /* tp_setattr */
3546 0, /* tp_reserved */
3547 0, /* tp_repr */
3548 0, /* tp_as_number */
3549 0, /* tp_as_sequence */
3550 0, /* tp_as_mapping */
3551 0, /* tp_hash */
3552 0, /* tp_call */
3553 0, /* tp_str */
3554 PyObject_GenericGetAttr, /* tp_getattro */
3555 0, /* tp_setattro */
3556 0, /* tp_as_buffer */
3557 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3558 Py_TPFLAGS_BASETYPE, /* tp_flags */
3559 accumulate_doc, /* tp_doc */
3560 (traverseproc)accumulate_traverse, /* tp_traverse */
3561 0, /* tp_clear */
3562 0, /* tp_richcompare */
3563 0, /* tp_weaklistoffset */
3564 PyObject_SelfIter, /* tp_iter */
3565 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003566 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003567 0, /* tp_members */
3568 0, /* tp_getset */
3569 0, /* tp_base */
3570 0, /* tp_dict */
3571 0, /* tp_descr_get */
3572 0, /* tp_descr_set */
3573 0, /* tp_dictoffset */
3574 0, /* tp_init */
3575 0, /* tp_alloc */
3576 accumulate_new, /* tp_new */
3577 PyObject_GC_Del, /* tp_free */
3578};
3579
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003580
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003581/* compress object ************************************************************/
3582
3583/* Equivalent to:
3584
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003585 def compress(data, selectors):
3586 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3587 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003588*/
3589
3590typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003591 PyObject_HEAD
3592 PyObject *data;
3593 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003594} compressobject;
3595
3596static PyTypeObject compress_type;
3597
3598static PyObject *
3599compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3600{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003601 PyObject *seq1, *seq2;
3602 PyObject *data=NULL, *selectors=NULL;
3603 compressobject *lz;
3604 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003605
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003606 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3607 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003608
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003609 data = PyObject_GetIter(seq1);
3610 if (data == NULL)
3611 goto fail;
3612 selectors = PyObject_GetIter(seq2);
3613 if (selectors == NULL)
3614 goto fail;
3615
3616 /* create compressobject structure */
3617 lz = (compressobject *)type->tp_alloc(type, 0);
3618 if (lz == NULL)
3619 goto fail;
3620 lz->data = data;
3621 lz->selectors = selectors;
3622 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003623
3624fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003625 Py_XDECREF(data);
3626 Py_XDECREF(selectors);
3627 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003628}
3629
3630static void
3631compress_dealloc(compressobject *lz)
3632{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003633 PyObject_GC_UnTrack(lz);
3634 Py_XDECREF(lz->data);
3635 Py_XDECREF(lz->selectors);
3636 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003637}
3638
3639static int
3640compress_traverse(compressobject *lz, visitproc visit, void *arg)
3641{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003642 Py_VISIT(lz->data);
3643 Py_VISIT(lz->selectors);
3644 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003645}
3646
3647static PyObject *
3648compress_next(compressobject *lz)
3649{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003650 PyObject *data = lz->data, *selectors = lz->selectors;
3651 PyObject *datum, *selector;
3652 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3653 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3654 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003655
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003656 while (1) {
3657 /* Steps: get datum, get selector, evaluate selector.
3658 Order is important (to match the pure python version
3659 in terms of which input gets a chance to raise an
3660 exception first).
3661 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003662
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003663 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003664 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003665 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003666
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003667 selector = selectornext(selectors);
3668 if (selector == NULL) {
3669 Py_DECREF(datum);
3670 return NULL;
3671 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003672
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003673 ok = PyObject_IsTrue(selector);
3674 Py_DECREF(selector);
3675 if (ok == 1)
3676 return datum;
3677 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003678 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003679 return NULL;
3680 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003681}
3682
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003683static PyObject *
3684compress_reduce(compressobject *lz)
3685{
3686 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3687 lz->data, lz->selectors);
3688 }
3689
3690static PyMethodDef compress_methods[] = {
3691 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3692 reduce_doc},
3693 {NULL, NULL} /* sentinel */
3694};
3695
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003696PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003697"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003698\n\
3699Return data elements corresponding to true selector elements.\n\
3700Forms a shorter iterator from selected data elements using the\n\
3701selectors to choose the data elements.");
3702
3703static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003704 PyVarObject_HEAD_INIT(NULL, 0)
3705 "itertools.compress", /* tp_name */
3706 sizeof(compressobject), /* tp_basicsize */
3707 0, /* tp_itemsize */
3708 /* methods */
3709 (destructor)compress_dealloc, /* tp_dealloc */
3710 0, /* tp_print */
3711 0, /* tp_getattr */
3712 0, /* tp_setattr */
3713 0, /* tp_reserved */
3714 0, /* tp_repr */
3715 0, /* tp_as_number */
3716 0, /* tp_as_sequence */
3717 0, /* tp_as_mapping */
3718 0, /* tp_hash */
3719 0, /* tp_call */
3720 0, /* tp_str */
3721 PyObject_GenericGetAttr, /* tp_getattro */
3722 0, /* tp_setattro */
3723 0, /* tp_as_buffer */
3724 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3725 Py_TPFLAGS_BASETYPE, /* tp_flags */
3726 compress_doc, /* tp_doc */
3727 (traverseproc)compress_traverse, /* tp_traverse */
3728 0, /* tp_clear */
3729 0, /* tp_richcompare */
3730 0, /* tp_weaklistoffset */
3731 PyObject_SelfIter, /* tp_iter */
3732 (iternextfunc)compress_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003733 compress_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003734 0, /* tp_members */
3735 0, /* tp_getset */
3736 0, /* tp_base */
3737 0, /* tp_dict */
3738 0, /* tp_descr_get */
3739 0, /* tp_descr_set */
3740 0, /* tp_dictoffset */
3741 0, /* tp_init */
3742 0, /* tp_alloc */
3743 compress_new, /* tp_new */
3744 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003745};
3746
3747
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003748/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003749
3750typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003751 PyObject_HEAD
3752 PyObject *func;
3753 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003754} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003755
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003756static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003757
3758static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003759filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003760{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003761 PyObject *func, *seq;
3762 PyObject *it;
3763 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003764
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003765 if (type == &filterfalse_type &&
3766 !_PyArg_NoKeywords("filterfalse()", kwds))
3767 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003768
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003769 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3770 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003771
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003772 /* Get iterator. */
3773 it = PyObject_GetIter(seq);
3774 if (it == NULL)
3775 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003776
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003777 /* create filterfalseobject structure */
3778 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3779 if (lz == NULL) {
3780 Py_DECREF(it);
3781 return NULL;
3782 }
3783 Py_INCREF(func);
3784 lz->func = func;
3785 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003786
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003787 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003788}
3789
3790static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003791filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003792{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003793 PyObject_GC_UnTrack(lz);
3794 Py_XDECREF(lz->func);
3795 Py_XDECREF(lz->it);
3796 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003797}
3798
3799static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003800filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003801{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003802 Py_VISIT(lz->it);
3803 Py_VISIT(lz->func);
3804 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003805}
3806
3807static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003808filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003809{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003810 PyObject *item;
3811 PyObject *it = lz->it;
3812 long ok;
3813 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003814
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003815 iternext = *Py_TYPE(it)->tp_iternext;
3816 for (;;) {
3817 item = iternext(it);
3818 if (item == NULL)
3819 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003820
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003821 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3822 ok = PyObject_IsTrue(item);
3823 } else {
3824 PyObject *good;
3825 good = PyObject_CallFunctionObjArgs(lz->func,
3826 item, NULL);
3827 if (good == NULL) {
3828 Py_DECREF(item);
3829 return NULL;
3830 }
3831 ok = PyObject_IsTrue(good);
3832 Py_DECREF(good);
3833 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003834 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003835 return item;
3836 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003837 if (ok < 0)
3838 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003839 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003840}
3841
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003842static PyObject *
3843filterfalse_reduce(filterfalseobject *lz)
3844{
3845 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3846 lz->func, lz->it);
3847 }
3848
3849static PyMethodDef filterfalse_methods[] = {
3850 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3851 reduce_doc},
3852 {NULL, NULL} /* sentinel */
3853};
3854
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003855PyDoc_STRVAR(filterfalse_doc,
3856"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003857\n\
3858Return those items of sequence for which function(item) is false.\n\
3859If function is None, return the items that are false.");
3860
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003861static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003862 PyVarObject_HEAD_INIT(NULL, 0)
3863 "itertools.filterfalse", /* tp_name */
3864 sizeof(filterfalseobject), /* tp_basicsize */
3865 0, /* tp_itemsize */
3866 /* methods */
3867 (destructor)filterfalse_dealloc, /* tp_dealloc */
3868 0, /* tp_print */
3869 0, /* tp_getattr */
3870 0, /* tp_setattr */
3871 0, /* tp_reserved */
3872 0, /* tp_repr */
3873 0, /* tp_as_number */
3874 0, /* tp_as_sequence */
3875 0, /* tp_as_mapping */
3876 0, /* tp_hash */
3877 0, /* tp_call */
3878 0, /* tp_str */
3879 PyObject_GenericGetAttr, /* tp_getattro */
3880 0, /* tp_setattro */
3881 0, /* tp_as_buffer */
3882 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3883 Py_TPFLAGS_BASETYPE, /* tp_flags */
3884 filterfalse_doc, /* tp_doc */
3885 (traverseproc)filterfalse_traverse, /* tp_traverse */
3886 0, /* tp_clear */
3887 0, /* tp_richcompare */
3888 0, /* tp_weaklistoffset */
3889 PyObject_SelfIter, /* tp_iter */
3890 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003891 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003892 0, /* tp_members */
3893 0, /* tp_getset */
3894 0, /* tp_base */
3895 0, /* tp_dict */
3896 0, /* tp_descr_get */
3897 0, /* tp_descr_set */
3898 0, /* tp_dictoffset */
3899 0, /* tp_init */
3900 0, /* tp_alloc */
3901 filterfalse_new, /* tp_new */
3902 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003903};
3904
3905
3906/* count object ************************************************************/
3907
3908typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003909 PyObject_HEAD
3910 Py_ssize_t cnt;
3911 PyObject *long_cnt;
3912 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003913} countobject;
3914
Raymond Hettinger30729212009-02-12 06:28:27 +00003915/* Counting logic and invariants:
3916
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003917fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003918
Serhiy Storchaka483405b2015-02-17 10:14:30 +02003919 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003920 Advances with: cnt += 1
3921 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003922
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003923slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003924
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003925 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3926 All counting is done with python objects (no overflows or underflows).
3927 Advances with: long_cnt += long_step
3928 Step may be zero -- effectively a slow version of repeat(cnt).
3929 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003930*/
3931
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003932static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003933
3934static PyObject *
3935count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3936{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003937 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003938 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003939 Py_ssize_t cnt = 0;
3940 PyObject *long_cnt = NULL;
3941 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003942 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003943 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003944
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003945 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3946 kwlist, &long_cnt, &long_step))
3947 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003948
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003949 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3950 (long_step != NULL && !PyNumber_Check(long_step))) {
3951 PyErr_SetString(PyExc_TypeError, "a number is required");
3952 return NULL;
3953 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003954
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003955 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
3956 (long_step == NULL || PyLong_Check(long_step));
3957
3958 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003959 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003960 if (fast_mode) {
3961 assert(PyLong_Check(long_cnt));
3962 cnt = PyLong_AsSsize_t(long_cnt);
3963 if (cnt == -1 && PyErr_Occurred()) {
3964 PyErr_Clear();
3965 fast_mode = 0;
3966 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003967 }
3968 Py_INCREF(long_cnt);
3969 } else {
3970 cnt = 0;
3971 long_cnt = PyLong_FromLong(0);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003972 if (long_cnt == NULL) {
3973 return NULL;
3974 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003975 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003976
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003977 /* If not specified, step defaults to 1 */
3978 if (long_step == NULL) {
3979 long_step = PyLong_FromLong(1);
3980 if (long_step == NULL) {
3981 Py_DECREF(long_cnt);
3982 return NULL;
3983 }
3984 } else
3985 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003986
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003987 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003988
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003989 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003990 if (fast_mode) {
3991 assert(PyLong_Check(long_step));
3992 step = PyLong_AsLong(long_step);
3993 if (step != 1) {
3994 fast_mode = 0;
3995 if (step == -1 && PyErr_Occurred())
3996 PyErr_Clear();
3997 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003998 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003999
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004000 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004001 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004002 else
4003 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004004
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004005 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4006 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4007 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004008 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004009
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004010 /* create countobject structure */
4011 lz = (countobject *)type->tp_alloc(type, 0);
4012 if (lz == NULL) {
4013 Py_XDECREF(long_cnt);
4014 return NULL;
4015 }
4016 lz->cnt = cnt;
4017 lz->long_cnt = long_cnt;
4018 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004019
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004020 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004021}
4022
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004023static void
4024count_dealloc(countobject *lz)
4025{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004026 PyObject_GC_UnTrack(lz);
4027 Py_XDECREF(lz->long_cnt);
4028 Py_XDECREF(lz->long_step);
4029 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004030}
4031
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004032static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004033count_traverse(countobject *lz, visitproc visit, void *arg)
4034{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004035 Py_VISIT(lz->long_cnt);
4036 Py_VISIT(lz->long_step);
4037 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004038}
4039
4040static PyObject *
4041count_nextlong(countobject *lz)
4042{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004043 PyObject *long_cnt;
4044 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004045
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004046 long_cnt = lz->long_cnt;
4047 if (long_cnt == NULL) {
4048 /* Switch to slow_mode */
4049 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4050 if (long_cnt == NULL)
4051 return NULL;
4052 }
4053 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004054
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004055 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4056 if (stepped_up == NULL)
4057 return NULL;
4058 lz->long_cnt = stepped_up;
4059 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004060}
4061
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004062static PyObject *
4063count_next(countobject *lz)
4064{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004065 if (lz->cnt == PY_SSIZE_T_MAX)
4066 return count_nextlong(lz);
4067 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004068}
4069
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004070static PyObject *
4071count_repr(countobject *lz)
4072{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004073 if (lz->cnt != PY_SSIZE_T_MAX)
4074 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004075
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004076 if (PyLong_Check(lz->long_step)) {
4077 long step = PyLong_AsLong(lz->long_step);
4078 if (step == -1 && PyErr_Occurred()) {
4079 PyErr_Clear();
4080 }
4081 if (step == 1) {
4082 /* Don't display step when it is an integer equal to 1 */
4083 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
4084 }
4085 }
4086 return PyUnicode_FromFormat("count(%R, %R)",
4087 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004088}
4089
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004090static PyObject *
4091count_reduce(countobject *lz)
4092{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004093 if (lz->cnt == PY_SSIZE_T_MAX)
4094 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4095 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004096}
4097
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004098static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004099 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004100 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004101 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004102};
4103
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004104PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00004105 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004106\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004107Return a count object whose .__next__() method returns consecutive values.\n\
4108Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00004109 def count(firstval=0, step=1):\n\
Eli Benderskyc554f722013-09-03 06:37:19 -07004110 x = firstval\n\
4111 while 1:\n\
4112 yield x\n\
4113 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004114
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004115static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004116 PyVarObject_HEAD_INIT(NULL, 0)
4117 "itertools.count", /* tp_name */
4118 sizeof(countobject), /* tp_basicsize */
4119 0, /* tp_itemsize */
4120 /* methods */
4121 (destructor)count_dealloc, /* tp_dealloc */
4122 0, /* tp_print */
4123 0, /* tp_getattr */
4124 0, /* tp_setattr */
4125 0, /* tp_reserved */
4126 (reprfunc)count_repr, /* tp_repr */
4127 0, /* tp_as_number */
4128 0, /* tp_as_sequence */
4129 0, /* tp_as_mapping */
4130 0, /* tp_hash */
4131 0, /* tp_call */
4132 0, /* tp_str */
4133 PyObject_GenericGetAttr, /* tp_getattro */
4134 0, /* tp_setattro */
4135 0, /* tp_as_buffer */
4136 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4137 Py_TPFLAGS_BASETYPE, /* tp_flags */
4138 count_doc, /* tp_doc */
4139 (traverseproc)count_traverse, /* tp_traverse */
4140 0, /* tp_clear */
4141 0, /* tp_richcompare */
4142 0, /* tp_weaklistoffset */
4143 PyObject_SelfIter, /* tp_iter */
4144 (iternextfunc)count_next, /* tp_iternext */
4145 count_methods, /* tp_methods */
4146 0, /* tp_members */
4147 0, /* tp_getset */
4148 0, /* tp_base */
4149 0, /* tp_dict */
4150 0, /* tp_descr_get */
4151 0, /* tp_descr_set */
4152 0, /* tp_dictoffset */
4153 0, /* tp_init */
4154 0, /* tp_alloc */
4155 count_new, /* tp_new */
4156 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004157};
4158
4159
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004160/* repeat object ************************************************************/
4161
4162typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004163 PyObject_HEAD
4164 PyObject *element;
4165 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004166} repeatobject;
4167
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004168static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004169
4170static PyObject *
4171repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4172{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004173 repeatobject *ro;
4174 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004175 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004176 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004177
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004178 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4179 &element, &cnt))
4180 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004181
Raymond Hettinger97d35552014-06-24 21:36:58 -07004182 if (kwds != NULL)
4183 n_kwds = PyDict_Size(kwds);
4184 /* Does user supply times argument? */
4185 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004186 cnt = 0;
4187
4188 ro = (repeatobject *)type->tp_alloc(type, 0);
4189 if (ro == NULL)
4190 return NULL;
4191 Py_INCREF(element);
4192 ro->element = element;
4193 ro->cnt = cnt;
4194 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004195}
4196
4197static void
4198repeat_dealloc(repeatobject *ro)
4199{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004200 PyObject_GC_UnTrack(ro);
4201 Py_XDECREF(ro->element);
4202 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004203}
4204
4205static int
4206repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4207{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004208 Py_VISIT(ro->element);
4209 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004210}
4211
4212static PyObject *
4213repeat_next(repeatobject *ro)
4214{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004215 if (ro->cnt == 0)
4216 return NULL;
4217 if (ro->cnt > 0)
4218 ro->cnt--;
4219 Py_INCREF(ro->element);
4220 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004221}
4222
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004223static PyObject *
4224repeat_repr(repeatobject *ro)
4225{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004226 if (ro->cnt == -1)
4227 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4228 else
4229 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4230}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004231
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004232static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004233repeat_len(repeatobject *ro)
4234{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004235 if (ro->cnt == -1) {
4236 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4237 return NULL;
4238 }
4239 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004240}
4241
Armin Rigof5b3e362006-02-11 21:32:43 +00004242PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004243
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004244static PyObject *
4245repeat_reduce(repeatobject *ro)
4246{
4247 /* unpickle this so that a new repeat iterator is constructed with an
4248 * object, then call __setstate__ on it to set cnt
4249 */
4250 if (ro->cnt >= 0)
4251 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4252 else
4253 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4254}
4255
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004256static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004257 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004258 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004259 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004260};
4261
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004262PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004263"repeat(object [,times]) -> create an iterator which returns the object\n\
4264for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004265endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004266
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004267static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004268 PyVarObject_HEAD_INIT(NULL, 0)
4269 "itertools.repeat", /* tp_name */
4270 sizeof(repeatobject), /* tp_basicsize */
4271 0, /* tp_itemsize */
4272 /* methods */
4273 (destructor)repeat_dealloc, /* tp_dealloc */
4274 0, /* tp_print */
4275 0, /* tp_getattr */
4276 0, /* tp_setattr */
4277 0, /* tp_reserved */
4278 (reprfunc)repeat_repr, /* tp_repr */
4279 0, /* tp_as_number */
4280 0, /* tp_as_sequence */
4281 0, /* tp_as_mapping */
4282 0, /* tp_hash */
4283 0, /* tp_call */
4284 0, /* tp_str */
4285 PyObject_GenericGetAttr, /* tp_getattro */
4286 0, /* tp_setattro */
4287 0, /* tp_as_buffer */
4288 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4289 Py_TPFLAGS_BASETYPE, /* tp_flags */
4290 repeat_doc, /* tp_doc */
4291 (traverseproc)repeat_traverse, /* tp_traverse */
4292 0, /* tp_clear */
4293 0, /* tp_richcompare */
4294 0, /* tp_weaklistoffset */
4295 PyObject_SelfIter, /* tp_iter */
4296 (iternextfunc)repeat_next, /* tp_iternext */
4297 repeat_methods, /* tp_methods */
4298 0, /* tp_members */
4299 0, /* tp_getset */
4300 0, /* tp_base */
4301 0, /* tp_dict */
4302 0, /* tp_descr_get */
4303 0, /* tp_descr_set */
4304 0, /* tp_dictoffset */
4305 0, /* tp_init */
4306 0, /* tp_alloc */
4307 repeat_new, /* tp_new */
4308 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004309};
4310
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004311/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004312
4313#include "Python.h"
4314
4315typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004316 PyObject_HEAD
4317 Py_ssize_t tuplesize;
4318 Py_ssize_t numactive;
4319 PyObject *ittuple; /* tuple of iterators */
4320 PyObject *result;
4321 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004322} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004323
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004324static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004325
4326static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004327zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004328{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004329 ziplongestobject *lz;
4330 Py_ssize_t i;
4331 PyObject *ittuple; /* tuple of iterators */
4332 PyObject *result;
4333 PyObject *fillvalue = Py_None;
4334 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004335
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004336 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
4337 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4338 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
4339 PyErr_SetString(PyExc_TypeError,
4340 "zip_longest() got an unexpected keyword argument");
4341 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004342 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004343 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004344
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004345 /* args must be a tuple */
4346 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00004347
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004348 /* obtain iterators */
4349 ittuple = PyTuple_New(tuplesize);
4350 if (ittuple == NULL)
4351 return NULL;
4352 for (i=0; i < tuplesize; ++i) {
4353 PyObject *item = PyTuple_GET_ITEM(args, i);
4354 PyObject *it = PyObject_GetIter(item);
4355 if (it == NULL) {
4356 if (PyErr_ExceptionMatches(PyExc_TypeError))
4357 PyErr_Format(PyExc_TypeError,
4358 "zip_longest argument #%zd must support iteration",
4359 i+1);
4360 Py_DECREF(ittuple);
4361 return NULL;
4362 }
4363 PyTuple_SET_ITEM(ittuple, i, it);
4364 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004365
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004366 /* create a result holder */
4367 result = PyTuple_New(tuplesize);
4368 if (result == NULL) {
4369 Py_DECREF(ittuple);
4370 return NULL;
4371 }
4372 for (i=0 ; i < tuplesize ; i++) {
4373 Py_INCREF(Py_None);
4374 PyTuple_SET_ITEM(result, i, Py_None);
4375 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004376
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004377 /* create ziplongestobject structure */
4378 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4379 if (lz == NULL) {
4380 Py_DECREF(ittuple);
4381 Py_DECREF(result);
4382 return NULL;
4383 }
4384 lz->ittuple = ittuple;
4385 lz->tuplesize = tuplesize;
4386 lz->numactive = tuplesize;
4387 lz->result = result;
4388 Py_INCREF(fillvalue);
4389 lz->fillvalue = fillvalue;
4390 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004391}
4392
4393static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004394zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004395{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004396 PyObject_GC_UnTrack(lz);
4397 Py_XDECREF(lz->ittuple);
4398 Py_XDECREF(lz->result);
4399 Py_XDECREF(lz->fillvalue);
4400 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004401}
4402
4403static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004404zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004405{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004406 Py_VISIT(lz->ittuple);
4407 Py_VISIT(lz->result);
4408 Py_VISIT(lz->fillvalue);
4409 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004410}
4411
4412static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004413zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004414{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004415 Py_ssize_t i;
4416 Py_ssize_t tuplesize = lz->tuplesize;
4417 PyObject *result = lz->result;
4418 PyObject *it;
4419 PyObject *item;
4420 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004421
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004422 if (tuplesize == 0)
4423 return NULL;
4424 if (lz->numactive == 0)
4425 return NULL;
4426 if (Py_REFCNT(result) == 1) {
4427 Py_INCREF(result);
4428 for (i=0 ; i < tuplesize ; i++) {
4429 it = PyTuple_GET_ITEM(lz->ittuple, i);
4430 if (it == NULL) {
4431 Py_INCREF(lz->fillvalue);
4432 item = lz->fillvalue;
4433 } else {
4434 item = PyIter_Next(it);
4435 if (item == NULL) {
4436 lz->numactive -= 1;
4437 if (lz->numactive == 0 || PyErr_Occurred()) {
4438 lz->numactive = 0;
4439 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004440 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004441 } else {
4442 Py_INCREF(lz->fillvalue);
4443 item = lz->fillvalue;
4444 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4445 Py_DECREF(it);
4446 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004447 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004448 }
4449 olditem = PyTuple_GET_ITEM(result, i);
4450 PyTuple_SET_ITEM(result, i, item);
4451 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004452 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004453 } else {
4454 result = PyTuple_New(tuplesize);
4455 if (result == NULL)
4456 return NULL;
4457 for (i=0 ; i < tuplesize ; i++) {
4458 it = PyTuple_GET_ITEM(lz->ittuple, i);
4459 if (it == NULL) {
4460 Py_INCREF(lz->fillvalue);
4461 item = lz->fillvalue;
4462 } else {
4463 item = PyIter_Next(it);
4464 if (item == NULL) {
4465 lz->numactive -= 1;
4466 if (lz->numactive == 0 || PyErr_Occurred()) {
4467 lz->numactive = 0;
4468 Py_DECREF(result);
4469 return NULL;
4470 } else {
4471 Py_INCREF(lz->fillvalue);
4472 item = lz->fillvalue;
4473 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4474 Py_DECREF(it);
4475 }
4476 }
4477 }
4478 PyTuple_SET_ITEM(result, i, item);
4479 }
4480 }
4481 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004482}
4483
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004484static PyObject *
4485zip_longest_reduce(ziplongestobject *lz)
4486{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004487
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004488 /* Create a new tuple with empty sequences where appropriate to pickle.
4489 * Then use setstate to set the fillvalue
4490 */
4491 int i;
4492 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
4493 if (args == NULL)
4494 return NULL;
4495 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4496 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4497 if (elem == NULL) {
4498 elem = PyTuple_New(0);
4499 if (elem == NULL) {
4500 Py_DECREF(args);
4501 return NULL;
4502 }
4503 } else
4504 Py_INCREF(elem);
4505 PyTuple_SET_ITEM(args, i, elem);
4506 }
4507 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4508}
4509
4510static PyObject *
4511zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4512{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004513 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004514 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004515 Py_RETURN_NONE;
4516}
4517
4518static PyMethodDef zip_longest_methods[] = {
4519 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4520 reduce_doc},
4521 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4522 setstate_doc},
4523 {NULL, NULL} /* sentinel */
4524};
4525
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004526PyDoc_STRVAR(zip_longest_doc,
4527"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004528\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004529Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004530the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004531method continues until the longest iterable in the argument sequence\n\
4532is exhausted and then it raises StopIteration. When the shorter iterables\n\
4533are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4534defaults to None or can be specified by a keyword argument.\n\
4535");
4536
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004537static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004538 PyVarObject_HEAD_INIT(NULL, 0)
4539 "itertools.zip_longest", /* tp_name */
4540 sizeof(ziplongestobject), /* tp_basicsize */
4541 0, /* tp_itemsize */
4542 /* methods */
4543 (destructor)zip_longest_dealloc, /* tp_dealloc */
4544 0, /* tp_print */
4545 0, /* tp_getattr */
4546 0, /* tp_setattr */
4547 0, /* tp_reserved */
4548 0, /* tp_repr */
4549 0, /* tp_as_number */
4550 0, /* tp_as_sequence */
4551 0, /* tp_as_mapping */
4552 0, /* tp_hash */
4553 0, /* tp_call */
4554 0, /* tp_str */
4555 PyObject_GenericGetAttr, /* tp_getattro */
4556 0, /* tp_setattro */
4557 0, /* tp_as_buffer */
4558 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4559 Py_TPFLAGS_BASETYPE, /* tp_flags */
4560 zip_longest_doc, /* tp_doc */
4561 (traverseproc)zip_longest_traverse, /* tp_traverse */
4562 0, /* tp_clear */
4563 0, /* tp_richcompare */
4564 0, /* tp_weaklistoffset */
4565 PyObject_SelfIter, /* tp_iter */
4566 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004567 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004568 0, /* tp_members */
4569 0, /* tp_getset */
4570 0, /* tp_base */
4571 0, /* tp_dict */
4572 0, /* tp_descr_get */
4573 0, /* tp_descr_set */
4574 0, /* tp_dictoffset */
4575 0, /* tp_init */
4576 0, /* tp_alloc */
4577 zip_longest_new, /* tp_new */
4578 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004579};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004580
4581/* module level code ********************************************************/
4582
4583PyDoc_STRVAR(module_doc,
4584"Functional tools for creating and using iterators.\n\
4585\n\
4586Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004587count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004588cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004589repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004590\n\
4591Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004592accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004593chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger8df58f72013-09-09 01:29:40 -05004594chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004595compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4596dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4597groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004598filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004599islice(seq, [start,] stop [, step]) --> elements from\n\
4600 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004601starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004602tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004603takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004604zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004605\n\
4606Combinatoric generators:\n\
4607product(p, q, ... [repeat=1]) --> cartesian product\n\
4608permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004609combinations(p, r)\n\
4610combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004611");
4612
4613
Raymond Hettingerad983e72003-11-12 14:32:26 +00004614static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004615 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4616 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004617};
4618
Martin v. Löwis1a214512008-06-11 05:26:20 +00004619
4620static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004621 PyModuleDef_HEAD_INIT,
4622 "itertools",
4623 module_doc,
4624 -1,
4625 module_methods,
4626 NULL,
4627 NULL,
4628 NULL,
4629 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004630};
4631
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004632PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004633PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004634{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004635 int i;
4636 PyObject *m;
4637 char *name;
4638 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004639 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004640 &combinations_type,
4641 &cwr_type,
4642 &cycle_type,
4643 &dropwhile_type,
4644 &takewhile_type,
4645 &islice_type,
4646 &starmap_type,
4647 &chain_type,
4648 &compress_type,
4649 &filterfalse_type,
4650 &count_type,
4651 &ziplongest_type,
4652 &permutations_type,
4653 &product_type,
4654 &repeat_type,
4655 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004656 &_grouper_type,
4657 &tee_type,
4658 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004659 NULL
4660 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004661
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004662 Py_TYPE(&teedataobject_type) = &PyType_Type;
4663 m = PyModule_Create(&itertoolsmodule);
4664 if (m == NULL)
4665 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004666
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004667 for (i=0 ; typelist[i] != NULL ; i++) {
4668 if (PyType_Ready(typelist[i]) < 0)
4669 return NULL;
4670 name = strchr(typelist[i]->tp_name, '.');
4671 assert (name != NULL);
4672 Py_INCREF(typelist[i]);
4673 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4674 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004675
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004676 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004677}