blob: d7a1ef07428058189ccbb7887d92d975d63a90b8 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
Raymond Hettingerca3788c2015-08-16 14:49:24 -07002#define PY_SSIZE_T_CLEAN
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003#include "Python.h"
Fred Drake08ebfec2004-10-17 19:36:57 +00004#include "structmember.h"
Raymond Hettinger96ef8112003-02-01 00:10:11 +00005
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00006/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00007 by Raymond D. Hettinger <python@rcn.com>
Raymond Hettinger96ef8112003-02-01 00:10:11 +00008*/
9
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000010
Raymond Hettingera6ea44a2015-08-17 23:55:28 -070011/* groupby object ************************************************************/
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000012
13typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000014 PyObject_HEAD
15 PyObject *it;
16 PyObject *keyfunc;
17 PyObject *tgtkey;
18 PyObject *currkey;
19 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000020} groupbyobject;
21
22static PyTypeObject groupby_type;
23static PyObject *_grouper_create(groupbyobject *, PyObject *);
24
25static PyObject *
26groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
27{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000028 static char *kwargs[] = {"iterable", "key", NULL};
29 groupbyobject *gbo;
30 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000031
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000032 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
33 &it, &keyfunc))
34 return NULL;
35
36 gbo = (groupbyobject *)type->tp_alloc(type, 0);
37 if (gbo == NULL)
38 return NULL;
39 gbo->tgtkey = NULL;
40 gbo->currkey = NULL;
41 gbo->currvalue = NULL;
42 gbo->keyfunc = keyfunc;
43 Py_INCREF(keyfunc);
44 gbo->it = PyObject_GetIter(it);
45 if (gbo->it == NULL) {
46 Py_DECREF(gbo);
47 return NULL;
48 }
49 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000050}
51
52static void
53groupby_dealloc(groupbyobject *gbo)
54{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000055 PyObject_GC_UnTrack(gbo);
56 Py_XDECREF(gbo->it);
57 Py_XDECREF(gbo->keyfunc);
58 Py_XDECREF(gbo->tgtkey);
59 Py_XDECREF(gbo->currkey);
60 Py_XDECREF(gbo->currvalue);
61 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000062}
63
64static int
65groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
66{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000067 Py_VISIT(gbo->it);
68 Py_VISIT(gbo->keyfunc);
69 Py_VISIT(gbo->tgtkey);
70 Py_VISIT(gbo->currkey);
71 Py_VISIT(gbo->currvalue);
72 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000073}
74
75static PyObject *
76groupby_next(groupbyobject *gbo)
77{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +020078 PyObject *newvalue, *newkey, *r, *grouper;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000079
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000080 /* skip to next iteration group */
81 for (;;) {
82 if (gbo->currkey == NULL)
83 /* pass */;
84 else if (gbo->tgtkey == NULL)
85 break;
86 else {
87 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000088
Raymond Hettingera6ea44a2015-08-17 23:55:28 -070089 rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000090 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 {
Victor Stinnerde4ae3d2016-12-04 22:59:09 +0100104 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000105 if (newkey == NULL) {
106 Py_DECREF(newvalue);
107 return NULL;
108 }
109 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000110
Serhiy Storchakaec397562016-04-06 09:50:03 +0300111 Py_XSETREF(gbo->currkey, newkey);
112 Py_XSETREF(gbo->currvalue, newvalue);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000113 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000114
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000115 Py_INCREF(gbo->currkey);
Serhiy Storchakaec397562016-04-06 09:50:03 +0300116 Py_XSETREF(gbo->tgtkey, gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000117
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000118 grouper = _grouper_create(gbo, gbo->tgtkey);
119 if (grouper == NULL)
120 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000121
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000122 r = PyTuple_Pack(2, gbo->currkey, grouper);
123 Py_DECREF(grouper);
124 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000125}
126
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000127static PyObject *
128groupby_reduce(groupbyobject *lz)
129{
130 /* reduce as a 'new' call with an optional 'setstate' if groupby
131 * has started
132 */
133 PyObject *value;
134 if (lz->tgtkey && lz->currkey && lz->currvalue)
135 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
136 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
137 else
138 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
139 lz->it, lz->keyfunc);
140
141 return value;
142}
143
144PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
145
146static PyObject *
147groupby_setstate(groupbyobject *lz, PyObject *state)
148{
149 PyObject *currkey, *currvalue, *tgtkey;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300150 if (!PyTuple_Check(state)) {
151 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000152 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300153 }
154 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey)) {
155 return NULL;
156 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200157 Py_INCREF(currkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300158 Py_XSETREF(lz->currkey, currkey);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200159 Py_INCREF(currvalue);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300160 Py_XSETREF(lz->currvalue, currvalue);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200161 Py_INCREF(tgtkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300162 Py_XSETREF(lz->tgtkey, tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000163 Py_RETURN_NONE;
164}
165
166PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
167
168static PyMethodDef groupby_methods[] = {
169 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
170 reduce_doc},
171 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
172 setstate_doc},
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700173 {NULL, NULL} /* sentinel */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000174};
175
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000176PyDoc_STRVAR(groupby_doc,
177"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
178(key, sub-iterator) grouped by each value of key(value).\n");
179
180static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000181 PyVarObject_HEAD_INIT(NULL, 0)
182 "itertools.groupby", /* tp_name */
183 sizeof(groupbyobject), /* tp_basicsize */
184 0, /* tp_itemsize */
185 /* methods */
186 (destructor)groupby_dealloc, /* tp_dealloc */
187 0, /* tp_print */
188 0, /* tp_getattr */
189 0, /* tp_setattr */
190 0, /* tp_reserved */
191 0, /* tp_repr */
192 0, /* tp_as_number */
193 0, /* tp_as_sequence */
194 0, /* tp_as_mapping */
195 0, /* tp_hash */
196 0, /* tp_call */
197 0, /* tp_str */
198 PyObject_GenericGetAttr, /* tp_getattro */
199 0, /* tp_setattro */
200 0, /* tp_as_buffer */
201 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
202 Py_TPFLAGS_BASETYPE, /* tp_flags */
203 groupby_doc, /* tp_doc */
204 (traverseproc)groupby_traverse, /* tp_traverse */
205 0, /* tp_clear */
206 0, /* tp_richcompare */
207 0, /* tp_weaklistoffset */
208 PyObject_SelfIter, /* tp_iter */
209 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000210 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000211 0, /* tp_members */
212 0, /* tp_getset */
213 0, /* tp_base */
214 0, /* tp_dict */
215 0, /* tp_descr_get */
216 0, /* tp_descr_set */
217 0, /* tp_dictoffset */
218 0, /* tp_init */
219 0, /* tp_alloc */
220 groupby_new, /* tp_new */
221 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000222};
223
224
225/* _grouper object (internal) ************************************************/
226
227typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000228 PyObject_HEAD
229 PyObject *parent;
230 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000231} _grouperobject;
232
233static PyTypeObject _grouper_type;
234
235static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000236_grouper_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
237{
238 PyObject *parent, *tgtkey;
239
240 if (!PyArg_ParseTuple(args, "O!O", &groupby_type, &parent, &tgtkey))
241 return NULL;
242
243 return _grouper_create((groupbyobject*) parent, tgtkey);
244}
245
246static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000247_grouper_create(groupbyobject *parent, PyObject *tgtkey)
248{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000249 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000250
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000251 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
252 if (igo == NULL)
253 return NULL;
254 igo->parent = (PyObject *)parent;
255 Py_INCREF(parent);
256 igo->tgtkey = tgtkey;
257 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000259 PyObject_GC_Track(igo);
260 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000261}
262
263static void
264_grouper_dealloc(_grouperobject *igo)
265{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000266 PyObject_GC_UnTrack(igo);
267 Py_DECREF(igo->parent);
268 Py_DECREF(igo->tgtkey);
269 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000270}
271
272static int
273_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
274{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000275 Py_VISIT(igo->parent);
276 Py_VISIT(igo->tgtkey);
277 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000278}
279
280static PyObject *
281_grouper_next(_grouperobject *igo)
282{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000283 groupbyobject *gbo = (groupbyobject *)igo->parent;
284 PyObject *newvalue, *newkey, *r;
285 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000286
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000287 if (gbo->currvalue == NULL) {
288 newvalue = PyIter_Next(gbo->it);
289 if (newvalue == NULL)
290 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000291
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000292 if (gbo->keyfunc == Py_None) {
293 newkey = newvalue;
294 Py_INCREF(newvalue);
295 } else {
Victor Stinnerde4ae3d2016-12-04 22:59:09 +0100296 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000297 if (newkey == NULL) {
298 Py_DECREF(newvalue);
299 return NULL;
300 }
301 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000302
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000303 assert(gbo->currkey == NULL);
304 gbo->currkey = newkey;
305 gbo->currvalue = newvalue;
306 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000307
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000308 assert(gbo->currkey != NULL);
309 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
310 if (rcmp <= 0)
311 /* got any error or current group is end */
312 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000313
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000314 r = gbo->currvalue;
315 gbo->currvalue = NULL;
316 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000317
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000318 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000319}
320
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000321static PyObject *
322_grouper_reduce(_grouperobject *lz)
323{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700324 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000325}
326
327static PyMethodDef _grouper_methods[] = {
328 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
329 reduce_doc},
330 {NULL, NULL} /* sentinel */
331};
332
333
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000334static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000335 PyVarObject_HEAD_INIT(NULL, 0)
336 "itertools._grouper", /* tp_name */
337 sizeof(_grouperobject), /* tp_basicsize */
338 0, /* tp_itemsize */
339 /* methods */
340 (destructor)_grouper_dealloc, /* tp_dealloc */
341 0, /* tp_print */
342 0, /* tp_getattr */
343 0, /* tp_setattr */
344 0, /* tp_reserved */
345 0, /* tp_repr */
346 0, /* tp_as_number */
347 0, /* tp_as_sequence */
348 0, /* tp_as_mapping */
349 0, /* tp_hash */
350 0, /* tp_call */
351 0, /* tp_str */
352 PyObject_GenericGetAttr, /* tp_getattro */
353 0, /* tp_setattro */
354 0, /* tp_as_buffer */
355 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
356 0, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700357 (traverseproc)_grouper_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000358 0, /* tp_clear */
359 0, /* tp_richcompare */
360 0, /* tp_weaklistoffset */
361 PyObject_SelfIter, /* tp_iter */
362 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000363 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000364 0, /* tp_members */
365 0, /* tp_getset */
366 0, /* tp_base */
367 0, /* tp_dict */
368 0, /* tp_descr_get */
369 0, /* tp_descr_set */
370 0, /* tp_dictoffset */
371 0, /* tp_init */
372 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000373 _grouper_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000374 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000375};
376
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000377
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700378/* tee object and with supporting function and objects ***********************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000379
Raymond Hettingerad983e72003-11-12 14:32:26 +0000380/* The teedataobject pre-allocates space for LINKCELLS number of objects.
381 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000382 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000383 the other structure members including PyHEAD overhead. The larger the
384 value, the less memory overhead per object and the less time spent
385 allocating/deallocating new links. The smaller the number, the less
386 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000387*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000388#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000389
390typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000391 PyObject_HEAD
392 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700393 int numread; /* 0 <= numread <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000394 PyObject *nextlink;
395 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000396} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000397
398typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000399 PyObject_HEAD
400 teedataobject *dataobj;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700401 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000402 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000403} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000404
Raymond Hettingerad983e72003-11-12 14:32:26 +0000405static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000406
407static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000408teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000409{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000410 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000411
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000412 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
413 if (tdo == NULL)
414 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000415
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000416 tdo->numread = 0;
417 tdo->nextlink = NULL;
418 Py_INCREF(it);
419 tdo->it = it;
420 PyObject_GC_Track(tdo);
421 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000422}
423
424static PyObject *
425teedataobject_jumplink(teedataobject *tdo)
426{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000427 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000428 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000429 Py_XINCREF(tdo->nextlink);
430 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000431}
432
433static PyObject *
434teedataobject_getitem(teedataobject *tdo, int i)
435{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000436 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000437
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000438 assert(i < LINKCELLS);
439 if (i < tdo->numread)
440 value = tdo->values[i];
441 else {
442 /* this is the lead iterator, so fetch more data */
443 assert(i == tdo->numread);
444 value = PyIter_Next(tdo->it);
445 if (value == NULL)
446 return NULL;
447 tdo->numread++;
448 tdo->values[i] = value;
449 }
450 Py_INCREF(value);
451 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000452}
453
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000454static int
455teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
456{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000457 int i;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700458
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000459 Py_VISIT(tdo->it);
460 for (i = 0; i < tdo->numread; i++)
461 Py_VISIT(tdo->values[i]);
462 Py_VISIT(tdo->nextlink);
463 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000464}
465
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200466static void
467teedataobject_safe_decref(PyObject *obj)
468{
469 while (obj && Py_TYPE(obj) == &teedataobject_type &&
470 Py_REFCNT(obj) == 1) {
471 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
472 ((teedataobject *)obj)->nextlink = NULL;
473 Py_DECREF(obj);
474 obj = nextlink;
475 }
476 Py_XDECREF(obj);
477}
478
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000479static int
480teedataobject_clear(teedataobject *tdo)
481{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000482 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200483 PyObject *tmp;
484
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000485 Py_CLEAR(tdo->it);
486 for (i=0 ; i<tdo->numread ; i++)
487 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200488 tmp = tdo->nextlink;
489 tdo->nextlink = NULL;
490 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000491 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000492}
493
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000494static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000495teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000496{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000497 PyObject_GC_UnTrack(tdo);
498 teedataobject_clear(tdo);
499 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000500}
501
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000502static PyObject *
503teedataobject_reduce(teedataobject *tdo)
504{
505 int i;
506 /* create a temporary list of already iterated values */
507 PyObject *values = PyList_New(tdo->numread);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700508
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000509 if (!values)
510 return NULL;
511 for (i=0 ; i<tdo->numread ; i++) {
512 Py_INCREF(tdo->values[i]);
513 PyList_SET_ITEM(values, i, tdo->values[i]);
514 }
515 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
516 values,
517 tdo->nextlink ? tdo->nextlink : Py_None);
518}
519
520static PyTypeObject teedataobject_type;
521
522static PyObject *
523teedataobject_new(PyTypeObject *type, PyObject *args, PyObject *kw)
524{
525 teedataobject *tdo;
526 PyObject *it, *values, *next;
527 Py_ssize_t i, len;
528
529 assert(type == &teedataobject_type);
530 if (!PyArg_ParseTuple(args, "OO!O", &it, &PyList_Type, &values, &next))
531 return NULL;
532
533 tdo = (teedataobject *)teedataobject_newinternal(it);
534 if (!tdo)
535 return NULL;
536
537 len = PyList_GET_SIZE(values);
538 if (len > LINKCELLS)
539 goto err;
540 for (i=0; i<len; i++) {
541 tdo->values[i] = PyList_GET_ITEM(values, i);
542 Py_INCREF(tdo->values[i]);
543 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200544 /* len <= LINKCELLS < INT_MAX */
545 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000546
547 if (len == LINKCELLS) {
548 if (next != Py_None) {
549 if (Py_TYPE(next) != &teedataobject_type)
550 goto err;
551 assert(tdo->nextlink == NULL);
552 Py_INCREF(next);
553 tdo->nextlink = next;
554 }
555 } else {
556 if (next != Py_None)
557 goto err; /* shouldn't have a next if we are not full */
558 }
559 return (PyObject*)tdo;
560
561err:
562 Py_XDECREF(tdo);
563 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
564 return NULL;
565}
566
567static PyMethodDef teedataobject_methods[] = {
568 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
569 reduce_doc},
570 {NULL, NULL} /* sentinel */
571};
572
Raymond Hettingerad983e72003-11-12 14:32:26 +0000573PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000574
Raymond Hettingerad983e72003-11-12 14:32:26 +0000575static PyTypeObject teedataobject_type = {
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700576 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000577 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000578 sizeof(teedataobject), /* tp_basicsize */
579 0, /* tp_itemsize */
580 /* methods */
581 (destructor)teedataobject_dealloc, /* tp_dealloc */
582 0, /* tp_print */
583 0, /* tp_getattr */
584 0, /* tp_setattr */
585 0, /* tp_reserved */
586 0, /* tp_repr */
587 0, /* tp_as_number */
588 0, /* tp_as_sequence */
589 0, /* tp_as_mapping */
590 0, /* tp_hash */
591 0, /* tp_call */
592 0, /* tp_str */
593 PyObject_GenericGetAttr, /* tp_getattro */
594 0, /* tp_setattro */
595 0, /* tp_as_buffer */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700596 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000597 teedataobject_doc, /* tp_doc */
598 (traverseproc)teedataobject_traverse, /* tp_traverse */
599 (inquiry)teedataobject_clear, /* tp_clear */
600 0, /* tp_richcompare */
601 0, /* tp_weaklistoffset */
602 0, /* tp_iter */
603 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000604 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000605 0, /* tp_members */
606 0, /* tp_getset */
607 0, /* tp_base */
608 0, /* tp_dict */
609 0, /* tp_descr_get */
610 0, /* tp_descr_set */
611 0, /* tp_dictoffset */
612 0, /* tp_init */
613 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000614 teedataobject_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000615 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000616};
617
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000618
619static PyTypeObject tee_type;
620
621static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000622tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000623{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000624 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000625
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000626 if (to->index >= LINKCELLS) {
627 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200628 if (link == NULL)
629 return NULL;
Serhiy Storchaka57a01d32016-04-10 18:05:40 +0300630 Py_SETREF(to->dataobj, (teedataobject *)link);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000631 to->index = 0;
632 }
633 value = teedataobject_getitem(to->dataobj, to->index);
634 if (value == NULL)
635 return NULL;
636 to->index++;
637 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000638}
639
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000640static int
641tee_traverse(teeobject *to, visitproc visit, void *arg)
642{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000643 Py_VISIT((PyObject *)to->dataobj);
644 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000645}
646
Raymond Hettingerad983e72003-11-12 14:32:26 +0000647static PyObject *
648tee_copy(teeobject *to)
649{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000650 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000651
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000652 newto = PyObject_GC_New(teeobject, &tee_type);
653 if (newto == NULL)
654 return NULL;
655 Py_INCREF(to->dataobj);
656 newto->dataobj = to->dataobj;
657 newto->index = to->index;
658 newto->weakreflist = NULL;
659 PyObject_GC_Track(newto);
660 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000661}
662
663PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
664
665static PyObject *
666tee_fromiterable(PyObject *iterable)
667{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000668 teeobject *to;
669 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000670
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000671 it = PyObject_GetIter(iterable);
672 if (it == NULL)
673 return NULL;
674 if (PyObject_TypeCheck(it, &tee_type)) {
675 to = (teeobject *)tee_copy((teeobject *)it);
676 goto done;
677 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000678
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000679 to = PyObject_GC_New(teeobject, &tee_type);
680 if (to == NULL)
681 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000682 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000683 if (!to->dataobj) {
684 PyObject_GC_Del(to);
685 to = NULL;
686 goto done;
687 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000688
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000689 to->index = 0;
690 to->weakreflist = NULL;
691 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000692done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000693 Py_XDECREF(it);
694 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000695}
696
697static PyObject *
698tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
699{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000700 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000701
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000702 if (!PyArg_UnpackTuple(args, "_tee", 1, 1, &iterable))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000703 return NULL;
704 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000705}
706
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000707static int
708tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000709{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000710 if (to->weakreflist != NULL)
711 PyObject_ClearWeakRefs((PyObject *) to);
712 Py_CLEAR(to->dataobj);
713 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000714}
715
716static void
717tee_dealloc(teeobject *to)
718{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000719 PyObject_GC_UnTrack(to);
720 tee_clear(to);
721 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000722}
723
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000724static PyObject *
725tee_reduce(teeobject *to)
726{
727 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
728}
729
730static PyObject *
731tee_setstate(teeobject *to, PyObject *state)
732{
733 teedataobject *tdo;
734 int index;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300735 if (!PyTuple_Check(state)) {
736 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000737 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300738 }
739 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
740 return NULL;
741 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000742 if (index < 0 || index > LINKCELLS) {
743 PyErr_SetString(PyExc_ValueError, "Index out of range");
744 return NULL;
745 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200746 Py_INCREF(tdo);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300747 Py_XSETREF(to->dataobj, tdo);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000748 to->index = index;
749 Py_RETURN_NONE;
750}
751
Raymond Hettingerad983e72003-11-12 14:32:26 +0000752PyDoc_STRVAR(teeobject_doc,
753"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000754
Raymond Hettingerad983e72003-11-12 14:32:26 +0000755static PyMethodDef tee_methods[] = {
Raymond Hettingera6a2d442015-08-16 14:38:07 -0700756 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
757 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
758 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000759 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000760};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000761
762static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000763 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000764 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000765 sizeof(teeobject), /* tp_basicsize */
766 0, /* tp_itemsize */
767 /* methods */
768 (destructor)tee_dealloc, /* tp_dealloc */
769 0, /* tp_print */
770 0, /* tp_getattr */
771 0, /* tp_setattr */
772 0, /* tp_reserved */
773 0, /* tp_repr */
774 0, /* tp_as_number */
775 0, /* tp_as_sequence */
776 0, /* tp_as_mapping */
777 0, /* tp_hash */
778 0, /* tp_call */
779 0, /* tp_str */
780 0, /* tp_getattro */
781 0, /* tp_setattro */
782 0, /* tp_as_buffer */
783 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
784 teeobject_doc, /* tp_doc */
785 (traverseproc)tee_traverse, /* tp_traverse */
786 (inquiry)tee_clear, /* tp_clear */
787 0, /* tp_richcompare */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700788 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000789 PyObject_SelfIter, /* tp_iter */
790 (iternextfunc)tee_next, /* tp_iternext */
791 tee_methods, /* tp_methods */
792 0, /* tp_members */
793 0, /* tp_getset */
794 0, /* tp_base */
795 0, /* tp_dict */
796 0, /* tp_descr_get */
797 0, /* tp_descr_set */
798 0, /* tp_dictoffset */
799 0, /* tp_init */
800 0, /* tp_alloc */
801 tee_new, /* tp_new */
802 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000803};
804
Raymond Hettingerad983e72003-11-12 14:32:26 +0000805static PyObject *
806tee(PyObject *self, PyObject *args)
807{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000808 Py_ssize_t i, n=2;
809 PyObject *it, *iterable, *copyable, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200810 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000811
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000812 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
813 return NULL;
814 if (n < 0) {
815 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
816 return NULL;
817 }
818 result = PyTuple_New(n);
819 if (result == NULL)
820 return NULL;
821 if (n == 0)
822 return result;
823 it = PyObject_GetIter(iterable);
824 if (it == NULL) {
825 Py_DECREF(result);
826 return NULL;
827 }
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200828 if (!_PyObject_HasAttrId(it, &PyId___copy__)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000829 copyable = tee_fromiterable(it);
830 Py_DECREF(it);
831 if (copyable == NULL) {
832 Py_DECREF(result);
833 return NULL;
834 }
835 } else
836 copyable = it;
837 PyTuple_SET_ITEM(result, 0, copyable);
838 for (i=1 ; i<n ; i++) {
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200839
840 copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000841 if (copyable == NULL) {
842 Py_DECREF(result);
843 return NULL;
844 }
845 PyTuple_SET_ITEM(result, i, copyable);
846 }
847 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000848}
849
850PyDoc_STRVAR(tee_doc,
851"tee(iterable, n=2) --> tuple of n independent iterators.");
852
853
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700854/* cycle object **************************************************************/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000855
856typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000857 PyObject_HEAD
858 PyObject *it;
859 PyObject *saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700860 Py_ssize_t index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000861 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000862} cycleobject;
863
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000864static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000865
866static PyObject *
867cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
868{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000869 PyObject *it;
870 PyObject *iterable;
871 PyObject *saved;
872 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000873
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +0300874 if (type == &cycle_type && !_PyArg_NoKeywords("cycle", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000875 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000876
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000877 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
878 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000879
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000880 /* Get iterator. */
881 it = PyObject_GetIter(iterable);
882 if (it == NULL)
883 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000884
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000885 saved = PyList_New(0);
886 if (saved == NULL) {
887 Py_DECREF(it);
888 return NULL;
889 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000890
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000891 /* create cycleobject structure */
892 lz = (cycleobject *)type->tp_alloc(type, 0);
893 if (lz == NULL) {
894 Py_DECREF(it);
895 Py_DECREF(saved);
896 return NULL;
897 }
898 lz->it = it;
899 lz->saved = saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700900 lz->index = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000901 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000902
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000903 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000904}
905
906static void
907cycle_dealloc(cycleobject *lz)
908{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000909 PyObject_GC_UnTrack(lz);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000910 Py_XDECREF(lz->it);
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700911 Py_XDECREF(lz->saved);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000912 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000913}
914
915static int
916cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
917{
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700918 if (lz->it)
919 Py_VISIT(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000920 Py_VISIT(lz->saved);
921 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000922}
923
924static PyObject *
925cycle_next(cycleobject *lz)
926{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000927 PyObject *item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000928
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700929 if (lz->it != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000930 item = PyIter_Next(lz->it);
931 if (item != NULL) {
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700932 if (lz->firstpass)
933 return item;
934 if (PyList_Append(lz->saved, item)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000935 Py_DECREF(item);
936 return NULL;
937 }
938 return item;
939 }
Raymond Hettinger98958fe2015-08-15 15:09:30 -0700940 /* Note: StopIteration is already cleared by PyIter_Next() */
941 if (PyErr_Occurred())
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000942 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700943 Py_CLEAR(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000944 }
Serhiy Storchakafff9a312017-03-21 08:53:25 +0200945 if (PyList_GET_SIZE(lz->saved) == 0)
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700946 return NULL;
947 item = PyList_GET_ITEM(lz->saved, lz->index);
948 lz->index++;
Serhiy Storchakafff9a312017-03-21 08:53:25 +0200949 if (lz->index >= PyList_GET_SIZE(lz->saved))
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700950 lz->index = 0;
951 Py_INCREF(item);
952 return item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000953}
954
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000955static PyObject *
956cycle_reduce(cycleobject *lz)
957{
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700958 /* Create a new cycle with the iterator tuple, then set the saved state */
959 if (lz->it == NULL) {
960 PyObject *it = PyObject_GetIter(lz->saved);
961 if (it == NULL)
962 return NULL;
963 if (lz->index != 0) {
964 _Py_IDENTIFIER(__setstate__);
965 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
966 "n", lz->index);
967 if (res == NULL) {
968 Py_DECREF(it);
969 return NULL;
970 }
971 Py_DECREF(res);
972 }
973 return Py_BuildValue("O(N)(Oi)", Py_TYPE(lz), it, lz->saved, 1);
974 }
975 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), lz->it, lz->saved,
976 lz->firstpass);
Raymond Hettingerc39786d2015-08-15 15:16:12 -0700977}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000978
979static PyObject *
980cycle_setstate(cycleobject *lz, PyObject *state)
981{
982 PyObject *saved=NULL;
983 int firstpass;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300984 if (!PyTuple_Check(state)) {
985 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000986 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300987 }
988 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
989 return NULL;
990 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700991 Py_INCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300992 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000993 lz->firstpass = firstpass != 0;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700994 lz->index = 0;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000995 Py_RETURN_NONE;
996}
997
998static PyMethodDef cycle_methods[] = {
999 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1000 reduce_doc},
1001 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1002 setstate_doc},
1003 {NULL, NULL} /* sentinel */
1004};
1005
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001006PyDoc_STRVAR(cycle_doc,
1007"cycle(iterable) --> cycle object\n\
1008\n\
1009Return elements from the iterable until it is exhausted.\n\
1010Then repeat the sequence indefinitely.");
1011
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001012static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001013 PyVarObject_HEAD_INIT(NULL, 0)
1014 "itertools.cycle", /* tp_name */
1015 sizeof(cycleobject), /* tp_basicsize */
1016 0, /* tp_itemsize */
1017 /* methods */
1018 (destructor)cycle_dealloc, /* tp_dealloc */
1019 0, /* tp_print */
1020 0, /* tp_getattr */
1021 0, /* tp_setattr */
1022 0, /* tp_reserved */
1023 0, /* tp_repr */
1024 0, /* tp_as_number */
1025 0, /* tp_as_sequence */
1026 0, /* tp_as_mapping */
1027 0, /* tp_hash */
1028 0, /* tp_call */
1029 0, /* tp_str */
1030 PyObject_GenericGetAttr, /* tp_getattro */
1031 0, /* tp_setattro */
1032 0, /* tp_as_buffer */
1033 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1034 Py_TPFLAGS_BASETYPE, /* tp_flags */
1035 cycle_doc, /* tp_doc */
1036 (traverseproc)cycle_traverse, /* tp_traverse */
1037 0, /* tp_clear */
1038 0, /* tp_richcompare */
1039 0, /* tp_weaklistoffset */
1040 PyObject_SelfIter, /* tp_iter */
1041 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001042 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001043 0, /* tp_members */
1044 0, /* tp_getset */
1045 0, /* tp_base */
1046 0, /* tp_dict */
1047 0, /* tp_descr_get */
1048 0, /* tp_descr_set */
1049 0, /* tp_dictoffset */
1050 0, /* tp_init */
1051 0, /* tp_alloc */
1052 cycle_new, /* tp_new */
1053 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001054};
1055
1056
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001057/* dropwhile object **********************************************************/
1058
1059typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001060 PyObject_HEAD
1061 PyObject *func;
1062 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001063 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001064} dropwhileobject;
1065
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001066static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001067
1068static PyObject *
1069dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1070{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001071 PyObject *func, *seq;
1072 PyObject *it;
1073 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001074
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001075 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001076 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001077
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001078 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
1079 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001080
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001081 /* Get iterator. */
1082 it = PyObject_GetIter(seq);
1083 if (it == NULL)
1084 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001085
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001086 /* create dropwhileobject structure */
1087 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1088 if (lz == NULL) {
1089 Py_DECREF(it);
1090 return NULL;
1091 }
1092 Py_INCREF(func);
1093 lz->func = func;
1094 lz->it = it;
1095 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001096
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001097 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001098}
1099
1100static void
1101dropwhile_dealloc(dropwhileobject *lz)
1102{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001103 PyObject_GC_UnTrack(lz);
1104 Py_XDECREF(lz->func);
1105 Py_XDECREF(lz->it);
1106 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001107}
1108
1109static int
1110dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1111{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001112 Py_VISIT(lz->it);
1113 Py_VISIT(lz->func);
1114 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001115}
1116
1117static PyObject *
1118dropwhile_next(dropwhileobject *lz)
1119{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001120 PyObject *item, *good;
1121 PyObject *it = lz->it;
1122 long ok;
1123 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001124
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001125 iternext = *Py_TYPE(it)->tp_iternext;
1126 for (;;) {
1127 item = iternext(it);
1128 if (item == NULL)
1129 return NULL;
1130 if (lz->start == 1)
1131 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001132
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01001133 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001134 if (good == NULL) {
1135 Py_DECREF(item);
1136 return NULL;
1137 }
1138 ok = PyObject_IsTrue(good);
1139 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001140 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001141 lz->start = 1;
1142 return item;
1143 }
1144 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001145 if (ok < 0)
1146 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001147 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001148}
1149
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001150static PyObject *
1151dropwhile_reduce(dropwhileobject *lz)
1152{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001153 return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->start);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001154}
1155
1156static PyObject *
1157dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1158{
1159 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001160 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001161 return NULL;
1162 lz->start = start;
1163 Py_RETURN_NONE;
1164}
1165
1166static PyMethodDef dropwhile_methods[] = {
1167 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1168 reduce_doc},
1169 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1170 setstate_doc},
1171 {NULL, NULL} /* sentinel */
1172};
1173
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001174PyDoc_STRVAR(dropwhile_doc,
1175"dropwhile(predicate, iterable) --> dropwhile object\n\
1176\n\
1177Drop items from the iterable while predicate(item) is true.\n\
1178Afterwards, return every element until the iterable is exhausted.");
1179
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001180static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001181 PyVarObject_HEAD_INIT(NULL, 0)
1182 "itertools.dropwhile", /* tp_name */
1183 sizeof(dropwhileobject), /* tp_basicsize */
1184 0, /* tp_itemsize */
1185 /* methods */
1186 (destructor)dropwhile_dealloc, /* tp_dealloc */
1187 0, /* tp_print */
1188 0, /* tp_getattr */
1189 0, /* tp_setattr */
1190 0, /* tp_reserved */
1191 0, /* tp_repr */
1192 0, /* tp_as_number */
1193 0, /* tp_as_sequence */
1194 0, /* tp_as_mapping */
1195 0, /* tp_hash */
1196 0, /* tp_call */
1197 0, /* tp_str */
1198 PyObject_GenericGetAttr, /* tp_getattro */
1199 0, /* tp_setattro */
1200 0, /* tp_as_buffer */
1201 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1202 Py_TPFLAGS_BASETYPE, /* tp_flags */
1203 dropwhile_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001204 (traverseproc)dropwhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001205 0, /* tp_clear */
1206 0, /* tp_richcompare */
1207 0, /* tp_weaklistoffset */
1208 PyObject_SelfIter, /* tp_iter */
1209 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001210 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001211 0, /* tp_members */
1212 0, /* tp_getset */
1213 0, /* tp_base */
1214 0, /* tp_dict */
1215 0, /* tp_descr_get */
1216 0, /* tp_descr_set */
1217 0, /* tp_dictoffset */
1218 0, /* tp_init */
1219 0, /* tp_alloc */
1220 dropwhile_new, /* tp_new */
1221 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001222};
1223
1224
1225/* takewhile object **********************************************************/
1226
1227typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001228 PyObject_HEAD
1229 PyObject *func;
1230 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001231 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001232} takewhileobject;
1233
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001234static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001235
1236static PyObject *
1237takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1238{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001239 PyObject *func, *seq;
1240 PyObject *it;
1241 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001242
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001243 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001244 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001245
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001246 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1247 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001248
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001249 /* Get iterator. */
1250 it = PyObject_GetIter(seq);
1251 if (it == NULL)
1252 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001253
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001254 /* create takewhileobject structure */
1255 lz = (takewhileobject *)type->tp_alloc(type, 0);
1256 if (lz == NULL) {
1257 Py_DECREF(it);
1258 return NULL;
1259 }
1260 Py_INCREF(func);
1261 lz->func = func;
1262 lz->it = it;
1263 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001264
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001265 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001266}
1267
1268static void
1269takewhile_dealloc(takewhileobject *lz)
1270{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001271 PyObject_GC_UnTrack(lz);
1272 Py_XDECREF(lz->func);
1273 Py_XDECREF(lz->it);
1274 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001275}
1276
1277static int
1278takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1279{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001280 Py_VISIT(lz->it);
1281 Py_VISIT(lz->func);
1282 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001283}
1284
1285static PyObject *
1286takewhile_next(takewhileobject *lz)
1287{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001288 PyObject *item, *good;
1289 PyObject *it = lz->it;
1290 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001291
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001292 if (lz->stop == 1)
1293 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001294
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001295 item = (*Py_TYPE(it)->tp_iternext)(it);
1296 if (item == NULL)
1297 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001298
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01001299 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001300 if (good == NULL) {
1301 Py_DECREF(item);
1302 return NULL;
1303 }
1304 ok = PyObject_IsTrue(good);
1305 Py_DECREF(good);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001306 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001307 return item;
1308 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001309 if (ok == 0)
1310 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001311 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001312}
1313
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001314static PyObject *
1315takewhile_reduce(takewhileobject *lz)
1316{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001317 return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->stop);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001318}
1319
1320static PyObject *
1321takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1322{
1323 int stop = PyObject_IsTrue(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001324
Antoine Pitrou721738f2012-08-15 23:20:39 +02001325 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001326 return NULL;
1327 lz->stop = stop;
1328 Py_RETURN_NONE;
1329}
1330
1331static PyMethodDef takewhile_reduce_methods[] = {
1332 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1333 reduce_doc},
1334 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1335 setstate_doc},
1336 {NULL, NULL} /* sentinel */
1337};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001338PyDoc_STRVAR(takewhile_doc,
1339"takewhile(predicate, iterable) --> takewhile object\n\
1340\n\
1341Return successive entries from an iterable as long as the \n\
1342predicate evaluates to true for each entry.");
1343
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001344static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001345 PyVarObject_HEAD_INIT(NULL, 0)
1346 "itertools.takewhile", /* tp_name */
1347 sizeof(takewhileobject), /* tp_basicsize */
1348 0, /* tp_itemsize */
1349 /* methods */
1350 (destructor)takewhile_dealloc, /* tp_dealloc */
1351 0, /* tp_print */
1352 0, /* tp_getattr */
1353 0, /* tp_setattr */
1354 0, /* tp_reserved */
1355 0, /* tp_repr */
1356 0, /* tp_as_number */
1357 0, /* tp_as_sequence */
1358 0, /* tp_as_mapping */
1359 0, /* tp_hash */
1360 0, /* tp_call */
1361 0, /* tp_str */
1362 PyObject_GenericGetAttr, /* tp_getattro */
1363 0, /* tp_setattro */
1364 0, /* tp_as_buffer */
1365 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1366 Py_TPFLAGS_BASETYPE, /* tp_flags */
1367 takewhile_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001368 (traverseproc)takewhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001369 0, /* tp_clear */
1370 0, /* tp_richcompare */
1371 0, /* tp_weaklistoffset */
1372 PyObject_SelfIter, /* tp_iter */
1373 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001374 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001375 0, /* tp_members */
1376 0, /* tp_getset */
1377 0, /* tp_base */
1378 0, /* tp_dict */
1379 0, /* tp_descr_get */
1380 0, /* tp_descr_set */
1381 0, /* tp_dictoffset */
1382 0, /* tp_init */
1383 0, /* tp_alloc */
1384 takewhile_new, /* tp_new */
1385 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001386};
1387
1388
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001389/* islice object *************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001390
1391typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001392 PyObject_HEAD
1393 PyObject *it;
1394 Py_ssize_t next;
1395 Py_ssize_t stop;
1396 Py_ssize_t step;
1397 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001398} isliceobject;
1399
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001400static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001401
1402static PyObject *
1403islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1404{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001405 PyObject *seq;
1406 Py_ssize_t start=0, stop=-1, step=1;
1407 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1408 Py_ssize_t numargs;
1409 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001410
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001411 if (type == &islice_type && !_PyArg_NoKeywords("islice", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001412 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001413
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001414 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1415 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001416
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001417 numargs = PyTuple_Size(args);
1418 if (numargs == 2) {
1419 if (a1 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001420 stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001421 if (stop == -1) {
1422 if (PyErr_Occurred())
1423 PyErr_Clear();
1424 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001425 "Stop argument for islice() must be None or "
1426 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001427 return NULL;
1428 }
1429 }
1430 } else {
1431 if (a1 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001432 start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001433 if (start == -1 && PyErr_Occurred())
1434 PyErr_Clear();
1435 if (a2 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001436 stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001437 if (stop == -1) {
1438 if (PyErr_Occurred())
1439 PyErr_Clear();
1440 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001441 "Stop argument for islice() must be None or "
1442 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001443 return NULL;
1444 }
1445 }
1446 }
1447 if (start<0 || stop<-1) {
1448 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001449 "Indices for islice() must be None or "
1450 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001451 return NULL;
1452 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001453
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001454 if (a3 != NULL) {
1455 if (a3 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001456 step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001457 if (step == -1 && PyErr_Occurred())
1458 PyErr_Clear();
1459 }
1460 if (step<1) {
1461 PyErr_SetString(PyExc_ValueError,
1462 "Step for islice() must be a positive integer or None.");
1463 return NULL;
1464 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001465
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001466 /* Get iterator. */
1467 it = PyObject_GetIter(seq);
1468 if (it == NULL)
1469 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001470
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001471 /* create isliceobject structure */
1472 lz = (isliceobject *)type->tp_alloc(type, 0);
1473 if (lz == NULL) {
1474 Py_DECREF(it);
1475 return NULL;
1476 }
1477 lz->it = it;
1478 lz->next = start;
1479 lz->stop = stop;
1480 lz->step = step;
1481 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001482
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001483 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001484}
1485
1486static void
1487islice_dealloc(isliceobject *lz)
1488{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001489 PyObject_GC_UnTrack(lz);
1490 Py_XDECREF(lz->it);
1491 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001492}
1493
1494static int
1495islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1496{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001497 Py_VISIT(lz->it);
1498 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001499}
1500
1501static PyObject *
1502islice_next(isliceobject *lz)
1503{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001504 PyObject *item;
1505 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001506 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001507 Py_ssize_t oldnext;
1508 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001509
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001510 if (it == NULL)
1511 return NULL;
1512
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001513 iternext = *Py_TYPE(it)->tp_iternext;
1514 while (lz->cnt < lz->next) {
1515 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 Py_DECREF(item);
1519 lz->cnt++;
1520 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001521 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001522 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001523 item = iternext(it);
1524 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001525 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001526 lz->cnt++;
1527 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001528 /* The (size_t) cast below avoids the danger of undefined
1529 behaviour from signed integer overflow. */
1530 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001531 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1532 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001533 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001534
1535empty:
1536 Py_CLEAR(lz->it);
1537 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001538}
1539
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001540static PyObject *
1541islice_reduce(isliceobject *lz)
1542{
1543 /* When unpickled, generate a new object with the same bounds,
1544 * then 'setstate' with the next and count
1545 */
1546 PyObject *stop;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001547
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001548 if (lz->it == NULL) {
1549 PyObject *empty_list;
1550 PyObject *empty_it;
1551 empty_list = PyList_New(0);
1552 if (empty_list == NULL)
1553 return NULL;
1554 empty_it = PyObject_GetIter(empty_list);
1555 Py_DECREF(empty_list);
1556 if (empty_it == NULL)
1557 return NULL;
1558 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1559 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001560 if (lz->stop == -1) {
1561 stop = Py_None;
1562 Py_INCREF(stop);
1563 } else {
1564 stop = PyLong_FromSsize_t(lz->stop);
1565 if (stop == NULL)
1566 return NULL;
1567 }
1568 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1569 lz->it, lz->next, stop, lz->step,
1570 lz->cnt);
1571}
1572
1573static PyObject *
1574islice_setstate(isliceobject *lz, PyObject *state)
1575{
1576 Py_ssize_t cnt = PyLong_AsSsize_t(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001577
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001578 if (cnt == -1 && PyErr_Occurred())
1579 return NULL;
1580 lz->cnt = cnt;
1581 Py_RETURN_NONE;
1582}
1583
1584static PyMethodDef islice_methods[] = {
1585 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1586 reduce_doc},
1587 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1588 setstate_doc},
1589 {NULL, NULL} /* sentinel */
1590};
1591
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001592PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001593"islice(iterable, stop) --> islice object\n\
1594islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001595\n\
1596Return an iterator whose next() method returns selected values from an\n\
1597iterable. If start is specified, will skip all preceding elements;\n\
1598otherwise, start defaults to zero. Step defaults to one. If\n\
1599specified as another value, step determines how many values are \n\
1600skipped between successive calls. Works like a slice() on a list\n\
1601but returns an iterator.");
1602
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001603static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001604 PyVarObject_HEAD_INIT(NULL, 0)
1605 "itertools.islice", /* tp_name */
1606 sizeof(isliceobject), /* tp_basicsize */
1607 0, /* tp_itemsize */
1608 /* methods */
1609 (destructor)islice_dealloc, /* tp_dealloc */
1610 0, /* tp_print */
1611 0, /* tp_getattr */
1612 0, /* tp_setattr */
1613 0, /* tp_reserved */
1614 0, /* tp_repr */
1615 0, /* tp_as_number */
1616 0, /* tp_as_sequence */
1617 0, /* tp_as_mapping */
1618 0, /* tp_hash */
1619 0, /* tp_call */
1620 0, /* tp_str */
1621 PyObject_GenericGetAttr, /* tp_getattro */
1622 0, /* tp_setattro */
1623 0, /* tp_as_buffer */
1624 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1625 Py_TPFLAGS_BASETYPE, /* tp_flags */
1626 islice_doc, /* tp_doc */
1627 (traverseproc)islice_traverse, /* tp_traverse */
1628 0, /* tp_clear */
1629 0, /* tp_richcompare */
1630 0, /* tp_weaklistoffset */
1631 PyObject_SelfIter, /* tp_iter */
1632 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001633 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001634 0, /* tp_members */
1635 0, /* tp_getset */
1636 0, /* tp_base */
1637 0, /* tp_dict */
1638 0, /* tp_descr_get */
1639 0, /* tp_descr_set */
1640 0, /* tp_dictoffset */
1641 0, /* tp_init */
1642 0, /* tp_alloc */
1643 islice_new, /* tp_new */
1644 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001645};
1646
1647
1648/* starmap object ************************************************************/
1649
1650typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001651 PyObject_HEAD
1652 PyObject *func;
1653 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001654} starmapobject;
1655
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001656static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001657
1658static PyObject *
1659starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1660{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001661 PyObject *func, *seq;
1662 PyObject *it;
1663 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001664
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001665 if (type == &starmap_type && !_PyArg_NoKeywords("starmap", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001666 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001667
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001668 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1669 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001670
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001671 /* Get iterator. */
1672 it = PyObject_GetIter(seq);
1673 if (it == NULL)
1674 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001675
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001676 /* create starmapobject structure */
1677 lz = (starmapobject *)type->tp_alloc(type, 0);
1678 if (lz == NULL) {
1679 Py_DECREF(it);
1680 return NULL;
1681 }
1682 Py_INCREF(func);
1683 lz->func = func;
1684 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001685
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001686 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001687}
1688
1689static void
1690starmap_dealloc(starmapobject *lz)
1691{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001692 PyObject_GC_UnTrack(lz);
1693 Py_XDECREF(lz->func);
1694 Py_XDECREF(lz->it);
1695 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001696}
1697
1698static int
1699starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1700{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001701 Py_VISIT(lz->it);
1702 Py_VISIT(lz->func);
1703 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001704}
1705
1706static PyObject *
1707starmap_next(starmapobject *lz)
1708{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001709 PyObject *args;
1710 PyObject *result;
1711 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001712
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001713 args = (*Py_TYPE(it)->tp_iternext)(it);
1714 if (args == NULL)
1715 return NULL;
1716 if (!PyTuple_CheckExact(args)) {
1717 PyObject *newargs = PySequence_Tuple(args);
1718 Py_DECREF(args);
1719 if (newargs == NULL)
1720 return NULL;
1721 args = newargs;
1722 }
1723 result = PyObject_Call(lz->func, args, NULL);
1724 Py_DECREF(args);
1725 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001726}
1727
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001728static PyObject *
1729starmap_reduce(starmapobject *lz)
1730{
1731 /* Just pickle the iterator */
1732 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1733}
1734
1735static PyMethodDef starmap_methods[] = {
1736 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1737 reduce_doc},
1738 {NULL, NULL} /* sentinel */
1739};
1740
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001741PyDoc_STRVAR(starmap_doc,
1742"starmap(function, sequence) --> starmap object\n\
1743\n\
1744Return an iterator whose values are returned from the function evaluated\n\
Serhiy Storchakad65c9492015-11-02 14:10:23 +02001745with an argument tuple taken from the given sequence.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001746
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001747static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001748 PyVarObject_HEAD_INIT(NULL, 0)
1749 "itertools.starmap", /* tp_name */
1750 sizeof(starmapobject), /* tp_basicsize */
1751 0, /* tp_itemsize */
1752 /* methods */
1753 (destructor)starmap_dealloc, /* tp_dealloc */
1754 0, /* tp_print */
1755 0, /* tp_getattr */
1756 0, /* tp_setattr */
1757 0, /* tp_reserved */
1758 0, /* tp_repr */
1759 0, /* tp_as_number */
1760 0, /* tp_as_sequence */
1761 0, /* tp_as_mapping */
1762 0, /* tp_hash */
1763 0, /* tp_call */
1764 0, /* tp_str */
1765 PyObject_GenericGetAttr, /* tp_getattro */
1766 0, /* tp_setattro */
1767 0, /* tp_as_buffer */
1768 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1769 Py_TPFLAGS_BASETYPE, /* tp_flags */
1770 starmap_doc, /* tp_doc */
1771 (traverseproc)starmap_traverse, /* tp_traverse */
1772 0, /* tp_clear */
1773 0, /* tp_richcompare */
1774 0, /* tp_weaklistoffset */
1775 PyObject_SelfIter, /* tp_iter */
1776 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001777 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001778 0, /* tp_members */
1779 0, /* tp_getset */
1780 0, /* tp_base */
1781 0, /* tp_dict */
1782 0, /* tp_descr_get */
1783 0, /* tp_descr_set */
1784 0, /* tp_dictoffset */
1785 0, /* tp_init */
1786 0, /* tp_alloc */
1787 starmap_new, /* tp_new */
1788 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001789};
1790
1791
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001792/* chain object **************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001793
1794typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001795 PyObject_HEAD
1796 PyObject *source; /* Iterator over input iterables */
1797 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001798} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001799
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001800static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001801
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001802static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001803chain_new_internal(PyTypeObject *type, PyObject *source)
1804{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001805 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001806
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001807 lz = (chainobject *)type->tp_alloc(type, 0);
1808 if (lz == NULL) {
1809 Py_DECREF(source);
1810 return NULL;
1811 }
1812
1813 lz->source = source;
1814 lz->active = NULL;
1815 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001816}
1817
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001818static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001819chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001820{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001821 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001822
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001823 if (type == &chain_type && !_PyArg_NoKeywords("chain", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001824 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001825
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001826 source = PyObject_GetIter(args);
1827 if (source == NULL)
1828 return NULL;
1829
1830 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001831}
1832
1833static PyObject *
1834chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1835{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001836 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001837
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001838 source = PyObject_GetIter(arg);
1839 if (source == NULL)
1840 return NULL;
1841
1842 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001843}
1844
1845static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001846chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001847{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001848 PyObject_GC_UnTrack(lz);
1849 Py_XDECREF(lz->active);
1850 Py_XDECREF(lz->source);
1851 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001852}
1853
Raymond Hettinger2012f172003-02-07 05:32:58 +00001854static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001855chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001856{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001857 Py_VISIT(lz->source);
1858 Py_VISIT(lz->active);
1859 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001860}
1861
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001862static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001863chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001864{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001865 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001866
T. Wouters5466d4a2017-03-30 09:58:35 -07001867 /* lz->source is the iterator of iterables. If it's NULL, we've already
1868 * consumed them all. lz->active is the current iterator. If it's NULL,
1869 * we should grab a new one from lz->source. */
1870 while (lz->source != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001871 if (lz->active == NULL) {
T. Wouters5466d4a2017-03-30 09:58:35 -07001872 PyObject *iterable = PyIter_Next(lz->source);
1873 if (iterable == NULL) {
1874 Py_CLEAR(lz->source);
1875 return NULL; /* no more input sources */
1876 }
1877 lz->active = PyObject_GetIter(iterable);
1878 Py_DECREF(iterable);
1879 if (lz->active == NULL) {
1880 Py_CLEAR(lz->source);
1881 return NULL; /* input not iterable */
1882 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001883 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001884 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1885 if (item != NULL)
1886 return item;
1887 if (PyErr_Occurred()) {
1888 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1889 PyErr_Clear();
1890 else
1891 return NULL; /* input raised an exception */
1892 }
1893 /* lz->active is consumed, try with the next iterable. */
1894 Py_CLEAR(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001895 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001896 /* Everything had been consumed already. */
1897 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001898}
1899
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001900static PyObject *
1901chain_reduce(chainobject *lz)
1902{
1903 if (lz->source) {
1904 /* we can't pickle function objects (itertools.from_iterable) so
1905 * we must use setstate to replace the iterable. One day we
1906 * will fix pickling of functions
1907 */
1908 if (lz->active) {
1909 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1910 } else {
1911 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1912 }
1913 } else {
1914 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1915 }
1916 return NULL;
1917}
1918
1919static PyObject *
1920chain_setstate(chainobject *lz, PyObject *state)
1921{
1922 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001923
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001924 if (!PyTuple_Check(state)) {
1925 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001926 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001927 }
1928 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
1929 return NULL;
1930 }
1931 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
1932 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
1933 return NULL;
1934 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001935
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02001936 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001937 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02001938 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001939 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001940 Py_RETURN_NONE;
1941}
1942
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001943PyDoc_STRVAR(chain_doc,
1944"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001945\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001946Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001947first iterable until it is exhausted, then elements from the next\n\
1948iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001949
Christian Heimesf16baeb2008-02-29 14:57:44 +00001950PyDoc_STRVAR(chain_from_iterable_doc,
1951"chain.from_iterable(iterable) --> chain object\n\
1952\n\
Martin Pantereb995702016-07-28 01:11:04 +00001953Alternate chain() constructor taking a single iterable argument\n\
Christian Heimesf16baeb2008-02-29 14:57:44 +00001954that evaluates lazily.");
1955
1956static PyMethodDef chain_methods[] = {
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001957 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1958 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001959 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1960 reduce_doc},
1961 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1962 setstate_doc},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001963 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001964};
1965
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001966static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001967 PyVarObject_HEAD_INIT(NULL, 0)
1968 "itertools.chain", /* tp_name */
1969 sizeof(chainobject), /* tp_basicsize */
1970 0, /* tp_itemsize */
1971 /* methods */
1972 (destructor)chain_dealloc, /* tp_dealloc */
1973 0, /* tp_print */
1974 0, /* tp_getattr */
1975 0, /* tp_setattr */
1976 0, /* tp_reserved */
1977 0, /* tp_repr */
1978 0, /* tp_as_number */
1979 0, /* tp_as_sequence */
1980 0, /* tp_as_mapping */
1981 0, /* tp_hash */
1982 0, /* tp_call */
1983 0, /* tp_str */
1984 PyObject_GenericGetAttr, /* tp_getattro */
1985 0, /* tp_setattro */
1986 0, /* tp_as_buffer */
1987 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1988 Py_TPFLAGS_BASETYPE, /* tp_flags */
1989 chain_doc, /* tp_doc */
1990 (traverseproc)chain_traverse, /* tp_traverse */
1991 0, /* tp_clear */
1992 0, /* tp_richcompare */
1993 0, /* tp_weaklistoffset */
1994 PyObject_SelfIter, /* tp_iter */
1995 (iternextfunc)chain_next, /* tp_iternext */
1996 chain_methods, /* tp_methods */
1997 0, /* tp_members */
1998 0, /* tp_getset */
1999 0, /* tp_base */
2000 0, /* tp_dict */
2001 0, /* tp_descr_get */
2002 0, /* tp_descr_set */
2003 0, /* tp_dictoffset */
2004 0, /* tp_init */
2005 0, /* tp_alloc */
2006 chain_new, /* tp_new */
2007 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002008};
2009
2010
Christian Heimesc3f30c42008-02-22 16:37:40 +00002011/* product object ************************************************************/
2012
2013typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002014 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002015 PyObject *pools; /* tuple of pool tuples */
2016 Py_ssize_t *indices; /* one index per pool */
2017 PyObject *result; /* most recently returned result tuple */
2018 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002019} productobject;
2020
2021static PyTypeObject product_type;
2022
2023static PyObject *
2024product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2025{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002026 productobject *lz;
2027 Py_ssize_t nargs, npools, repeat=1;
2028 PyObject *pools = NULL;
2029 Py_ssize_t *indices = NULL;
2030 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002031
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002032 if (kwds != NULL) {
2033 char *kwlist[] = {"repeat", 0};
2034 PyObject *tmpargs = PyTuple_New(0);
2035 if (tmpargs == NULL)
2036 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002037 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2038 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002039 Py_DECREF(tmpargs);
2040 return NULL;
2041 }
2042 Py_DECREF(tmpargs);
2043 if (repeat < 0) {
2044 PyErr_SetString(PyExc_ValueError,
2045 "repeat argument cannot be negative");
2046 return NULL;
2047 }
2048 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002049
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002050 assert(PyTuple_CheckExact(args));
2051 if (repeat == 0) {
2052 nargs = 0;
2053 } else {
2054 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002055 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002056 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2057 return NULL;
2058 }
2059 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002060 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002061
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002062 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002063 if (indices == NULL) {
2064 PyErr_NoMemory();
2065 goto error;
2066 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002067
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002068 pools = PyTuple_New(npools);
2069 if (pools == NULL)
2070 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002071
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002072 for (i=0; i < nargs ; ++i) {
2073 PyObject *item = PyTuple_GET_ITEM(args, i);
2074 PyObject *pool = PySequence_Tuple(item);
2075 if (pool == NULL)
2076 goto error;
2077 PyTuple_SET_ITEM(pools, i, pool);
2078 indices[i] = 0;
2079 }
2080 for ( ; i < npools; ++i) {
2081 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2082 Py_INCREF(pool);
2083 PyTuple_SET_ITEM(pools, i, pool);
2084 indices[i] = 0;
2085 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002086
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002087 /* create productobject structure */
2088 lz = (productobject *)type->tp_alloc(type, 0);
2089 if (lz == NULL)
2090 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002091
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002092 lz->pools = pools;
2093 lz->indices = indices;
2094 lz->result = NULL;
2095 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002096
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002097 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002098
2099error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002100 if (indices != NULL)
2101 PyMem_Free(indices);
2102 Py_XDECREF(pools);
2103 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002104}
2105
2106static void
2107product_dealloc(productobject *lz)
2108{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002109 PyObject_GC_UnTrack(lz);
2110 Py_XDECREF(lz->pools);
2111 Py_XDECREF(lz->result);
2112 if (lz->indices != NULL)
2113 PyMem_Free(lz->indices);
2114 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002115}
2116
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002117static PyObject *
2118product_sizeof(productobject *lz, void *unused)
2119{
2120 Py_ssize_t res;
2121
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002122 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002123 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2124 return PyLong_FromSsize_t(res);
2125}
2126
2127PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2128
Christian Heimesc3f30c42008-02-22 16:37:40 +00002129static int
2130product_traverse(productobject *lz, visitproc visit, void *arg)
2131{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002132 Py_VISIT(lz->pools);
2133 Py_VISIT(lz->result);
2134 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002135}
2136
2137static PyObject *
2138product_next(productobject *lz)
2139{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002140 PyObject *pool;
2141 PyObject *elem;
2142 PyObject *oldelem;
2143 PyObject *pools = lz->pools;
2144 PyObject *result = lz->result;
2145 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2146 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002147
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002148 if (lz->stopped)
2149 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002150
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002151 if (result == NULL) {
2152 /* On the first pass, return an initial tuple filled with the
2153 first element from each pool. */
2154 result = PyTuple_New(npools);
2155 if (result == NULL)
2156 goto empty;
2157 lz->result = result;
2158 for (i=0; i < npools; i++) {
2159 pool = PyTuple_GET_ITEM(pools, i);
2160 if (PyTuple_GET_SIZE(pool) == 0)
2161 goto empty;
2162 elem = PyTuple_GET_ITEM(pool, 0);
2163 Py_INCREF(elem);
2164 PyTuple_SET_ITEM(result, i, elem);
2165 }
2166 } else {
2167 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002168
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002169 /* Copy the previous result tuple or re-use it if available */
2170 if (Py_REFCNT(result) > 1) {
2171 PyObject *old_result = result;
2172 result = PyTuple_New(npools);
2173 if (result == NULL)
2174 goto empty;
2175 lz->result = result;
2176 for (i=0; i < npools; i++) {
2177 elem = PyTuple_GET_ITEM(old_result, i);
2178 Py_INCREF(elem);
2179 PyTuple_SET_ITEM(result, i, elem);
2180 }
2181 Py_DECREF(old_result);
2182 }
2183 /* Now, we've got the only copy so we can update it in-place */
2184 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002185
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002186 /* Update the pool indices right-to-left. Only advance to the
2187 next pool when the previous one rolls-over */
2188 for (i=npools-1 ; i >= 0 ; i--) {
2189 pool = PyTuple_GET_ITEM(pools, i);
2190 indices[i]++;
2191 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2192 /* Roll-over and advance to next pool */
2193 indices[i] = 0;
2194 elem = PyTuple_GET_ITEM(pool, 0);
2195 Py_INCREF(elem);
2196 oldelem = PyTuple_GET_ITEM(result, i);
2197 PyTuple_SET_ITEM(result, i, elem);
2198 Py_DECREF(oldelem);
2199 } else {
2200 /* No rollover. Just increment and stop here. */
2201 elem = PyTuple_GET_ITEM(pool, indices[i]);
2202 Py_INCREF(elem);
2203 oldelem = PyTuple_GET_ITEM(result, i);
2204 PyTuple_SET_ITEM(result, i, elem);
2205 Py_DECREF(oldelem);
2206 break;
2207 }
2208 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002209
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002210 /* If i is negative, then the indices have all rolled-over
2211 and we're done. */
2212 if (i < 0)
2213 goto empty;
2214 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002215
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002216 Py_INCREF(result);
2217 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002218
2219empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002220 lz->stopped = 1;
2221 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002222}
2223
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002224static PyObject *
2225product_reduce(productobject *lz)
2226{
2227 if (lz->stopped) {
2228 return Py_BuildValue("O(())", Py_TYPE(lz));
2229 } else if (lz->result == NULL) {
2230 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2231 } else {
2232 PyObject *indices;
2233 Py_ssize_t n, i;
2234
2235 /* we must pickle the indices use them for setstate, and
2236 * additionally indicate that the iterator has started
2237 */
2238 n = PyTuple_GET_SIZE(lz->pools);
2239 indices = PyTuple_New(n);
2240 if (indices == NULL)
2241 return NULL;
2242 for (i=0; i<n; i++){
2243 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2244 if (!index) {
2245 Py_DECREF(indices);
2246 return NULL;
2247 }
2248 PyTuple_SET_ITEM(indices, i, index);
2249 }
2250 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2251 }
2252}
2253
2254static PyObject *
2255product_setstate(productobject *lz, PyObject *state)
2256{
2257 PyObject *result;
2258 Py_ssize_t n, i;
2259
2260 n = PyTuple_GET_SIZE(lz->pools);
2261 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2262 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2263 return NULL;
2264 }
2265 for (i=0; i<n; i++)
2266 {
2267 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2268 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002269 PyObject* pool;
2270 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002271 if (index < 0 && PyErr_Occurred())
2272 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002273 pool = PyTuple_GET_ITEM(lz->pools, i);
2274 poolsize = PyTuple_GET_SIZE(pool);
2275 if (poolsize == 0) {
2276 lz->stopped = 1;
2277 Py_RETURN_NONE;
2278 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002279 /* clamp the index */
2280 if (index < 0)
2281 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002282 else if (index > poolsize-1)
2283 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002284 lz->indices[i] = index;
2285 }
2286
2287 result = PyTuple_New(n);
2288 if (!result)
2289 return NULL;
2290 for (i=0; i<n; i++) {
2291 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2292 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2293 Py_INCREF(element);
2294 PyTuple_SET_ITEM(result, i, element);
2295 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002296 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002297 Py_RETURN_NONE;
2298}
2299
2300static PyMethodDef product_methods[] = {
2301 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2302 reduce_doc},
2303 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2304 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002305 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2306 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002307 {NULL, NULL} /* sentinel */
2308};
2309
Christian Heimesc3f30c42008-02-22 16:37:40 +00002310PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002311"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002312\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002313Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002314For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2315The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2316cycle in a manner similar to an odometer (with the rightmost element changing\n\
2317on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002318To compute the product of an iterable with itself, specify the number\n\
2319of repetitions with the optional repeat keyword argument. For example,\n\
2320product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002321product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2322product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2323
2324static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002325 PyVarObject_HEAD_INIT(NULL, 0)
2326 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002327 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002328 0, /* tp_itemsize */
2329 /* methods */
2330 (destructor)product_dealloc, /* tp_dealloc */
2331 0, /* tp_print */
2332 0, /* tp_getattr */
2333 0, /* tp_setattr */
2334 0, /* tp_reserved */
2335 0, /* tp_repr */
2336 0, /* tp_as_number */
2337 0, /* tp_as_sequence */
2338 0, /* tp_as_mapping */
2339 0, /* tp_hash */
2340 0, /* tp_call */
2341 0, /* tp_str */
2342 PyObject_GenericGetAttr, /* tp_getattro */
2343 0, /* tp_setattro */
2344 0, /* tp_as_buffer */
2345 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2346 Py_TPFLAGS_BASETYPE, /* tp_flags */
2347 product_doc, /* tp_doc */
2348 (traverseproc)product_traverse, /* tp_traverse */
2349 0, /* tp_clear */
2350 0, /* tp_richcompare */
2351 0, /* tp_weaklistoffset */
2352 PyObject_SelfIter, /* tp_iter */
2353 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002354 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002355 0, /* tp_members */
2356 0, /* tp_getset */
2357 0, /* tp_base */
2358 0, /* tp_dict */
2359 0, /* tp_descr_get */
2360 0, /* tp_descr_set */
2361 0, /* tp_dictoffset */
2362 0, /* tp_init */
2363 0, /* tp_alloc */
2364 product_new, /* tp_new */
2365 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002366};
2367
2368
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002369/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002370
2371typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002372 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002373 PyObject *pool; /* input converted to a tuple */
2374 Py_ssize_t *indices; /* one index per result element */
2375 PyObject *result; /* most recently returned result tuple */
2376 Py_ssize_t r; /* size of result tuple */
2377 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002378} combinationsobject;
2379
2380static PyTypeObject combinations_type;
2381
2382static PyObject *
2383combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2384{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002385 combinationsobject *co;
2386 Py_ssize_t n;
2387 Py_ssize_t r;
2388 PyObject *pool = NULL;
2389 PyObject *iterable = NULL;
2390 Py_ssize_t *indices = NULL;
2391 Py_ssize_t i;
2392 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002393
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002394 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2395 &iterable, &r))
2396 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002397
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002398 pool = PySequence_Tuple(iterable);
2399 if (pool == NULL)
2400 goto error;
2401 n = PyTuple_GET_SIZE(pool);
2402 if (r < 0) {
2403 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2404 goto error;
2405 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002406
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002407 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002408 if (indices == NULL) {
2409 PyErr_NoMemory();
2410 goto error;
2411 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002412
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002413 for (i=0 ; i<r ; i++)
2414 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002415
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002416 /* create combinationsobject structure */
2417 co = (combinationsobject *)type->tp_alloc(type, 0);
2418 if (co == NULL)
2419 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002420
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002421 co->pool = pool;
2422 co->indices = indices;
2423 co->result = NULL;
2424 co->r = r;
2425 co->stopped = r > n ? 1 : 0;
2426
2427 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002428
2429error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002430 if (indices != NULL)
2431 PyMem_Free(indices);
2432 Py_XDECREF(pool);
2433 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002434}
2435
2436static void
2437combinations_dealloc(combinationsobject *co)
2438{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002439 PyObject_GC_UnTrack(co);
2440 Py_XDECREF(co->pool);
2441 Py_XDECREF(co->result);
2442 if (co->indices != NULL)
2443 PyMem_Free(co->indices);
2444 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002445}
2446
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002447static PyObject *
2448combinations_sizeof(combinationsobject *co, void *unused)
2449{
2450 Py_ssize_t res;
2451
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002452 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002453 res += co->r * sizeof(Py_ssize_t);
2454 return PyLong_FromSsize_t(res);
2455}
2456
Christian Heimes380f7f22008-02-28 11:19:05 +00002457static int
2458combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2459{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002460 Py_VISIT(co->pool);
2461 Py_VISIT(co->result);
2462 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002463}
2464
2465static PyObject *
2466combinations_next(combinationsobject *co)
2467{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002468 PyObject *elem;
2469 PyObject *oldelem;
2470 PyObject *pool = co->pool;
2471 Py_ssize_t *indices = co->indices;
2472 PyObject *result = co->result;
2473 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2474 Py_ssize_t r = co->r;
2475 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002476
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002477 if (co->stopped)
2478 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002479
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002480 if (result == NULL) {
2481 /* On the first pass, initialize result tuple using the indices */
2482 result = PyTuple_New(r);
2483 if (result == NULL)
2484 goto empty;
2485 co->result = result;
2486 for (i=0; i<r ; i++) {
2487 index = indices[i];
2488 elem = PyTuple_GET_ITEM(pool, index);
2489 Py_INCREF(elem);
2490 PyTuple_SET_ITEM(result, i, elem);
2491 }
2492 } else {
2493 /* Copy the previous result tuple or re-use it if available */
2494 if (Py_REFCNT(result) > 1) {
2495 PyObject *old_result = result;
2496 result = PyTuple_New(r);
2497 if (result == NULL)
2498 goto empty;
2499 co->result = result;
2500 for (i=0; i<r ; i++) {
2501 elem = PyTuple_GET_ITEM(old_result, i);
2502 Py_INCREF(elem);
2503 PyTuple_SET_ITEM(result, i, elem);
2504 }
2505 Py_DECREF(old_result);
2506 }
2507 /* Now, we've got the only copy so we can update it in-place
2508 * CPython's empty tuple is a singleton and cached in
2509 * PyTuple's freelist.
2510 */
2511 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002512
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002513 /* Scan indices right-to-left until finding one that is not
2514 at its maximum (i + n - r). */
2515 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2516 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002517
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002518 /* If i is negative, then the indices are all at
2519 their maximum value and we're done. */
2520 if (i < 0)
2521 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002522
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002523 /* Increment the current index which we know is not at its
2524 maximum. Then move back to the right setting each index
2525 to its lowest possible value (one higher than the index
2526 to its left -- this maintains the sort order invariant). */
2527 indices[i]++;
2528 for (j=i+1 ; j<r ; j++)
2529 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002530
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002531 /* Update the result tuple for the new indices
2532 starting with i, the leftmost index that changed */
2533 for ( ; i<r ; i++) {
2534 index = indices[i];
2535 elem = PyTuple_GET_ITEM(pool, index);
2536 Py_INCREF(elem);
2537 oldelem = PyTuple_GET_ITEM(result, i);
2538 PyTuple_SET_ITEM(result, i, elem);
2539 Py_DECREF(oldelem);
2540 }
2541 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002542
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002543 Py_INCREF(result);
2544 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002545
2546empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002547 co->stopped = 1;
2548 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002549}
2550
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002551static PyObject *
2552combinations_reduce(combinationsobject *lz)
2553{
2554 if (lz->result == NULL) {
2555 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2556 } else if (lz->stopped) {
2557 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2558 } else {
2559 PyObject *indices;
2560 Py_ssize_t i;
2561
2562 /* we must pickle the indices and use them for setstate */
2563 indices = PyTuple_New(lz->r);
2564 if (!indices)
2565 return NULL;
2566 for (i=0; i<lz->r; i++)
2567 {
2568 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2569 if (!index) {
2570 Py_DECREF(indices);
2571 return NULL;
2572 }
2573 PyTuple_SET_ITEM(indices, i, index);
2574 }
2575
2576 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2577 }
2578}
2579
2580static PyObject *
2581combinations_setstate(combinationsobject *lz, PyObject *state)
2582{
2583 PyObject *result;
2584 Py_ssize_t i;
2585 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2586
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002587 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002588 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2589 return NULL;
2590 }
2591
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002592 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002593 Py_ssize_t max;
2594 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2595 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002596
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002597 if (index == -1 && PyErr_Occurred())
2598 return NULL; /* not an integer */
2599 max = i + n - lz->r;
2600 /* clamp the index (beware of negative max) */
2601 if (index > max)
2602 index = max;
2603 if (index < 0)
2604 index = 0;
2605 lz->indices[i] = index;
2606 }
2607
2608 result = PyTuple_New(lz->r);
2609 if (result == NULL)
2610 return NULL;
2611 for (i=0; i<lz->r; i++) {
2612 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2613 Py_INCREF(element);
2614 PyTuple_SET_ITEM(result, i, element);
2615 }
2616
Serhiy Storchaka48842712016-04-06 09:45:48 +03002617 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002618 Py_RETURN_NONE;
2619}
2620
2621static PyMethodDef combinations_methods[] = {
2622 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2623 reduce_doc},
2624 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2625 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002626 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2627 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002628 {NULL, NULL} /* sentinel */
2629};
2630
Christian Heimes380f7f22008-02-28 11:19:05 +00002631PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002632"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002633\n\
2634Return successive r-length combinations of elements in the iterable.\n\n\
2635combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2636
2637static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002638 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002639 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002640 sizeof(combinationsobject), /* tp_basicsize */
2641 0, /* tp_itemsize */
2642 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002643 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002644 0, /* tp_print */
2645 0, /* tp_getattr */
2646 0, /* tp_setattr */
2647 0, /* tp_reserved */
2648 0, /* tp_repr */
2649 0, /* tp_as_number */
2650 0, /* tp_as_sequence */
2651 0, /* tp_as_mapping */
2652 0, /* tp_hash */
2653 0, /* tp_call */
2654 0, /* tp_str */
2655 PyObject_GenericGetAttr, /* tp_getattro */
2656 0, /* tp_setattro */
2657 0, /* tp_as_buffer */
2658 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2659 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002660 combinations_doc, /* tp_doc */
2661 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002662 0, /* tp_clear */
2663 0, /* tp_richcompare */
2664 0, /* tp_weaklistoffset */
2665 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002666 (iternextfunc)combinations_next, /* tp_iternext */
2667 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002668 0, /* tp_members */
2669 0, /* tp_getset */
2670 0, /* tp_base */
2671 0, /* tp_dict */
2672 0, /* tp_descr_get */
2673 0, /* tp_descr_set */
2674 0, /* tp_dictoffset */
2675 0, /* tp_init */
2676 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002677 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002678 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002679};
2680
2681
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002682/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002683
2684/* Equivalent to:
2685
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002686 def combinations_with_replacement(iterable, r):
2687 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2688 # number items returned: (n+r-1)! / r! / (n-1)!
2689 pool = tuple(iterable)
2690 n = len(pool)
2691 indices = [0] * r
2692 yield tuple(pool[i] for i in indices)
2693 while 1:
2694 for i in reversed(range(r)):
2695 if indices[i] != n - 1:
2696 break
2697 else:
2698 return
2699 indices[i:] = [indices[i] + 1] * (r - i)
2700 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002701
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002702 def combinations_with_replacement2(iterable, r):
2703 'Alternate version that filters from product()'
2704 pool = tuple(iterable)
2705 n = len(pool)
2706 for indices in product(range(n), repeat=r):
2707 if sorted(indices) == list(indices):
2708 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002709*/
2710typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002711 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002712 PyObject *pool; /* input converted to a tuple */
2713 Py_ssize_t *indices; /* one index per result element */
2714 PyObject *result; /* most recently returned result tuple */
2715 Py_ssize_t r; /* size of result tuple */
2716 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002717} cwrobject;
2718
2719static PyTypeObject cwr_type;
2720
2721static PyObject *
2722cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2723{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002724 cwrobject *co;
2725 Py_ssize_t n;
2726 Py_ssize_t r;
2727 PyObject *pool = NULL;
2728 PyObject *iterable = NULL;
2729 Py_ssize_t *indices = NULL;
2730 Py_ssize_t i;
2731 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002732
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002733 if (!PyArg_ParseTupleAndKeywords(args, kwds,
2734 "On:combinations_with_replacement",
2735 kwargs, &iterable, &r))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002736 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002737
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002738 pool = PySequence_Tuple(iterable);
2739 if (pool == NULL)
2740 goto error;
2741 n = PyTuple_GET_SIZE(pool);
2742 if (r < 0) {
2743 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2744 goto error;
2745 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002746
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002747 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002748 if (indices == NULL) {
2749 PyErr_NoMemory();
2750 goto error;
2751 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002752
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002753 for (i=0 ; i<r ; i++)
2754 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002755
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002756 /* create cwrobject structure */
2757 co = (cwrobject *)type->tp_alloc(type, 0);
2758 if (co == NULL)
2759 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002760
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002761 co->pool = pool;
2762 co->indices = indices;
2763 co->result = NULL;
2764 co->r = r;
2765 co->stopped = !n && r;
2766
2767 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002768
2769error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002770 if (indices != NULL)
2771 PyMem_Free(indices);
2772 Py_XDECREF(pool);
2773 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002774}
2775
2776static void
2777cwr_dealloc(cwrobject *co)
2778{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002779 PyObject_GC_UnTrack(co);
2780 Py_XDECREF(co->pool);
2781 Py_XDECREF(co->result);
2782 if (co->indices != NULL)
2783 PyMem_Free(co->indices);
2784 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002785}
2786
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002787static PyObject *
2788cwr_sizeof(cwrobject *co, void *unused)
2789{
2790 Py_ssize_t res;
2791
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002792 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002793 res += co->r * sizeof(Py_ssize_t);
2794 return PyLong_FromSsize_t(res);
2795}
2796
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002797static int
2798cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2799{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002800 Py_VISIT(co->pool);
2801 Py_VISIT(co->result);
2802 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002803}
2804
2805static PyObject *
2806cwr_next(cwrobject *co)
2807{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002808 PyObject *elem;
2809 PyObject *oldelem;
2810 PyObject *pool = co->pool;
2811 Py_ssize_t *indices = co->indices;
2812 PyObject *result = co->result;
2813 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2814 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002815 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002816
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002817 if (co->stopped)
2818 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002819
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002820 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002821 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002822 result = PyTuple_New(r);
2823 if (result == NULL)
2824 goto empty;
2825 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002826 if (n > 0) {
2827 elem = PyTuple_GET_ITEM(pool, 0);
2828 for (i=0; i<r ; i++) {
2829 assert(indices[i] == 0);
2830 Py_INCREF(elem);
2831 PyTuple_SET_ITEM(result, i, elem);
2832 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002833 }
2834 } else {
2835 /* Copy the previous result tuple or re-use it if available */
2836 if (Py_REFCNT(result) > 1) {
2837 PyObject *old_result = result;
2838 result = PyTuple_New(r);
2839 if (result == NULL)
2840 goto empty;
2841 co->result = result;
2842 for (i=0; i<r ; i++) {
2843 elem = PyTuple_GET_ITEM(old_result, i);
2844 Py_INCREF(elem);
2845 PyTuple_SET_ITEM(result, i, elem);
2846 }
2847 Py_DECREF(old_result);
2848 }
2849 /* Now, we've got the only copy so we can update it in-place CPython's
2850 empty tuple is a singleton and cached in PyTuple's freelist. */
2851 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002852
Tim Peters9edb1682013-09-03 11:49:31 -05002853 /* Scan indices right-to-left until finding one that is not
2854 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002855 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2856 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002857
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002858 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002859 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002860 if (i < 0)
2861 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002862
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002863 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002864 maximum. Then set all to the right to the same value. */
2865 index = indices[i] + 1;
2866 assert(index < n);
2867 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002868 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002869 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002870 Py_INCREF(elem);
2871 oldelem = PyTuple_GET_ITEM(result, i);
2872 PyTuple_SET_ITEM(result, i, elem);
2873 Py_DECREF(oldelem);
2874 }
2875 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002876
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002877 Py_INCREF(result);
2878 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002879
2880empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002881 co->stopped = 1;
2882 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002883}
2884
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002885static PyObject *
2886cwr_reduce(cwrobject *lz)
2887{
2888 if (lz->result == NULL) {
2889 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2890 } else if (lz->stopped) {
2891 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2892 } else {
2893 PyObject *indices;
2894 Py_ssize_t i;
2895
2896 /* we must pickle the indices and use them for setstate */
2897 indices = PyTuple_New(lz->r);
2898 if (!indices)
2899 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002900 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002901 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2902 if (!index) {
2903 Py_DECREF(indices);
2904 return NULL;
2905 }
2906 PyTuple_SET_ITEM(indices, i, index);
2907 }
2908
2909 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2910 }
2911}
2912
2913static PyObject *
2914cwr_setstate(cwrobject *lz, PyObject *state)
2915{
2916 PyObject *result;
2917 Py_ssize_t n, i;
2918
2919 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2920 {
2921 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2922 return NULL;
2923 }
2924
2925 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002926 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002927 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2928 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002929
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002930 if (index < 0 && PyErr_Occurred())
2931 return NULL; /* not an integer */
2932 /* clamp the index */
2933 if (index < 0)
2934 index = 0;
2935 else if (index > n-1)
2936 index = n-1;
2937 lz->indices[i] = index;
2938 }
2939 result = PyTuple_New(lz->r);
2940 if (result == NULL)
2941 return NULL;
2942 for (i=0; i<lz->r; i++) {
2943 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2944 Py_INCREF(element);
2945 PyTuple_SET_ITEM(result, i, element);
2946 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002947 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002948 Py_RETURN_NONE;
2949}
2950
2951static PyMethodDef cwr_methods[] = {
2952 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2953 reduce_doc},
2954 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2955 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002956 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
2957 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002958 {NULL, NULL} /* sentinel */
2959};
2960
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002961PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002962"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002963\n\
2964Return successive r-length combinations of elements in the iterable\n\
2965allowing individual elements to have successive repeats.\n\
2966combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2967
2968static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002969 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002970 "itertools.combinations_with_replacement", /* tp_name */
2971 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002972 0, /* tp_itemsize */
2973 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002974 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002975 0, /* tp_print */
2976 0, /* tp_getattr */
2977 0, /* tp_setattr */
2978 0, /* tp_reserved */
2979 0, /* tp_repr */
2980 0, /* tp_as_number */
2981 0, /* tp_as_sequence */
2982 0, /* tp_as_mapping */
2983 0, /* tp_hash */
2984 0, /* tp_call */
2985 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002986 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002987 0, /* tp_setattro */
2988 0, /* tp_as_buffer */
2989 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002990 Py_TPFLAGS_BASETYPE, /* tp_flags */
2991 cwr_doc, /* tp_doc */
2992 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002993 0, /* tp_clear */
2994 0, /* tp_richcompare */
2995 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002996 PyObject_SelfIter, /* tp_iter */
2997 (iternextfunc)cwr_next, /* tp_iternext */
2998 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002999 0, /* tp_members */
3000 0, /* tp_getset */
3001 0, /* tp_base */
3002 0, /* tp_dict */
3003 0, /* tp_descr_get */
3004 0, /* tp_descr_set */
3005 0, /* tp_dictoffset */
3006 0, /* tp_init */
3007 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003008 cwr_new, /* tp_new */
3009 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003010};
3011
3012
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003013/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003014
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003015def permutations(iterable, r=None):
3016 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
3017 pool = tuple(iterable)
3018 n = len(pool)
3019 r = n if r is None else r
3020 indices = range(n)
3021 cycles = range(n-r+1, n+1)[::-1]
3022 yield tuple(pool[i] for i in indices[:r])
3023 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003024 for i in reversed(range(r)):
3025 cycles[i] -= 1
3026 if cycles[i] == 0:
3027 indices[i:] = indices[i+1:] + indices[i:i+1]
3028 cycles[i] = n - i
3029 else:
3030 j = cycles[i]
3031 indices[i], indices[-j] = indices[-j], indices[i]
3032 yield tuple(pool[i] for i in indices[:r])
3033 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003034 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003035 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003036*/
3037
3038typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003039 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003040 PyObject *pool; /* input converted to a tuple */
3041 Py_ssize_t *indices; /* one index per element in the pool */
3042 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3043 PyObject *result; /* most recently returned result tuple */
3044 Py_ssize_t r; /* size of result tuple */
3045 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003046} permutationsobject;
3047
3048static PyTypeObject permutations_type;
3049
3050static PyObject *
3051permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3052{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003053 permutationsobject *po;
3054 Py_ssize_t n;
3055 Py_ssize_t r;
3056 PyObject *robj = Py_None;
3057 PyObject *pool = NULL;
3058 PyObject *iterable = NULL;
3059 Py_ssize_t *indices = NULL;
3060 Py_ssize_t *cycles = NULL;
3061 Py_ssize_t i;
3062 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003063
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003064 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
3065 &iterable, &robj))
3066 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003067
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003068 pool = PySequence_Tuple(iterable);
3069 if (pool == NULL)
3070 goto error;
3071 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003072
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003073 r = n;
3074 if (robj != Py_None) {
3075 if (!PyLong_Check(robj)) {
3076 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3077 goto error;
3078 }
3079 r = PyLong_AsSsize_t(robj);
3080 if (r == -1 && PyErr_Occurred())
3081 goto error;
3082 }
3083 if (r < 0) {
3084 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3085 goto error;
3086 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003087
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003088 indices = PyMem_New(Py_ssize_t, n);
3089 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003090 if (indices == NULL || cycles == NULL) {
3091 PyErr_NoMemory();
3092 goto error;
3093 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003094
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003095 for (i=0 ; i<n ; i++)
3096 indices[i] = i;
3097 for (i=0 ; i<r ; i++)
3098 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003099
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003100 /* create permutationsobject structure */
3101 po = (permutationsobject *)type->tp_alloc(type, 0);
3102 if (po == NULL)
3103 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003104
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003105 po->pool = pool;
3106 po->indices = indices;
3107 po->cycles = cycles;
3108 po->result = NULL;
3109 po->r = r;
3110 po->stopped = r > n ? 1 : 0;
3111
3112 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003113
3114error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003115 if (indices != NULL)
3116 PyMem_Free(indices);
3117 if (cycles != NULL)
3118 PyMem_Free(cycles);
3119 Py_XDECREF(pool);
3120 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003121}
3122
3123static void
3124permutations_dealloc(permutationsobject *po)
3125{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003126 PyObject_GC_UnTrack(po);
3127 Py_XDECREF(po->pool);
3128 Py_XDECREF(po->result);
3129 PyMem_Free(po->indices);
3130 PyMem_Free(po->cycles);
3131 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003132}
3133
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003134static PyObject *
3135permutations_sizeof(permutationsobject *po, void *unused)
3136{
3137 Py_ssize_t res;
3138
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003139 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003140 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3141 res += po->r * sizeof(Py_ssize_t);
3142 return PyLong_FromSsize_t(res);
3143}
3144
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003145static int
3146permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3147{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003148 Py_VISIT(po->pool);
3149 Py_VISIT(po->result);
3150 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003151}
3152
3153static PyObject *
3154permutations_next(permutationsobject *po)
3155{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003156 PyObject *elem;
3157 PyObject *oldelem;
3158 PyObject *pool = po->pool;
3159 Py_ssize_t *indices = po->indices;
3160 Py_ssize_t *cycles = po->cycles;
3161 PyObject *result = po->result;
3162 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3163 Py_ssize_t r = po->r;
3164 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003165
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003166 if (po->stopped)
3167 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003168
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003169 if (result == NULL) {
3170 /* On the first pass, initialize result tuple using the indices */
3171 result = PyTuple_New(r);
3172 if (result == NULL)
3173 goto empty;
3174 po->result = result;
3175 for (i=0; i<r ; i++) {
3176 index = indices[i];
3177 elem = PyTuple_GET_ITEM(pool, index);
3178 Py_INCREF(elem);
3179 PyTuple_SET_ITEM(result, i, elem);
3180 }
3181 } else {
3182 if (n == 0)
3183 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003184
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003185 /* Copy the previous result tuple or re-use it if available */
3186 if (Py_REFCNT(result) > 1) {
3187 PyObject *old_result = result;
3188 result = PyTuple_New(r);
3189 if (result == NULL)
3190 goto empty;
3191 po->result = result;
3192 for (i=0; i<r ; i++) {
3193 elem = PyTuple_GET_ITEM(old_result, i);
3194 Py_INCREF(elem);
3195 PyTuple_SET_ITEM(result, i, elem);
3196 }
3197 Py_DECREF(old_result);
3198 }
3199 /* Now, we've got the only copy so we can update it in-place */
3200 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003201
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003202 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3203 for (i=r-1 ; i>=0 ; i--) {
3204 cycles[i] -= 1;
3205 if (cycles[i] == 0) {
3206 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3207 index = indices[i];
3208 for (j=i ; j<n-1 ; j++)
3209 indices[j] = indices[j+1];
3210 indices[n-1] = index;
3211 cycles[i] = n - i;
3212 } else {
3213 j = cycles[i];
3214 index = indices[i];
3215 indices[i] = indices[n-j];
3216 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003217
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003218 for (k=i; k<r ; k++) {
3219 /* start with i, the leftmost element that changed */
3220 /* yield tuple(pool[k] for k in indices[:r]) */
3221 index = indices[k];
3222 elem = PyTuple_GET_ITEM(pool, index);
3223 Py_INCREF(elem);
3224 oldelem = PyTuple_GET_ITEM(result, k);
3225 PyTuple_SET_ITEM(result, k, elem);
3226 Py_DECREF(oldelem);
3227 }
3228 break;
3229 }
3230 }
3231 /* If i is negative, then the cycles have all
3232 rolled-over and we're done. */
3233 if (i < 0)
3234 goto empty;
3235 }
3236 Py_INCREF(result);
3237 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003238
3239empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003240 po->stopped = 1;
3241 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003242}
3243
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003244static PyObject *
3245permutations_reduce(permutationsobject *po)
3246{
3247 if (po->result == NULL) {
3248 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3249 } else if (po->stopped) {
3250 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3251 } else {
3252 PyObject *indices=NULL, *cycles=NULL;
3253 Py_ssize_t n, i;
3254
3255 /* we must pickle the indices and cycles and use them for setstate */
3256 n = PyTuple_GET_SIZE(po->pool);
3257 indices = PyTuple_New(n);
3258 if (indices == NULL)
3259 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003260 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003261 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3262 if (!index)
3263 goto err;
3264 PyTuple_SET_ITEM(indices, i, index);
3265 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003266
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003267 cycles = PyTuple_New(po->r);
3268 if (cycles == NULL)
3269 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003270 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003271 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3272 if (!index)
3273 goto err;
3274 PyTuple_SET_ITEM(cycles, i, index);
3275 }
3276 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3277 po->pool, po->r,
3278 indices, cycles);
3279 err:
3280 Py_XDECREF(indices);
3281 Py_XDECREF(cycles);
3282 return NULL;
3283 }
3284}
3285
3286static PyObject *
3287permutations_setstate(permutationsobject *po, PyObject *state)
3288{
3289 PyObject *indices, *cycles, *result;
3290 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003291
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003292 if (!PyTuple_Check(state)) {
3293 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3294 return NULL;
3295 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003296 if (!PyArg_ParseTuple(state, "O!O!",
3297 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003298 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003299 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003300 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003301
3302 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003303 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003304 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3305 return NULL;
3306 }
3307
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003308 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003309 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3310 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3311 if (index < 0 && PyErr_Occurred())
3312 return NULL; /* not an integer */
3313 /* clamp the index */
3314 if (index < 0)
3315 index = 0;
3316 else if (index > n-1)
3317 index = n-1;
3318 po->indices[i] = index;
3319 }
3320
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003321 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003322 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3323 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3324 if (index < 0 && PyErr_Occurred())
3325 return NULL; /* not an integer */
3326 if (index < 1)
3327 index = 1;
3328 else if (index > n-i)
3329 index = n-i;
3330 po->cycles[i] = index;
3331 }
3332 result = PyTuple_New(po->r);
3333 if (result == NULL)
3334 return NULL;
3335 for (i=0; i<po->r; i++) {
3336 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3337 Py_INCREF(element);
3338 PyTuple_SET_ITEM(result, i, element);
3339 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003340 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003341 Py_RETURN_NONE;
3342}
3343
3344static PyMethodDef permuations_methods[] = {
3345 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3346 reduce_doc},
3347 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3348 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003349 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3350 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003351 {NULL, NULL} /* sentinel */
3352};
3353
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003354PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003355"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003356\n\
3357Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003358permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003359
3360static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003361 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003362 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003363 sizeof(permutationsobject), /* tp_basicsize */
3364 0, /* tp_itemsize */
3365 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003366 (destructor)permutations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003367 0, /* tp_print */
3368 0, /* tp_getattr */
3369 0, /* tp_setattr */
3370 0, /* tp_reserved */
3371 0, /* tp_repr */
3372 0, /* tp_as_number */
3373 0, /* tp_as_sequence */
3374 0, /* tp_as_mapping */
3375 0, /* tp_hash */
3376 0, /* tp_call */
3377 0, /* tp_str */
3378 PyObject_GenericGetAttr, /* tp_getattro */
3379 0, /* tp_setattro */
3380 0, /* tp_as_buffer */
3381 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3382 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003383 permutations_doc, /* tp_doc */
3384 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003385 0, /* tp_clear */
3386 0, /* tp_richcompare */
3387 0, /* tp_weaklistoffset */
3388 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003389 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003390 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003391 0, /* tp_members */
3392 0, /* tp_getset */
3393 0, /* tp_base */
3394 0, /* tp_dict */
3395 0, /* tp_descr_get */
3396 0, /* tp_descr_set */
3397 0, /* tp_dictoffset */
3398 0, /* tp_init */
3399 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003400 permutations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003401 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003402};
3403
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003404/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003405
3406typedef struct {
3407 PyObject_HEAD
3408 PyObject *total;
3409 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003410 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003411} accumulateobject;
3412
3413static PyTypeObject accumulate_type;
3414
3415static PyObject *
3416accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3417{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003418 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003419 PyObject *iterable;
3420 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003421 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003422 accumulateobject *lz;
3423
Raymond Hettinger5d446132011-03-27 18:52:10 -07003424 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3425 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003426 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003427
3428 /* Get iterator. */
3429 it = PyObject_GetIter(iterable);
3430 if (it == NULL)
3431 return NULL;
3432
Raymond Hettinger482ba772010-12-01 22:48:00 +00003433 /* create accumulateobject structure */
3434 lz = (accumulateobject *)type->tp_alloc(type, 0);
3435 if (lz == NULL) {
3436 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003437 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003438 }
3439
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003440 if (binop != Py_None) {
3441 Py_XINCREF(binop);
3442 lz->binop = binop;
3443 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003444 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003445 lz->it = it;
3446 return (PyObject *)lz;
3447}
3448
3449static void
3450accumulate_dealloc(accumulateobject *lz)
3451{
3452 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003453 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003454 Py_XDECREF(lz->total);
3455 Py_XDECREF(lz->it);
3456 Py_TYPE(lz)->tp_free(lz);
3457}
3458
3459static int
3460accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3461{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003462 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003463 Py_VISIT(lz->it);
3464 Py_VISIT(lz->total);
3465 return 0;
3466}
3467
3468static PyObject *
3469accumulate_next(accumulateobject *lz)
3470{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003471 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003472
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003473 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003474 if (val == NULL)
3475 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003476
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003477 if (lz->total == NULL) {
3478 Py_INCREF(val);
3479 lz->total = val;
3480 return lz->total;
3481 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003482
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003483 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003484 newtotal = PyNumber_Add(lz->total, val);
3485 else
3486 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003487 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003488 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003489 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003490
Raymond Hettinger482ba772010-12-01 22:48:00 +00003491 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003492 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003493 return newtotal;
3494}
3495
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003496static PyObject *
3497accumulate_reduce(accumulateobject *lz)
3498{
Serhiy Storchakad5516252016-03-06 14:00:45 +02003499 if (lz->total == Py_None) {
3500 PyObject *it;
3501
3502 if (PyType_Ready(&chain_type) < 0)
3503 return NULL;
3504 if (PyType_Ready(&islice_type) < 0)
3505 return NULL;
3506 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3507 lz->total, lz->it);
3508 if (it == NULL)
3509 return NULL;
3510 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3511 it, lz->binop ? lz->binop : Py_None);
3512 if (it == NULL)
3513 return NULL;
3514 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3515 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003516 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3517 lz->it, lz->binop?lz->binop:Py_None,
3518 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003519}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003520
3521static PyObject *
3522accumulate_setstate(accumulateobject *lz, PyObject *state)
3523{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003524 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003525 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003526 Py_RETURN_NONE;
3527}
3528
3529static PyMethodDef accumulate_methods[] = {
3530 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3531 reduce_doc},
3532 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3533 setstate_doc},
3534 {NULL, NULL} /* sentinel */
3535};
3536
Raymond Hettinger482ba772010-12-01 22:48:00 +00003537PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003538"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003539\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003540Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003541
3542static PyTypeObject accumulate_type = {
3543 PyVarObject_HEAD_INIT(NULL, 0)
3544 "itertools.accumulate", /* tp_name */
3545 sizeof(accumulateobject), /* tp_basicsize */
3546 0, /* tp_itemsize */
3547 /* methods */
3548 (destructor)accumulate_dealloc, /* tp_dealloc */
3549 0, /* tp_print */
3550 0, /* tp_getattr */
3551 0, /* tp_setattr */
3552 0, /* tp_reserved */
3553 0, /* tp_repr */
3554 0, /* tp_as_number */
3555 0, /* tp_as_sequence */
3556 0, /* tp_as_mapping */
3557 0, /* tp_hash */
3558 0, /* tp_call */
3559 0, /* tp_str */
3560 PyObject_GenericGetAttr, /* tp_getattro */
3561 0, /* tp_setattro */
3562 0, /* tp_as_buffer */
3563 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3564 Py_TPFLAGS_BASETYPE, /* tp_flags */
3565 accumulate_doc, /* tp_doc */
3566 (traverseproc)accumulate_traverse, /* tp_traverse */
3567 0, /* tp_clear */
3568 0, /* tp_richcompare */
3569 0, /* tp_weaklistoffset */
3570 PyObject_SelfIter, /* tp_iter */
3571 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003572 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003573 0, /* tp_members */
3574 0, /* tp_getset */
3575 0, /* tp_base */
3576 0, /* tp_dict */
3577 0, /* tp_descr_get */
3578 0, /* tp_descr_set */
3579 0, /* tp_dictoffset */
3580 0, /* tp_init */
3581 0, /* tp_alloc */
3582 accumulate_new, /* tp_new */
3583 PyObject_GC_Del, /* tp_free */
3584};
3585
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003586
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003587/* compress object ************************************************************/
3588
3589/* Equivalent to:
3590
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003591 def compress(data, selectors):
3592 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3593 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003594*/
3595
3596typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003597 PyObject_HEAD
3598 PyObject *data;
3599 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003600} compressobject;
3601
3602static PyTypeObject compress_type;
3603
3604static PyObject *
3605compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3606{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003607 PyObject *seq1, *seq2;
3608 PyObject *data=NULL, *selectors=NULL;
3609 compressobject *lz;
3610 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003611
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003612 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3613 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003614
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003615 data = PyObject_GetIter(seq1);
3616 if (data == NULL)
3617 goto fail;
3618 selectors = PyObject_GetIter(seq2);
3619 if (selectors == NULL)
3620 goto fail;
3621
3622 /* create compressobject structure */
3623 lz = (compressobject *)type->tp_alloc(type, 0);
3624 if (lz == NULL)
3625 goto fail;
3626 lz->data = data;
3627 lz->selectors = selectors;
3628 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003629
3630fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003631 Py_XDECREF(data);
3632 Py_XDECREF(selectors);
3633 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003634}
3635
3636static void
3637compress_dealloc(compressobject *lz)
3638{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003639 PyObject_GC_UnTrack(lz);
3640 Py_XDECREF(lz->data);
3641 Py_XDECREF(lz->selectors);
3642 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003643}
3644
3645static int
3646compress_traverse(compressobject *lz, visitproc visit, void *arg)
3647{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003648 Py_VISIT(lz->data);
3649 Py_VISIT(lz->selectors);
3650 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003651}
3652
3653static PyObject *
3654compress_next(compressobject *lz)
3655{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003656 PyObject *data = lz->data, *selectors = lz->selectors;
3657 PyObject *datum, *selector;
3658 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3659 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3660 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003661
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003662 while (1) {
3663 /* Steps: get datum, get selector, evaluate selector.
3664 Order is important (to match the pure python version
3665 in terms of which input gets a chance to raise an
3666 exception first).
3667 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003668
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003669 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003670 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003671 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003672
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003673 selector = selectornext(selectors);
3674 if (selector == NULL) {
3675 Py_DECREF(datum);
3676 return NULL;
3677 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003678
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003679 ok = PyObject_IsTrue(selector);
3680 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003681 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003682 return datum;
3683 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003684 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003685 return NULL;
3686 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003687}
3688
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003689static PyObject *
3690compress_reduce(compressobject *lz)
3691{
3692 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3693 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003694}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003695
3696static PyMethodDef compress_methods[] = {
3697 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3698 reduce_doc},
3699 {NULL, NULL} /* sentinel */
3700};
3701
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003702PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003703"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003704\n\
3705Return data elements corresponding to true selector elements.\n\
3706Forms a shorter iterator from selected data elements using the\n\
3707selectors to choose the data elements.");
3708
3709static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003710 PyVarObject_HEAD_INIT(NULL, 0)
3711 "itertools.compress", /* tp_name */
3712 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003713 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003714 /* methods */
3715 (destructor)compress_dealloc, /* tp_dealloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003716 0, /* tp_print */
3717 0, /* tp_getattr */
3718 0, /* tp_setattr */
3719 0, /* tp_reserved */
3720 0, /* tp_repr */
3721 0, /* tp_as_number */
3722 0, /* tp_as_sequence */
3723 0, /* tp_as_mapping */
3724 0, /* tp_hash */
3725 0, /* tp_call */
3726 0, /* tp_str */
3727 PyObject_GenericGetAttr, /* tp_getattro */
3728 0, /* tp_setattro */
3729 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003730 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003731 Py_TPFLAGS_BASETYPE, /* tp_flags */
3732 compress_doc, /* tp_doc */
3733 (traverseproc)compress_traverse, /* tp_traverse */
3734 0, /* tp_clear */
3735 0, /* tp_richcompare */
3736 0, /* tp_weaklistoffset */
3737 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003738 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003739 compress_methods, /* tp_methods */
3740 0, /* tp_members */
3741 0, /* tp_getset */
3742 0, /* tp_base */
3743 0, /* tp_dict */
3744 0, /* tp_descr_get */
3745 0, /* tp_descr_set */
3746 0, /* tp_dictoffset */
3747 0, /* tp_init */
3748 0, /* tp_alloc */
3749 compress_new, /* tp_new */
3750 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003751};
3752
3753
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003754/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003755
3756typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003757 PyObject_HEAD
3758 PyObject *func;
3759 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003760} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003761
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003762static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003763
3764static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003765filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003766{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003767 PyObject *func, *seq;
3768 PyObject *it;
3769 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003770
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003771 if (type == &filterfalse_type &&
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03003772 !_PyArg_NoKeywords("filterfalse", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003773 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003774
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003775 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3776 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003777
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003778 /* Get iterator. */
3779 it = PyObject_GetIter(seq);
3780 if (it == NULL)
3781 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003782
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003783 /* create filterfalseobject structure */
3784 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3785 if (lz == NULL) {
3786 Py_DECREF(it);
3787 return NULL;
3788 }
3789 Py_INCREF(func);
3790 lz->func = func;
3791 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003792
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003793 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003794}
3795
3796static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003797filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003798{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003799 PyObject_GC_UnTrack(lz);
3800 Py_XDECREF(lz->func);
3801 Py_XDECREF(lz->it);
3802 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003803}
3804
3805static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003806filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003807{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003808 Py_VISIT(lz->it);
3809 Py_VISIT(lz->func);
3810 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003811}
3812
3813static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003814filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003815{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003816 PyObject *item;
3817 PyObject *it = lz->it;
3818 long ok;
3819 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003820
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003821 iternext = *Py_TYPE(it)->tp_iternext;
3822 for (;;) {
3823 item = iternext(it);
3824 if (item == NULL)
3825 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003826
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003827 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3828 ok = PyObject_IsTrue(item);
3829 } else {
3830 PyObject *good;
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01003831 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003832 if (good == NULL) {
3833 Py_DECREF(item);
3834 return NULL;
3835 }
3836 ok = PyObject_IsTrue(good);
3837 Py_DECREF(good);
3838 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003839 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003840 return item;
3841 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003842 if (ok < 0)
3843 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003844 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003845}
3846
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003847static PyObject *
3848filterfalse_reduce(filterfalseobject *lz)
3849{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003850 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3851}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003852
3853static PyMethodDef filterfalse_methods[] = {
3854 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3855 reduce_doc},
3856 {NULL, NULL} /* sentinel */
3857};
3858
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003859PyDoc_STRVAR(filterfalse_doc,
3860"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003861\n\
3862Return those items of sequence for which function(item) is false.\n\
3863If function is None, return the items that are false.");
3864
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003865static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003866 PyVarObject_HEAD_INIT(NULL, 0)
3867 "itertools.filterfalse", /* tp_name */
3868 sizeof(filterfalseobject), /* tp_basicsize */
3869 0, /* tp_itemsize */
3870 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003871 (destructor)filterfalse_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003872 0, /* tp_print */
3873 0, /* tp_getattr */
3874 0, /* tp_setattr */
3875 0, /* tp_reserved */
3876 0, /* tp_repr */
3877 0, /* tp_as_number */
3878 0, /* tp_as_sequence */
3879 0, /* tp_as_mapping */
3880 0, /* tp_hash */
3881 0, /* tp_call */
3882 0, /* tp_str */
3883 PyObject_GenericGetAttr, /* tp_getattro */
3884 0, /* tp_setattro */
3885 0, /* tp_as_buffer */
3886 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3887 Py_TPFLAGS_BASETYPE, /* tp_flags */
3888 filterfalse_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003889 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003890 0, /* tp_clear */
3891 0, /* tp_richcompare */
3892 0, /* tp_weaklistoffset */
3893 PyObject_SelfIter, /* tp_iter */
3894 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003895 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003896 0, /* tp_members */
3897 0, /* tp_getset */
3898 0, /* tp_base */
3899 0, /* tp_dict */
3900 0, /* tp_descr_get */
3901 0, /* tp_descr_set */
3902 0, /* tp_dictoffset */
3903 0, /* tp_init */
3904 0, /* tp_alloc */
3905 filterfalse_new, /* tp_new */
3906 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003907};
3908
3909
3910/* count object ************************************************************/
3911
3912typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003913 PyObject_HEAD
3914 Py_ssize_t cnt;
3915 PyObject *long_cnt;
3916 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003917} countobject;
3918
Raymond Hettinger30729212009-02-12 06:28:27 +00003919/* Counting logic and invariants:
3920
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003921fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003922
Serhiy Storchaka483405b2015-02-17 10:14:30 +02003923 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003924 Advances with: cnt += 1
3925 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003926
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003927slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003928
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003929 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3930 All counting is done with python objects (no overflows or underflows).
3931 Advances with: long_cnt += long_step
3932 Step may be zero -- effectively a slow version of repeat(cnt).
3933 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003934*/
3935
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003936static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003937
3938static PyObject *
3939count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3940{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003941 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003942 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003943 Py_ssize_t cnt = 0;
3944 PyObject *long_cnt = NULL;
3945 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003946 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003947 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003948
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003949 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3950 kwlist, &long_cnt, &long_step))
3951 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003952
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003953 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3954 (long_step != NULL && !PyNumber_Check(long_step))) {
3955 PyErr_SetString(PyExc_TypeError, "a number is required");
3956 return NULL;
3957 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003958
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003959 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
3960 (long_step == NULL || PyLong_Check(long_step));
3961
3962 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003963 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003964 if (fast_mode) {
3965 assert(PyLong_Check(long_cnt));
3966 cnt = PyLong_AsSsize_t(long_cnt);
3967 if (cnt == -1 && PyErr_Occurred()) {
3968 PyErr_Clear();
3969 fast_mode = 0;
3970 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003971 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003972 } else {
3973 cnt = 0;
Serhiy Storchakaba85d692017-03-30 09:09:41 +03003974 long_cnt = _PyLong_Zero;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003975 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03003976 Py_INCREF(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003977
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003978 /* If not specified, step defaults to 1 */
Serhiy Storchakaba85d692017-03-30 09:09:41 +03003979 if (long_step == NULL)
3980 long_step = _PyLong_One;
3981 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003982
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003983 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003984
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003985 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003986 if (fast_mode) {
3987 assert(PyLong_Check(long_step));
3988 step = PyLong_AsLong(long_step);
3989 if (step != 1) {
3990 fast_mode = 0;
3991 if (step == -1 && PyErr_Occurred())
3992 PyErr_Clear();
3993 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003994 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003995
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003996 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003997 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003998 else
3999 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004000
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004001 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4002 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4003 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004004 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004005
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004006 /* create countobject structure */
4007 lz = (countobject *)type->tp_alloc(type, 0);
4008 if (lz == NULL) {
4009 Py_XDECREF(long_cnt);
4010 return NULL;
4011 }
4012 lz->cnt = cnt;
4013 lz->long_cnt = long_cnt;
4014 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004015
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004016 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004017}
4018
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004019static void
4020count_dealloc(countobject *lz)
4021{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004022 PyObject_GC_UnTrack(lz);
4023 Py_XDECREF(lz->long_cnt);
4024 Py_XDECREF(lz->long_step);
4025 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004026}
4027
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004028static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004029count_traverse(countobject *lz, visitproc visit, void *arg)
4030{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004031 Py_VISIT(lz->long_cnt);
4032 Py_VISIT(lz->long_step);
4033 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004034}
4035
4036static PyObject *
4037count_nextlong(countobject *lz)
4038{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004039 PyObject *long_cnt;
4040 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004041
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004042 long_cnt = lz->long_cnt;
4043 if (long_cnt == NULL) {
4044 /* Switch to slow_mode */
4045 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4046 if (long_cnt == NULL)
4047 return NULL;
4048 }
4049 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004050
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004051 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4052 if (stepped_up == NULL)
4053 return NULL;
4054 lz->long_cnt = stepped_up;
4055 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004056}
4057
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004058static PyObject *
4059count_next(countobject *lz)
4060{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004061 if (lz->cnt == PY_SSIZE_T_MAX)
4062 return count_nextlong(lz);
4063 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004064}
4065
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004066static PyObject *
4067count_repr(countobject *lz)
4068{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004069 if (lz->cnt != PY_SSIZE_T_MAX)
4070 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004071
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004072 if (PyLong_Check(lz->long_step)) {
4073 long step = PyLong_AsLong(lz->long_step);
4074 if (step == -1 && PyErr_Occurred()) {
4075 PyErr_Clear();
4076 }
4077 if (step == 1) {
4078 /* Don't display step when it is an integer equal to 1 */
4079 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
4080 }
4081 }
4082 return PyUnicode_FromFormat("count(%R, %R)",
4083 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004084}
4085
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004086static PyObject *
4087count_reduce(countobject *lz)
4088{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004089 if (lz->cnt == PY_SSIZE_T_MAX)
4090 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4091 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004092}
4093
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004094static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004095 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004096 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004097 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004098};
4099
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004100PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00004101 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004102\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004103Return a count object whose .__next__() method returns consecutive values.\n\
4104Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00004105 def count(firstval=0, step=1):\n\
Eli Benderskyc554f722013-09-03 06:37:19 -07004106 x = firstval\n\
4107 while 1:\n\
4108 yield x\n\
4109 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004110
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004111static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004112 PyVarObject_HEAD_INIT(NULL, 0)
4113 "itertools.count", /* tp_name */
4114 sizeof(countobject), /* tp_basicsize */
4115 0, /* tp_itemsize */
4116 /* methods */
4117 (destructor)count_dealloc, /* tp_dealloc */
4118 0, /* tp_print */
4119 0, /* tp_getattr */
4120 0, /* tp_setattr */
4121 0, /* tp_reserved */
4122 (reprfunc)count_repr, /* tp_repr */
4123 0, /* tp_as_number */
4124 0, /* tp_as_sequence */
4125 0, /* tp_as_mapping */
4126 0, /* tp_hash */
4127 0, /* tp_call */
4128 0, /* tp_str */
4129 PyObject_GenericGetAttr, /* tp_getattro */
4130 0, /* tp_setattro */
4131 0, /* tp_as_buffer */
4132 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004133 Py_TPFLAGS_BASETYPE, /* tp_flags */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004134 count_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004135 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004136 0, /* tp_clear */
4137 0, /* tp_richcompare */
4138 0, /* tp_weaklistoffset */
4139 PyObject_SelfIter, /* tp_iter */
4140 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004141 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004142 0, /* tp_members */
4143 0, /* tp_getset */
4144 0, /* tp_base */
4145 0, /* tp_dict */
4146 0, /* tp_descr_get */
4147 0, /* tp_descr_set */
4148 0, /* tp_dictoffset */
4149 0, /* tp_init */
4150 0, /* tp_alloc */
4151 count_new, /* tp_new */
4152 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004153};
4154
4155
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004156/* repeat object ************************************************************/
4157
4158typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004159 PyObject_HEAD
4160 PyObject *element;
4161 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004162} repeatobject;
4163
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004164static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004165
4166static PyObject *
4167repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4168{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004169 repeatobject *ro;
4170 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004171 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004172 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004173
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004174 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4175 &element, &cnt))
4176 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004177
Raymond Hettinger97d35552014-06-24 21:36:58 -07004178 if (kwds != NULL)
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004179 n_kwds = PyDict_GET_SIZE(kwds);
Raymond Hettinger97d35552014-06-24 21:36:58 -07004180 /* Does user supply times argument? */
4181 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004182 cnt = 0;
4183
4184 ro = (repeatobject *)type->tp_alloc(type, 0);
4185 if (ro == NULL)
4186 return NULL;
4187 Py_INCREF(element);
4188 ro->element = element;
4189 ro->cnt = cnt;
4190 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004191}
4192
4193static void
4194repeat_dealloc(repeatobject *ro)
4195{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004196 PyObject_GC_UnTrack(ro);
4197 Py_XDECREF(ro->element);
4198 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004199}
4200
4201static int
4202repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4203{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004204 Py_VISIT(ro->element);
4205 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004206}
4207
4208static PyObject *
4209repeat_next(repeatobject *ro)
4210{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004211 if (ro->cnt == 0)
4212 return NULL;
4213 if (ro->cnt > 0)
4214 ro->cnt--;
4215 Py_INCREF(ro->element);
4216 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004217}
4218
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004219static PyObject *
4220repeat_repr(repeatobject *ro)
4221{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004222 if (ro->cnt == -1)
4223 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4224 else
4225 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4226}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004227
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004228static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004229repeat_len(repeatobject *ro)
4230{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004231 if (ro->cnt == -1) {
4232 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4233 return NULL;
4234 }
4235 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004236}
4237
Armin Rigof5b3e362006-02-11 21:32:43 +00004238PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004239
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004240static PyObject *
4241repeat_reduce(repeatobject *ro)
4242{
4243 /* unpickle this so that a new repeat iterator is constructed with an
4244 * object, then call __setstate__ on it to set cnt
4245 */
4246 if (ro->cnt >= 0)
4247 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4248 else
4249 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4250}
4251
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004252static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004253 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004254 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004255 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004256};
4257
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004258PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004259"repeat(object [,times]) -> create an iterator which returns the object\n\
4260for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004261endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004262
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004263static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004264 PyVarObject_HEAD_INIT(NULL, 0)
4265 "itertools.repeat", /* tp_name */
4266 sizeof(repeatobject), /* tp_basicsize */
4267 0, /* tp_itemsize */
4268 /* methods */
4269 (destructor)repeat_dealloc, /* tp_dealloc */
4270 0, /* tp_print */
4271 0, /* tp_getattr */
4272 0, /* tp_setattr */
4273 0, /* tp_reserved */
4274 (reprfunc)repeat_repr, /* tp_repr */
4275 0, /* tp_as_number */
4276 0, /* tp_as_sequence */
4277 0, /* tp_as_mapping */
4278 0, /* tp_hash */
4279 0, /* tp_call */
4280 0, /* tp_str */
4281 PyObject_GenericGetAttr, /* tp_getattro */
4282 0, /* tp_setattro */
4283 0, /* tp_as_buffer */
4284 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4285 Py_TPFLAGS_BASETYPE, /* tp_flags */
4286 repeat_doc, /* tp_doc */
4287 (traverseproc)repeat_traverse, /* tp_traverse */
4288 0, /* tp_clear */
4289 0, /* tp_richcompare */
4290 0, /* tp_weaklistoffset */
4291 PyObject_SelfIter, /* tp_iter */
4292 (iternextfunc)repeat_next, /* tp_iternext */
4293 repeat_methods, /* tp_methods */
4294 0, /* tp_members */
4295 0, /* tp_getset */
4296 0, /* tp_base */
4297 0, /* tp_dict */
4298 0, /* tp_descr_get */
4299 0, /* tp_descr_set */
4300 0, /* tp_dictoffset */
4301 0, /* tp_init */
4302 0, /* tp_alloc */
4303 repeat_new, /* tp_new */
4304 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004305};
4306
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004307/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004308
4309typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004310 PyObject_HEAD
4311 Py_ssize_t tuplesize;
4312 Py_ssize_t numactive;
4313 PyObject *ittuple; /* tuple of iterators */
4314 PyObject *result;
4315 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004316} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004317
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004318static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004319
4320static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004321zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004322{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004323 ziplongestobject *lz;
4324 Py_ssize_t i;
4325 PyObject *ittuple; /* tuple of iterators */
4326 PyObject *result;
4327 PyObject *fillvalue = Py_None;
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004328 Py_ssize_t tuplesize;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004329
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004330 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004331 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004332 if (fillvalue == NULL || PyDict_GET_SIZE(kwds) > 1) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004333 PyErr_SetString(PyExc_TypeError,
4334 "zip_longest() got an unexpected keyword argument");
4335 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004336 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004337 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004338
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004339 /* args must be a tuple */
4340 assert(PyTuple_Check(args));
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004341 tuplesize = PyTuple_GET_SIZE(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004342
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004343 /* obtain iterators */
4344 ittuple = PyTuple_New(tuplesize);
4345 if (ittuple == NULL)
4346 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004347 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004348 PyObject *item = PyTuple_GET_ITEM(args, i);
4349 PyObject *it = PyObject_GetIter(item);
4350 if (it == NULL) {
4351 if (PyErr_ExceptionMatches(PyExc_TypeError))
4352 PyErr_Format(PyExc_TypeError,
4353 "zip_longest argument #%zd must support iteration",
4354 i+1);
4355 Py_DECREF(ittuple);
4356 return NULL;
4357 }
4358 PyTuple_SET_ITEM(ittuple, i, it);
4359 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004360
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004361 /* create a result holder */
4362 result = PyTuple_New(tuplesize);
4363 if (result == NULL) {
4364 Py_DECREF(ittuple);
4365 return NULL;
4366 }
4367 for (i=0 ; i < tuplesize ; i++) {
4368 Py_INCREF(Py_None);
4369 PyTuple_SET_ITEM(result, i, Py_None);
4370 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004371
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004372 /* create ziplongestobject structure */
4373 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4374 if (lz == NULL) {
4375 Py_DECREF(ittuple);
4376 Py_DECREF(result);
4377 return NULL;
4378 }
4379 lz->ittuple = ittuple;
4380 lz->tuplesize = tuplesize;
4381 lz->numactive = tuplesize;
4382 lz->result = result;
4383 Py_INCREF(fillvalue);
4384 lz->fillvalue = fillvalue;
4385 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004386}
4387
4388static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004389zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004390{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004391 PyObject_GC_UnTrack(lz);
4392 Py_XDECREF(lz->ittuple);
4393 Py_XDECREF(lz->result);
4394 Py_XDECREF(lz->fillvalue);
4395 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004396}
4397
4398static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004399zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004400{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004401 Py_VISIT(lz->ittuple);
4402 Py_VISIT(lz->result);
4403 Py_VISIT(lz->fillvalue);
4404 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004405}
4406
4407static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004408zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004409{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004410 Py_ssize_t i;
4411 Py_ssize_t tuplesize = lz->tuplesize;
4412 PyObject *result = lz->result;
4413 PyObject *it;
4414 PyObject *item;
4415 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004416
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004417 if (tuplesize == 0)
4418 return NULL;
4419 if (lz->numactive == 0)
4420 return NULL;
4421 if (Py_REFCNT(result) == 1) {
4422 Py_INCREF(result);
4423 for (i=0 ; i < tuplesize ; i++) {
4424 it = PyTuple_GET_ITEM(lz->ittuple, i);
4425 if (it == NULL) {
4426 Py_INCREF(lz->fillvalue);
4427 item = lz->fillvalue;
4428 } else {
4429 item = PyIter_Next(it);
4430 if (item == NULL) {
4431 lz->numactive -= 1;
4432 if (lz->numactive == 0 || PyErr_Occurred()) {
4433 lz->numactive = 0;
4434 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004435 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004436 } else {
4437 Py_INCREF(lz->fillvalue);
4438 item = lz->fillvalue;
4439 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4440 Py_DECREF(it);
4441 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004442 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004443 }
4444 olditem = PyTuple_GET_ITEM(result, i);
4445 PyTuple_SET_ITEM(result, i, item);
4446 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004447 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004448 } else {
4449 result = PyTuple_New(tuplesize);
4450 if (result == NULL)
4451 return NULL;
4452 for (i=0 ; i < tuplesize ; i++) {
4453 it = PyTuple_GET_ITEM(lz->ittuple, i);
4454 if (it == NULL) {
4455 Py_INCREF(lz->fillvalue);
4456 item = lz->fillvalue;
4457 } else {
4458 item = PyIter_Next(it);
4459 if (item == NULL) {
4460 lz->numactive -= 1;
4461 if (lz->numactive == 0 || PyErr_Occurred()) {
4462 lz->numactive = 0;
4463 Py_DECREF(result);
4464 return NULL;
4465 } else {
4466 Py_INCREF(lz->fillvalue);
4467 item = lz->fillvalue;
4468 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4469 Py_DECREF(it);
4470 }
4471 }
4472 }
4473 PyTuple_SET_ITEM(result, i, item);
4474 }
4475 }
4476 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004477}
4478
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004479static PyObject *
4480zip_longest_reduce(ziplongestobject *lz)
4481{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004482
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004483 /* Create a new tuple with empty sequences where appropriate to pickle.
4484 * Then use setstate to set the fillvalue
4485 */
4486 int i;
4487 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004488
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004489 if (args == NULL)
4490 return NULL;
4491 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4492 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4493 if (elem == NULL) {
4494 elem = PyTuple_New(0);
4495 if (elem == NULL) {
4496 Py_DECREF(args);
4497 return NULL;
4498 }
4499 } else
4500 Py_INCREF(elem);
4501 PyTuple_SET_ITEM(args, i, elem);
4502 }
4503 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4504}
4505
4506static PyObject *
4507zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4508{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004509 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004510 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004511 Py_RETURN_NONE;
4512}
4513
4514static PyMethodDef zip_longest_methods[] = {
4515 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4516 reduce_doc},
4517 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4518 setstate_doc},
4519 {NULL, NULL} /* sentinel */
4520};
4521
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004522PyDoc_STRVAR(zip_longest_doc,
4523"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004524\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004525Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004526the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004527method continues until the longest iterable in the argument sequence\n\
4528is exhausted and then it raises StopIteration. When the shorter iterables\n\
4529are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4530defaults to None or can be specified by a keyword argument.\n\
4531");
4532
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004533static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004534 PyVarObject_HEAD_INIT(NULL, 0)
4535 "itertools.zip_longest", /* tp_name */
4536 sizeof(ziplongestobject), /* tp_basicsize */
4537 0, /* tp_itemsize */
4538 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004539 (destructor)zip_longest_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004540 0, /* tp_print */
4541 0, /* tp_getattr */
4542 0, /* tp_setattr */
4543 0, /* tp_reserved */
4544 0, /* tp_repr */
4545 0, /* tp_as_number */
4546 0, /* tp_as_sequence */
4547 0, /* tp_as_mapping */
4548 0, /* tp_hash */
4549 0, /* tp_call */
4550 0, /* tp_str */
4551 PyObject_GenericGetAttr, /* tp_getattro */
4552 0, /* tp_setattro */
4553 0, /* tp_as_buffer */
4554 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4555 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004556 zip_longest_doc, /* tp_doc */
4557 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004558 0, /* tp_clear */
4559 0, /* tp_richcompare */
4560 0, /* tp_weaklistoffset */
4561 PyObject_SelfIter, /* tp_iter */
4562 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004563 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004564 0, /* tp_members */
4565 0, /* tp_getset */
4566 0, /* tp_base */
4567 0, /* tp_dict */
4568 0, /* tp_descr_get */
4569 0, /* tp_descr_set */
4570 0, /* tp_dictoffset */
4571 0, /* tp_init */
4572 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004573 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004574 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004575};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004576
4577/* module level code ********************************************************/
4578
4579PyDoc_STRVAR(module_doc,
4580"Functional tools for creating and using iterators.\n\
4581\n\
4582Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004583count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004584cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004585repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004586\n\
4587Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004588accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004589chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger8df58f72013-09-09 01:29:40 -05004590chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004591compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4592dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4593groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004594filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004595islice(seq, [start,] stop [, step]) --> elements from\n\
4596 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004597starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004598tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004599takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004600zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004601\n\
4602Combinatoric generators:\n\
4603product(p, q, ... [repeat=1]) --> cartesian product\n\
4604permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004605combinations(p, r)\n\
4606combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004607");
4608
4609
Raymond Hettingerad983e72003-11-12 14:32:26 +00004610static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004611 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4612 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004613};
4614
Martin v. Löwis1a214512008-06-11 05:26:20 +00004615
4616static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004617 PyModuleDef_HEAD_INIT,
4618 "itertools",
4619 module_doc,
4620 -1,
4621 module_methods,
4622 NULL,
4623 NULL,
4624 NULL,
4625 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004626};
4627
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004628PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004629PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004630{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004631 int i;
4632 PyObject *m;
4633 char *name;
4634 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004635 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004636 &combinations_type,
4637 &cwr_type,
4638 &cycle_type,
4639 &dropwhile_type,
4640 &takewhile_type,
4641 &islice_type,
4642 &starmap_type,
4643 &chain_type,
4644 &compress_type,
4645 &filterfalse_type,
4646 &count_type,
4647 &ziplongest_type,
4648 &permutations_type,
4649 &product_type,
4650 &repeat_type,
4651 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004652 &_grouper_type,
4653 &tee_type,
4654 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004655 NULL
4656 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004657
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004658 Py_TYPE(&teedataobject_type) = &PyType_Type;
4659 m = PyModule_Create(&itertoolsmodule);
4660 if (m == NULL)
4661 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004662
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004663 for (i=0 ; typelist[i] != NULL ; i++) {
4664 if (PyType_Ready(typelist[i]) < 0)
4665 return NULL;
4666 name = strchr(typelist[i]->tp_name, '.');
4667 assert (name != NULL);
4668 Py_INCREF(typelist[i]);
4669 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4670 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004671
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004672 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004673}