blob: 6bf04cbee31c35ce76e1b05887472d6ca716e728 [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 {
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700104 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 {
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700296 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
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000874 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
875 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 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700945 if (Py_SIZE(lz->saved) == 0)
946 return NULL;
947 item = PyList_GET_ITEM(lz->saved, lz->index);
948 lz->index++;
949 if (lz->index >= Py_SIZE(lz->saved))
950 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
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001075 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
1076 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
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001133 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1134 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
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001243 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1244 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
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001299 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1300 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
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001411 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1412 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) {
1420 stop = PyLong_AsSsize_t(a1);
1421 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)
1432 start = PyLong_AsSsize_t(a1);
1433 if (start == -1 && PyErr_Occurred())
1434 PyErr_Clear();
1435 if (a2 != Py_None) {
1436 stop = PyLong_AsSsize_t(a2);
1437 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)
1456 step = PyLong_AsSsize_t(a3);
1457 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
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001665 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1666 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
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001823 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1824 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
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001867 if (lz->source == NULL)
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001868 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001869
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001870 if (lz->active == NULL) {
1871 PyObject *iterable = PyIter_Next(lz->source);
1872 if (iterable == NULL) {
1873 Py_CLEAR(lz->source);
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001874 return NULL; /* no more input sources */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001875 }
1876 lz->active = PyObject_GetIter(iterable);
1877 Py_DECREF(iterable);
1878 if (lz->active == NULL) {
1879 Py_CLEAR(lz->source);
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001880 return NULL; /* input not iterable */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001881 }
1882 }
Raymond Hettingerb5244a32015-08-16 14:24:20 -07001883 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001884 if (item != NULL)
1885 return item;
1886 if (PyErr_Occurred()) {
1887 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1888 PyErr_Clear();
1889 else
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001890 return NULL; /* input raised an exception */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001891 }
1892 Py_CLEAR(lz->active);
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001893 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001894}
1895
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001896static PyObject *
1897chain_reduce(chainobject *lz)
1898{
1899 if (lz->source) {
1900 /* we can't pickle function objects (itertools.from_iterable) so
1901 * we must use setstate to replace the iterable. One day we
1902 * will fix pickling of functions
1903 */
1904 if (lz->active) {
1905 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1906 } else {
1907 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1908 }
1909 } else {
1910 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1911 }
1912 return NULL;
1913}
1914
1915static PyObject *
1916chain_setstate(chainobject *lz, PyObject *state)
1917{
1918 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001919
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001920 if (!PyTuple_Check(state)) {
1921 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001922 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001923 }
1924 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
1925 return NULL;
1926 }
1927 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
1928 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
1929 return NULL;
1930 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001931
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02001932 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001933 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02001934 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001935 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001936 Py_RETURN_NONE;
1937}
1938
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001939PyDoc_STRVAR(chain_doc,
1940"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001941\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001942Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001943first iterable until it is exhausted, then elements from the next\n\
1944iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001945
Christian Heimesf16baeb2008-02-29 14:57:44 +00001946PyDoc_STRVAR(chain_from_iterable_doc,
1947"chain.from_iterable(iterable) --> chain object\n\
1948\n\
Martin Pantereb995702016-07-28 01:11:04 +00001949Alternate chain() constructor taking a single iterable argument\n\
Christian Heimesf16baeb2008-02-29 14:57:44 +00001950that evaluates lazily.");
1951
1952static PyMethodDef chain_methods[] = {
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001953 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1954 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001955 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1956 reduce_doc},
1957 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1958 setstate_doc},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001959 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001960};
1961
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001962static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001963 PyVarObject_HEAD_INIT(NULL, 0)
1964 "itertools.chain", /* tp_name */
1965 sizeof(chainobject), /* tp_basicsize */
1966 0, /* tp_itemsize */
1967 /* methods */
1968 (destructor)chain_dealloc, /* tp_dealloc */
1969 0, /* tp_print */
1970 0, /* tp_getattr */
1971 0, /* tp_setattr */
1972 0, /* tp_reserved */
1973 0, /* tp_repr */
1974 0, /* tp_as_number */
1975 0, /* tp_as_sequence */
1976 0, /* tp_as_mapping */
1977 0, /* tp_hash */
1978 0, /* tp_call */
1979 0, /* tp_str */
1980 PyObject_GenericGetAttr, /* tp_getattro */
1981 0, /* tp_setattro */
1982 0, /* tp_as_buffer */
1983 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1984 Py_TPFLAGS_BASETYPE, /* tp_flags */
1985 chain_doc, /* tp_doc */
1986 (traverseproc)chain_traverse, /* tp_traverse */
1987 0, /* tp_clear */
1988 0, /* tp_richcompare */
1989 0, /* tp_weaklistoffset */
1990 PyObject_SelfIter, /* tp_iter */
1991 (iternextfunc)chain_next, /* tp_iternext */
1992 chain_methods, /* tp_methods */
1993 0, /* tp_members */
1994 0, /* tp_getset */
1995 0, /* tp_base */
1996 0, /* tp_dict */
1997 0, /* tp_descr_get */
1998 0, /* tp_descr_set */
1999 0, /* tp_dictoffset */
2000 0, /* tp_init */
2001 0, /* tp_alloc */
2002 chain_new, /* tp_new */
2003 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002004};
2005
2006
Christian Heimesc3f30c42008-02-22 16:37:40 +00002007/* product object ************************************************************/
2008
2009typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002010 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002011 PyObject *pools; /* tuple of pool tuples */
2012 Py_ssize_t *indices; /* one index per pool */
2013 PyObject *result; /* most recently returned result tuple */
2014 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002015} productobject;
2016
2017static PyTypeObject product_type;
2018
2019static PyObject *
2020product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2021{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002022 productobject *lz;
2023 Py_ssize_t nargs, npools, repeat=1;
2024 PyObject *pools = NULL;
2025 Py_ssize_t *indices = NULL;
2026 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002027
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002028 if (kwds != NULL) {
2029 char *kwlist[] = {"repeat", 0};
2030 PyObject *tmpargs = PyTuple_New(0);
2031 if (tmpargs == NULL)
2032 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002033 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2034 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002035 Py_DECREF(tmpargs);
2036 return NULL;
2037 }
2038 Py_DECREF(tmpargs);
2039 if (repeat < 0) {
2040 PyErr_SetString(PyExc_ValueError,
2041 "repeat argument cannot be negative");
2042 return NULL;
2043 }
2044 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002045
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002046 assert(PyTuple_CheckExact(args));
2047 if (repeat == 0) {
2048 nargs = 0;
2049 } else {
2050 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002051 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002052 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2053 return NULL;
2054 }
2055 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002056 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002057
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002058 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002059 if (indices == NULL) {
2060 PyErr_NoMemory();
2061 goto error;
2062 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002063
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002064 pools = PyTuple_New(npools);
2065 if (pools == NULL)
2066 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002067
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002068 for (i=0; i < nargs ; ++i) {
2069 PyObject *item = PyTuple_GET_ITEM(args, i);
2070 PyObject *pool = PySequence_Tuple(item);
2071 if (pool == NULL)
2072 goto error;
2073 PyTuple_SET_ITEM(pools, i, pool);
2074 indices[i] = 0;
2075 }
2076 for ( ; i < npools; ++i) {
2077 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2078 Py_INCREF(pool);
2079 PyTuple_SET_ITEM(pools, i, pool);
2080 indices[i] = 0;
2081 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002082
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002083 /* create productobject structure */
2084 lz = (productobject *)type->tp_alloc(type, 0);
2085 if (lz == NULL)
2086 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002087
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002088 lz->pools = pools;
2089 lz->indices = indices;
2090 lz->result = NULL;
2091 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002092
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002093 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002094
2095error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002096 if (indices != NULL)
2097 PyMem_Free(indices);
2098 Py_XDECREF(pools);
2099 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002100}
2101
2102static void
2103product_dealloc(productobject *lz)
2104{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002105 PyObject_GC_UnTrack(lz);
2106 Py_XDECREF(lz->pools);
2107 Py_XDECREF(lz->result);
2108 if (lz->indices != NULL)
2109 PyMem_Free(lz->indices);
2110 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002111}
2112
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002113static PyObject *
2114product_sizeof(productobject *lz, void *unused)
2115{
2116 Py_ssize_t res;
2117
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002118 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002119 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2120 return PyLong_FromSsize_t(res);
2121}
2122
2123PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2124
Christian Heimesc3f30c42008-02-22 16:37:40 +00002125static int
2126product_traverse(productobject *lz, visitproc visit, void *arg)
2127{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002128 Py_VISIT(lz->pools);
2129 Py_VISIT(lz->result);
2130 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002131}
2132
2133static PyObject *
2134product_next(productobject *lz)
2135{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002136 PyObject *pool;
2137 PyObject *elem;
2138 PyObject *oldelem;
2139 PyObject *pools = lz->pools;
2140 PyObject *result = lz->result;
2141 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2142 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002143
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002144 if (lz->stopped)
2145 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002146
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002147 if (result == NULL) {
2148 /* On the first pass, return an initial tuple filled with the
2149 first element from each pool. */
2150 result = PyTuple_New(npools);
2151 if (result == NULL)
2152 goto empty;
2153 lz->result = result;
2154 for (i=0; i < npools; i++) {
2155 pool = PyTuple_GET_ITEM(pools, i);
2156 if (PyTuple_GET_SIZE(pool) == 0)
2157 goto empty;
2158 elem = PyTuple_GET_ITEM(pool, 0);
2159 Py_INCREF(elem);
2160 PyTuple_SET_ITEM(result, i, elem);
2161 }
2162 } else {
2163 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002164
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002165 /* Copy the previous result tuple or re-use it if available */
2166 if (Py_REFCNT(result) > 1) {
2167 PyObject *old_result = result;
2168 result = PyTuple_New(npools);
2169 if (result == NULL)
2170 goto empty;
2171 lz->result = result;
2172 for (i=0; i < npools; i++) {
2173 elem = PyTuple_GET_ITEM(old_result, i);
2174 Py_INCREF(elem);
2175 PyTuple_SET_ITEM(result, i, elem);
2176 }
2177 Py_DECREF(old_result);
2178 }
2179 /* Now, we've got the only copy so we can update it in-place */
2180 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002181
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002182 /* Update the pool indices right-to-left. Only advance to the
2183 next pool when the previous one rolls-over */
2184 for (i=npools-1 ; i >= 0 ; i--) {
2185 pool = PyTuple_GET_ITEM(pools, i);
2186 indices[i]++;
2187 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2188 /* Roll-over and advance to next pool */
2189 indices[i] = 0;
2190 elem = PyTuple_GET_ITEM(pool, 0);
2191 Py_INCREF(elem);
2192 oldelem = PyTuple_GET_ITEM(result, i);
2193 PyTuple_SET_ITEM(result, i, elem);
2194 Py_DECREF(oldelem);
2195 } else {
2196 /* No rollover. Just increment and stop here. */
2197 elem = PyTuple_GET_ITEM(pool, indices[i]);
2198 Py_INCREF(elem);
2199 oldelem = PyTuple_GET_ITEM(result, i);
2200 PyTuple_SET_ITEM(result, i, elem);
2201 Py_DECREF(oldelem);
2202 break;
2203 }
2204 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002205
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002206 /* If i is negative, then the indices have all rolled-over
2207 and we're done. */
2208 if (i < 0)
2209 goto empty;
2210 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002211
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002212 Py_INCREF(result);
2213 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002214
2215empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002216 lz->stopped = 1;
2217 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002218}
2219
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002220static PyObject *
2221product_reduce(productobject *lz)
2222{
2223 if (lz->stopped) {
2224 return Py_BuildValue("O(())", Py_TYPE(lz));
2225 } else if (lz->result == NULL) {
2226 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2227 } else {
2228 PyObject *indices;
2229 Py_ssize_t n, i;
2230
2231 /* we must pickle the indices use them for setstate, and
2232 * additionally indicate that the iterator has started
2233 */
2234 n = PyTuple_GET_SIZE(lz->pools);
2235 indices = PyTuple_New(n);
2236 if (indices == NULL)
2237 return NULL;
2238 for (i=0; i<n; i++){
2239 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2240 if (!index) {
2241 Py_DECREF(indices);
2242 return NULL;
2243 }
2244 PyTuple_SET_ITEM(indices, i, index);
2245 }
2246 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2247 }
2248}
2249
2250static PyObject *
2251product_setstate(productobject *lz, PyObject *state)
2252{
2253 PyObject *result;
2254 Py_ssize_t n, i;
2255
2256 n = PyTuple_GET_SIZE(lz->pools);
2257 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2258 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2259 return NULL;
2260 }
2261 for (i=0; i<n; i++)
2262 {
2263 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2264 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002265 PyObject* pool;
2266 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002267 if (index < 0 && PyErr_Occurred())
2268 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002269 pool = PyTuple_GET_ITEM(lz->pools, i);
2270 poolsize = PyTuple_GET_SIZE(pool);
2271 if (poolsize == 0) {
2272 lz->stopped = 1;
2273 Py_RETURN_NONE;
2274 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002275 /* clamp the index */
2276 if (index < 0)
2277 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002278 else if (index > poolsize-1)
2279 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002280 lz->indices[i] = index;
2281 }
2282
2283 result = PyTuple_New(n);
2284 if (!result)
2285 return NULL;
2286 for (i=0; i<n; i++) {
2287 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2288 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2289 Py_INCREF(element);
2290 PyTuple_SET_ITEM(result, i, element);
2291 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002292 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002293 Py_RETURN_NONE;
2294}
2295
2296static PyMethodDef product_methods[] = {
2297 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2298 reduce_doc},
2299 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2300 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002301 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2302 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002303 {NULL, NULL} /* sentinel */
2304};
2305
Christian Heimesc3f30c42008-02-22 16:37:40 +00002306PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002307"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002308\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002309Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002310For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2311The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2312cycle in a manner similar to an odometer (with the rightmost element changing\n\
2313on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002314To compute the product of an iterable with itself, specify the number\n\
2315of repetitions with the optional repeat keyword argument. For example,\n\
2316product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002317product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2318product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2319
2320static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002321 PyVarObject_HEAD_INIT(NULL, 0)
2322 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002323 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002324 0, /* tp_itemsize */
2325 /* methods */
2326 (destructor)product_dealloc, /* tp_dealloc */
2327 0, /* tp_print */
2328 0, /* tp_getattr */
2329 0, /* tp_setattr */
2330 0, /* tp_reserved */
2331 0, /* tp_repr */
2332 0, /* tp_as_number */
2333 0, /* tp_as_sequence */
2334 0, /* tp_as_mapping */
2335 0, /* tp_hash */
2336 0, /* tp_call */
2337 0, /* tp_str */
2338 PyObject_GenericGetAttr, /* tp_getattro */
2339 0, /* tp_setattro */
2340 0, /* tp_as_buffer */
2341 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2342 Py_TPFLAGS_BASETYPE, /* tp_flags */
2343 product_doc, /* tp_doc */
2344 (traverseproc)product_traverse, /* tp_traverse */
2345 0, /* tp_clear */
2346 0, /* tp_richcompare */
2347 0, /* tp_weaklistoffset */
2348 PyObject_SelfIter, /* tp_iter */
2349 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002350 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002351 0, /* tp_members */
2352 0, /* tp_getset */
2353 0, /* tp_base */
2354 0, /* tp_dict */
2355 0, /* tp_descr_get */
2356 0, /* tp_descr_set */
2357 0, /* tp_dictoffset */
2358 0, /* tp_init */
2359 0, /* tp_alloc */
2360 product_new, /* tp_new */
2361 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002362};
2363
2364
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002365/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002366
2367typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002368 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002369 PyObject *pool; /* input converted to a tuple */
2370 Py_ssize_t *indices; /* one index per result element */
2371 PyObject *result; /* most recently returned result tuple */
2372 Py_ssize_t r; /* size of result tuple */
2373 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002374} combinationsobject;
2375
2376static PyTypeObject combinations_type;
2377
2378static PyObject *
2379combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2380{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002381 combinationsobject *co;
2382 Py_ssize_t n;
2383 Py_ssize_t r;
2384 PyObject *pool = NULL;
2385 PyObject *iterable = NULL;
2386 Py_ssize_t *indices = NULL;
2387 Py_ssize_t i;
2388 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002389
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002390 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2391 &iterable, &r))
2392 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002393
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002394 pool = PySequence_Tuple(iterable);
2395 if (pool == NULL)
2396 goto error;
2397 n = PyTuple_GET_SIZE(pool);
2398 if (r < 0) {
2399 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2400 goto error;
2401 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002402
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002403 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002404 if (indices == NULL) {
2405 PyErr_NoMemory();
2406 goto error;
2407 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002408
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002409 for (i=0 ; i<r ; i++)
2410 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002411
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002412 /* create combinationsobject structure */
2413 co = (combinationsobject *)type->tp_alloc(type, 0);
2414 if (co == NULL)
2415 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002416
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002417 co->pool = pool;
2418 co->indices = indices;
2419 co->result = NULL;
2420 co->r = r;
2421 co->stopped = r > n ? 1 : 0;
2422
2423 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002424
2425error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002426 if (indices != NULL)
2427 PyMem_Free(indices);
2428 Py_XDECREF(pool);
2429 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002430}
2431
2432static void
2433combinations_dealloc(combinationsobject *co)
2434{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002435 PyObject_GC_UnTrack(co);
2436 Py_XDECREF(co->pool);
2437 Py_XDECREF(co->result);
2438 if (co->indices != NULL)
2439 PyMem_Free(co->indices);
2440 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002441}
2442
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002443static PyObject *
2444combinations_sizeof(combinationsobject *co, void *unused)
2445{
2446 Py_ssize_t res;
2447
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002448 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002449 res += co->r * sizeof(Py_ssize_t);
2450 return PyLong_FromSsize_t(res);
2451}
2452
Christian Heimes380f7f22008-02-28 11:19:05 +00002453static int
2454combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2455{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002456 Py_VISIT(co->pool);
2457 Py_VISIT(co->result);
2458 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002459}
2460
2461static PyObject *
2462combinations_next(combinationsobject *co)
2463{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002464 PyObject *elem;
2465 PyObject *oldelem;
2466 PyObject *pool = co->pool;
2467 Py_ssize_t *indices = co->indices;
2468 PyObject *result = co->result;
2469 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2470 Py_ssize_t r = co->r;
2471 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002472
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002473 if (co->stopped)
2474 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002475
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002476 if (result == NULL) {
2477 /* On the first pass, initialize result tuple using the indices */
2478 result = PyTuple_New(r);
2479 if (result == NULL)
2480 goto empty;
2481 co->result = result;
2482 for (i=0; i<r ; i++) {
2483 index = indices[i];
2484 elem = PyTuple_GET_ITEM(pool, index);
2485 Py_INCREF(elem);
2486 PyTuple_SET_ITEM(result, i, elem);
2487 }
2488 } else {
2489 /* Copy the previous result tuple or re-use it if available */
2490 if (Py_REFCNT(result) > 1) {
2491 PyObject *old_result = result;
2492 result = PyTuple_New(r);
2493 if (result == NULL)
2494 goto empty;
2495 co->result = result;
2496 for (i=0; i<r ; i++) {
2497 elem = PyTuple_GET_ITEM(old_result, i);
2498 Py_INCREF(elem);
2499 PyTuple_SET_ITEM(result, i, elem);
2500 }
2501 Py_DECREF(old_result);
2502 }
2503 /* Now, we've got the only copy so we can update it in-place
2504 * CPython's empty tuple is a singleton and cached in
2505 * PyTuple's freelist.
2506 */
2507 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002508
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002509 /* Scan indices right-to-left until finding one that is not
2510 at its maximum (i + n - r). */
2511 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2512 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002513
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002514 /* If i is negative, then the indices are all at
2515 their maximum value and we're done. */
2516 if (i < 0)
2517 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002518
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002519 /* Increment the current index which we know is not at its
2520 maximum. Then move back to the right setting each index
2521 to its lowest possible value (one higher than the index
2522 to its left -- this maintains the sort order invariant). */
2523 indices[i]++;
2524 for (j=i+1 ; j<r ; j++)
2525 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002526
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002527 /* Update the result tuple for the new indices
2528 starting with i, the leftmost index that changed */
2529 for ( ; i<r ; i++) {
2530 index = indices[i];
2531 elem = PyTuple_GET_ITEM(pool, index);
2532 Py_INCREF(elem);
2533 oldelem = PyTuple_GET_ITEM(result, i);
2534 PyTuple_SET_ITEM(result, i, elem);
2535 Py_DECREF(oldelem);
2536 }
2537 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002538
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002539 Py_INCREF(result);
2540 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002541
2542empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002543 co->stopped = 1;
2544 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002545}
2546
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002547static PyObject *
2548combinations_reduce(combinationsobject *lz)
2549{
2550 if (lz->result == NULL) {
2551 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2552 } else if (lz->stopped) {
2553 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2554 } else {
2555 PyObject *indices;
2556 Py_ssize_t i;
2557
2558 /* we must pickle the indices and use them for setstate */
2559 indices = PyTuple_New(lz->r);
2560 if (!indices)
2561 return NULL;
2562 for (i=0; i<lz->r; i++)
2563 {
2564 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2565 if (!index) {
2566 Py_DECREF(indices);
2567 return NULL;
2568 }
2569 PyTuple_SET_ITEM(indices, i, index);
2570 }
2571
2572 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2573 }
2574}
2575
2576static PyObject *
2577combinations_setstate(combinationsobject *lz, PyObject *state)
2578{
2579 PyObject *result;
2580 Py_ssize_t i;
2581 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2582
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002583 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002584 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2585 return NULL;
2586 }
2587
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002588 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002589 Py_ssize_t max;
2590 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2591 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002592
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002593 if (index == -1 && PyErr_Occurred())
2594 return NULL; /* not an integer */
2595 max = i + n - lz->r;
2596 /* clamp the index (beware of negative max) */
2597 if (index > max)
2598 index = max;
2599 if (index < 0)
2600 index = 0;
2601 lz->indices[i] = index;
2602 }
2603
2604 result = PyTuple_New(lz->r);
2605 if (result == NULL)
2606 return NULL;
2607 for (i=0; i<lz->r; i++) {
2608 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2609 Py_INCREF(element);
2610 PyTuple_SET_ITEM(result, i, element);
2611 }
2612
Serhiy Storchaka48842712016-04-06 09:45:48 +03002613 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002614 Py_RETURN_NONE;
2615}
2616
2617static PyMethodDef combinations_methods[] = {
2618 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2619 reduce_doc},
2620 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2621 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002622 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2623 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002624 {NULL, NULL} /* sentinel */
2625};
2626
Christian Heimes380f7f22008-02-28 11:19:05 +00002627PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002628"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002629\n\
2630Return successive r-length combinations of elements in the iterable.\n\n\
2631combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2632
2633static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002634 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002635 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002636 sizeof(combinationsobject), /* tp_basicsize */
2637 0, /* tp_itemsize */
2638 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002639 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002640 0, /* tp_print */
2641 0, /* tp_getattr */
2642 0, /* tp_setattr */
2643 0, /* tp_reserved */
2644 0, /* tp_repr */
2645 0, /* tp_as_number */
2646 0, /* tp_as_sequence */
2647 0, /* tp_as_mapping */
2648 0, /* tp_hash */
2649 0, /* tp_call */
2650 0, /* tp_str */
2651 PyObject_GenericGetAttr, /* tp_getattro */
2652 0, /* tp_setattro */
2653 0, /* tp_as_buffer */
2654 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2655 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002656 combinations_doc, /* tp_doc */
2657 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002658 0, /* tp_clear */
2659 0, /* tp_richcompare */
2660 0, /* tp_weaklistoffset */
2661 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002662 (iternextfunc)combinations_next, /* tp_iternext */
2663 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002664 0, /* tp_members */
2665 0, /* tp_getset */
2666 0, /* tp_base */
2667 0, /* tp_dict */
2668 0, /* tp_descr_get */
2669 0, /* tp_descr_set */
2670 0, /* tp_dictoffset */
2671 0, /* tp_init */
2672 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002673 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002674 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002675};
2676
2677
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002678/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002679
2680/* Equivalent to:
2681
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002682 def combinations_with_replacement(iterable, r):
2683 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2684 # number items returned: (n+r-1)! / r! / (n-1)!
2685 pool = tuple(iterable)
2686 n = len(pool)
2687 indices = [0] * r
2688 yield tuple(pool[i] for i in indices)
2689 while 1:
2690 for i in reversed(range(r)):
2691 if indices[i] != n - 1:
2692 break
2693 else:
2694 return
2695 indices[i:] = [indices[i] + 1] * (r - i)
2696 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002697
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002698 def combinations_with_replacement2(iterable, r):
2699 'Alternate version that filters from product()'
2700 pool = tuple(iterable)
2701 n = len(pool)
2702 for indices in product(range(n), repeat=r):
2703 if sorted(indices) == list(indices):
2704 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002705*/
2706typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002707 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002708 PyObject *pool; /* input converted to a tuple */
2709 Py_ssize_t *indices; /* one index per result element */
2710 PyObject *result; /* most recently returned result tuple */
2711 Py_ssize_t r; /* size of result tuple */
2712 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002713} cwrobject;
2714
2715static PyTypeObject cwr_type;
2716
2717static PyObject *
2718cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2719{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002720 cwrobject *co;
2721 Py_ssize_t n;
2722 Py_ssize_t r;
2723 PyObject *pool = NULL;
2724 PyObject *iterable = NULL;
2725 Py_ssize_t *indices = NULL;
2726 Py_ssize_t i;
2727 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002728
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002729 if (!PyArg_ParseTupleAndKeywords(args, kwds,
2730 "On:combinations_with_replacement",
2731 kwargs, &iterable, &r))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002732 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002733
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002734 pool = PySequence_Tuple(iterable);
2735 if (pool == NULL)
2736 goto error;
2737 n = PyTuple_GET_SIZE(pool);
2738 if (r < 0) {
2739 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2740 goto error;
2741 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002742
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002743 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002744 if (indices == NULL) {
2745 PyErr_NoMemory();
2746 goto error;
2747 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002748
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002749 for (i=0 ; i<r ; i++)
2750 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002751
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002752 /* create cwrobject structure */
2753 co = (cwrobject *)type->tp_alloc(type, 0);
2754 if (co == NULL)
2755 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002756
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002757 co->pool = pool;
2758 co->indices = indices;
2759 co->result = NULL;
2760 co->r = r;
2761 co->stopped = !n && r;
2762
2763 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002764
2765error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002766 if (indices != NULL)
2767 PyMem_Free(indices);
2768 Py_XDECREF(pool);
2769 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002770}
2771
2772static void
2773cwr_dealloc(cwrobject *co)
2774{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002775 PyObject_GC_UnTrack(co);
2776 Py_XDECREF(co->pool);
2777 Py_XDECREF(co->result);
2778 if (co->indices != NULL)
2779 PyMem_Free(co->indices);
2780 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002781}
2782
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002783static PyObject *
2784cwr_sizeof(cwrobject *co, void *unused)
2785{
2786 Py_ssize_t res;
2787
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002788 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002789 res += co->r * sizeof(Py_ssize_t);
2790 return PyLong_FromSsize_t(res);
2791}
2792
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002793static int
2794cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2795{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002796 Py_VISIT(co->pool);
2797 Py_VISIT(co->result);
2798 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002799}
2800
2801static PyObject *
2802cwr_next(cwrobject *co)
2803{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002804 PyObject *elem;
2805 PyObject *oldelem;
2806 PyObject *pool = co->pool;
2807 Py_ssize_t *indices = co->indices;
2808 PyObject *result = co->result;
2809 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2810 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002811 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002812
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002813 if (co->stopped)
2814 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002815
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002816 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002817 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002818 result = PyTuple_New(r);
2819 if (result == NULL)
2820 goto empty;
2821 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002822 if (n > 0) {
2823 elem = PyTuple_GET_ITEM(pool, 0);
2824 for (i=0; i<r ; i++) {
2825 assert(indices[i] == 0);
2826 Py_INCREF(elem);
2827 PyTuple_SET_ITEM(result, i, elem);
2828 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002829 }
2830 } else {
2831 /* Copy the previous result tuple or re-use it if available */
2832 if (Py_REFCNT(result) > 1) {
2833 PyObject *old_result = result;
2834 result = PyTuple_New(r);
2835 if (result == NULL)
2836 goto empty;
2837 co->result = result;
2838 for (i=0; i<r ; i++) {
2839 elem = PyTuple_GET_ITEM(old_result, i);
2840 Py_INCREF(elem);
2841 PyTuple_SET_ITEM(result, i, elem);
2842 }
2843 Py_DECREF(old_result);
2844 }
2845 /* Now, we've got the only copy so we can update it in-place CPython's
2846 empty tuple is a singleton and cached in PyTuple's freelist. */
2847 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002848
Tim Peters9edb1682013-09-03 11:49:31 -05002849 /* Scan indices right-to-left until finding one that is not
2850 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002851 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2852 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002853
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002854 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002855 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002856 if (i < 0)
2857 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002858
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002859 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002860 maximum. Then set all to the right to the same value. */
2861 index = indices[i] + 1;
2862 assert(index < n);
2863 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002864 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002865 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002866 Py_INCREF(elem);
2867 oldelem = PyTuple_GET_ITEM(result, i);
2868 PyTuple_SET_ITEM(result, i, elem);
2869 Py_DECREF(oldelem);
2870 }
2871 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002872
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002873 Py_INCREF(result);
2874 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002875
2876empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002877 co->stopped = 1;
2878 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002879}
2880
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002881static PyObject *
2882cwr_reduce(cwrobject *lz)
2883{
2884 if (lz->result == NULL) {
2885 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2886 } else if (lz->stopped) {
2887 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2888 } else {
2889 PyObject *indices;
2890 Py_ssize_t i;
2891
2892 /* we must pickle the indices and use them for setstate */
2893 indices = PyTuple_New(lz->r);
2894 if (!indices)
2895 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002896 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002897 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2898 if (!index) {
2899 Py_DECREF(indices);
2900 return NULL;
2901 }
2902 PyTuple_SET_ITEM(indices, i, index);
2903 }
2904
2905 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2906 }
2907}
2908
2909static PyObject *
2910cwr_setstate(cwrobject *lz, PyObject *state)
2911{
2912 PyObject *result;
2913 Py_ssize_t n, i;
2914
2915 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2916 {
2917 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2918 return NULL;
2919 }
2920
2921 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002922 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002923 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2924 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002925
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002926 if (index < 0 && PyErr_Occurred())
2927 return NULL; /* not an integer */
2928 /* clamp the index */
2929 if (index < 0)
2930 index = 0;
2931 else if (index > n-1)
2932 index = n-1;
2933 lz->indices[i] = index;
2934 }
2935 result = PyTuple_New(lz->r);
2936 if (result == NULL)
2937 return NULL;
2938 for (i=0; i<lz->r; i++) {
2939 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2940 Py_INCREF(element);
2941 PyTuple_SET_ITEM(result, i, element);
2942 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002943 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002944 Py_RETURN_NONE;
2945}
2946
2947static PyMethodDef cwr_methods[] = {
2948 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2949 reduce_doc},
2950 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2951 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002952 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
2953 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002954 {NULL, NULL} /* sentinel */
2955};
2956
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002957PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002958"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002959\n\
2960Return successive r-length combinations of elements in the iterable\n\
2961allowing individual elements to have successive repeats.\n\
2962combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2963
2964static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002965 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002966 "itertools.combinations_with_replacement", /* tp_name */
2967 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002968 0, /* tp_itemsize */
2969 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002970 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002971 0, /* tp_print */
2972 0, /* tp_getattr */
2973 0, /* tp_setattr */
2974 0, /* tp_reserved */
2975 0, /* tp_repr */
2976 0, /* tp_as_number */
2977 0, /* tp_as_sequence */
2978 0, /* tp_as_mapping */
2979 0, /* tp_hash */
2980 0, /* tp_call */
2981 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002982 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002983 0, /* tp_setattro */
2984 0, /* tp_as_buffer */
2985 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002986 Py_TPFLAGS_BASETYPE, /* tp_flags */
2987 cwr_doc, /* tp_doc */
2988 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002989 0, /* tp_clear */
2990 0, /* tp_richcompare */
2991 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002992 PyObject_SelfIter, /* tp_iter */
2993 (iternextfunc)cwr_next, /* tp_iternext */
2994 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002995 0, /* tp_members */
2996 0, /* tp_getset */
2997 0, /* tp_base */
2998 0, /* tp_dict */
2999 0, /* tp_descr_get */
3000 0, /* tp_descr_set */
3001 0, /* tp_dictoffset */
3002 0, /* tp_init */
3003 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003004 cwr_new, /* tp_new */
3005 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003006};
3007
3008
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003009/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003010
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003011def permutations(iterable, r=None):
3012 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
3013 pool = tuple(iterable)
3014 n = len(pool)
3015 r = n if r is None else r
3016 indices = range(n)
3017 cycles = range(n-r+1, n+1)[::-1]
3018 yield tuple(pool[i] for i in indices[:r])
3019 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003020 for i in reversed(range(r)):
3021 cycles[i] -= 1
3022 if cycles[i] == 0:
3023 indices[i:] = indices[i+1:] + indices[i:i+1]
3024 cycles[i] = n - i
3025 else:
3026 j = cycles[i]
3027 indices[i], indices[-j] = indices[-j], indices[i]
3028 yield tuple(pool[i] for i in indices[:r])
3029 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003030 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003031 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003032*/
3033
3034typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003035 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003036 PyObject *pool; /* input converted to a tuple */
3037 Py_ssize_t *indices; /* one index per element in the pool */
3038 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3039 PyObject *result; /* most recently returned result tuple */
3040 Py_ssize_t r; /* size of result tuple */
3041 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003042} permutationsobject;
3043
3044static PyTypeObject permutations_type;
3045
3046static PyObject *
3047permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3048{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003049 permutationsobject *po;
3050 Py_ssize_t n;
3051 Py_ssize_t r;
3052 PyObject *robj = Py_None;
3053 PyObject *pool = NULL;
3054 PyObject *iterable = NULL;
3055 Py_ssize_t *indices = NULL;
3056 Py_ssize_t *cycles = NULL;
3057 Py_ssize_t i;
3058 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003059
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003060 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
3061 &iterable, &robj))
3062 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003063
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003064 pool = PySequence_Tuple(iterable);
3065 if (pool == NULL)
3066 goto error;
3067 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003068
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003069 r = n;
3070 if (robj != Py_None) {
3071 if (!PyLong_Check(robj)) {
3072 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3073 goto error;
3074 }
3075 r = PyLong_AsSsize_t(robj);
3076 if (r == -1 && PyErr_Occurred())
3077 goto error;
3078 }
3079 if (r < 0) {
3080 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3081 goto error;
3082 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003083
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003084 indices = PyMem_New(Py_ssize_t, n);
3085 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003086 if (indices == NULL || cycles == NULL) {
3087 PyErr_NoMemory();
3088 goto error;
3089 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003090
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003091 for (i=0 ; i<n ; i++)
3092 indices[i] = i;
3093 for (i=0 ; i<r ; i++)
3094 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003095
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003096 /* create permutationsobject structure */
3097 po = (permutationsobject *)type->tp_alloc(type, 0);
3098 if (po == NULL)
3099 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003100
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003101 po->pool = pool;
3102 po->indices = indices;
3103 po->cycles = cycles;
3104 po->result = NULL;
3105 po->r = r;
3106 po->stopped = r > n ? 1 : 0;
3107
3108 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003109
3110error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003111 if (indices != NULL)
3112 PyMem_Free(indices);
3113 if (cycles != NULL)
3114 PyMem_Free(cycles);
3115 Py_XDECREF(pool);
3116 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003117}
3118
3119static void
3120permutations_dealloc(permutationsobject *po)
3121{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003122 PyObject_GC_UnTrack(po);
3123 Py_XDECREF(po->pool);
3124 Py_XDECREF(po->result);
3125 PyMem_Free(po->indices);
3126 PyMem_Free(po->cycles);
3127 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003128}
3129
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003130static PyObject *
3131permutations_sizeof(permutationsobject *po, void *unused)
3132{
3133 Py_ssize_t res;
3134
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003135 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003136 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3137 res += po->r * sizeof(Py_ssize_t);
3138 return PyLong_FromSsize_t(res);
3139}
3140
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003141static int
3142permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3143{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003144 Py_VISIT(po->pool);
3145 Py_VISIT(po->result);
3146 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003147}
3148
3149static PyObject *
3150permutations_next(permutationsobject *po)
3151{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003152 PyObject *elem;
3153 PyObject *oldelem;
3154 PyObject *pool = po->pool;
3155 Py_ssize_t *indices = po->indices;
3156 Py_ssize_t *cycles = po->cycles;
3157 PyObject *result = po->result;
3158 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3159 Py_ssize_t r = po->r;
3160 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003161
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003162 if (po->stopped)
3163 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003164
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003165 if (result == NULL) {
3166 /* On the first pass, initialize result tuple using the indices */
3167 result = PyTuple_New(r);
3168 if (result == NULL)
3169 goto empty;
3170 po->result = result;
3171 for (i=0; i<r ; i++) {
3172 index = indices[i];
3173 elem = PyTuple_GET_ITEM(pool, index);
3174 Py_INCREF(elem);
3175 PyTuple_SET_ITEM(result, i, elem);
3176 }
3177 } else {
3178 if (n == 0)
3179 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003180
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003181 /* Copy the previous result tuple or re-use it if available */
3182 if (Py_REFCNT(result) > 1) {
3183 PyObject *old_result = result;
3184 result = PyTuple_New(r);
3185 if (result == NULL)
3186 goto empty;
3187 po->result = result;
3188 for (i=0; i<r ; i++) {
3189 elem = PyTuple_GET_ITEM(old_result, i);
3190 Py_INCREF(elem);
3191 PyTuple_SET_ITEM(result, i, elem);
3192 }
3193 Py_DECREF(old_result);
3194 }
3195 /* Now, we've got the only copy so we can update it in-place */
3196 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003197
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003198 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3199 for (i=r-1 ; i>=0 ; i--) {
3200 cycles[i] -= 1;
3201 if (cycles[i] == 0) {
3202 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3203 index = indices[i];
3204 for (j=i ; j<n-1 ; j++)
3205 indices[j] = indices[j+1];
3206 indices[n-1] = index;
3207 cycles[i] = n - i;
3208 } else {
3209 j = cycles[i];
3210 index = indices[i];
3211 indices[i] = indices[n-j];
3212 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003213
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003214 for (k=i; k<r ; k++) {
3215 /* start with i, the leftmost element that changed */
3216 /* yield tuple(pool[k] for k in indices[:r]) */
3217 index = indices[k];
3218 elem = PyTuple_GET_ITEM(pool, index);
3219 Py_INCREF(elem);
3220 oldelem = PyTuple_GET_ITEM(result, k);
3221 PyTuple_SET_ITEM(result, k, elem);
3222 Py_DECREF(oldelem);
3223 }
3224 break;
3225 }
3226 }
3227 /* If i is negative, then the cycles have all
3228 rolled-over and we're done. */
3229 if (i < 0)
3230 goto empty;
3231 }
3232 Py_INCREF(result);
3233 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003234
3235empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003236 po->stopped = 1;
3237 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003238}
3239
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003240static PyObject *
3241permutations_reduce(permutationsobject *po)
3242{
3243 if (po->result == NULL) {
3244 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3245 } else if (po->stopped) {
3246 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3247 } else {
3248 PyObject *indices=NULL, *cycles=NULL;
3249 Py_ssize_t n, i;
3250
3251 /* we must pickle the indices and cycles and use them for setstate */
3252 n = PyTuple_GET_SIZE(po->pool);
3253 indices = PyTuple_New(n);
3254 if (indices == NULL)
3255 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003256 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003257 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3258 if (!index)
3259 goto err;
3260 PyTuple_SET_ITEM(indices, i, index);
3261 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003262
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003263 cycles = PyTuple_New(po->r);
3264 if (cycles == NULL)
3265 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003266 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003267 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3268 if (!index)
3269 goto err;
3270 PyTuple_SET_ITEM(cycles, i, index);
3271 }
3272 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3273 po->pool, po->r,
3274 indices, cycles);
3275 err:
3276 Py_XDECREF(indices);
3277 Py_XDECREF(cycles);
3278 return NULL;
3279 }
3280}
3281
3282static PyObject *
3283permutations_setstate(permutationsobject *po, PyObject *state)
3284{
3285 PyObject *indices, *cycles, *result;
3286 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003287
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003288 if (!PyTuple_Check(state)) {
3289 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3290 return NULL;
3291 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003292 if (!PyArg_ParseTuple(state, "O!O!",
3293 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003294 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003295 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003296 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003297
3298 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003299 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003300 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3301 return NULL;
3302 }
3303
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003304 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003305 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3306 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3307 if (index < 0 && PyErr_Occurred())
3308 return NULL; /* not an integer */
3309 /* clamp the index */
3310 if (index < 0)
3311 index = 0;
3312 else if (index > n-1)
3313 index = n-1;
3314 po->indices[i] = index;
3315 }
3316
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003317 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003318 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3319 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3320 if (index < 0 && PyErr_Occurred())
3321 return NULL; /* not an integer */
3322 if (index < 1)
3323 index = 1;
3324 else if (index > n-i)
3325 index = n-i;
3326 po->cycles[i] = index;
3327 }
3328 result = PyTuple_New(po->r);
3329 if (result == NULL)
3330 return NULL;
3331 for (i=0; i<po->r; i++) {
3332 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3333 Py_INCREF(element);
3334 PyTuple_SET_ITEM(result, i, element);
3335 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003336 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003337 Py_RETURN_NONE;
3338}
3339
3340static PyMethodDef permuations_methods[] = {
3341 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3342 reduce_doc},
3343 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3344 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003345 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3346 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003347 {NULL, NULL} /* sentinel */
3348};
3349
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003350PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003351"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003352\n\
3353Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003354permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003355
3356static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003357 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003358 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003359 sizeof(permutationsobject), /* tp_basicsize */
3360 0, /* tp_itemsize */
3361 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003362 (destructor)permutations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003363 0, /* tp_print */
3364 0, /* tp_getattr */
3365 0, /* tp_setattr */
3366 0, /* tp_reserved */
3367 0, /* tp_repr */
3368 0, /* tp_as_number */
3369 0, /* tp_as_sequence */
3370 0, /* tp_as_mapping */
3371 0, /* tp_hash */
3372 0, /* tp_call */
3373 0, /* tp_str */
3374 PyObject_GenericGetAttr, /* tp_getattro */
3375 0, /* tp_setattro */
3376 0, /* tp_as_buffer */
3377 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3378 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003379 permutations_doc, /* tp_doc */
3380 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003381 0, /* tp_clear */
3382 0, /* tp_richcompare */
3383 0, /* tp_weaklistoffset */
3384 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003385 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003386 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003387 0, /* tp_members */
3388 0, /* tp_getset */
3389 0, /* tp_base */
3390 0, /* tp_dict */
3391 0, /* tp_descr_get */
3392 0, /* tp_descr_set */
3393 0, /* tp_dictoffset */
3394 0, /* tp_init */
3395 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003396 permutations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003397 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003398};
3399
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003400/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003401
3402typedef struct {
3403 PyObject_HEAD
3404 PyObject *total;
3405 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003406 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003407} accumulateobject;
3408
3409static PyTypeObject accumulate_type;
3410
3411static PyObject *
3412accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3413{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003414 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003415 PyObject *iterable;
3416 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003417 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003418 accumulateobject *lz;
3419
Raymond Hettinger5d446132011-03-27 18:52:10 -07003420 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3421 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003422 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003423
3424 /* Get iterator. */
3425 it = PyObject_GetIter(iterable);
3426 if (it == NULL)
3427 return NULL;
3428
Raymond Hettinger482ba772010-12-01 22:48:00 +00003429 /* create accumulateobject structure */
3430 lz = (accumulateobject *)type->tp_alloc(type, 0);
3431 if (lz == NULL) {
3432 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003433 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003434 }
3435
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003436 if (binop != Py_None) {
3437 Py_XINCREF(binop);
3438 lz->binop = binop;
3439 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003440 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003441 lz->it = it;
3442 return (PyObject *)lz;
3443}
3444
3445static void
3446accumulate_dealloc(accumulateobject *lz)
3447{
3448 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003449 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003450 Py_XDECREF(lz->total);
3451 Py_XDECREF(lz->it);
3452 Py_TYPE(lz)->tp_free(lz);
3453}
3454
3455static int
3456accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3457{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003458 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003459 Py_VISIT(lz->it);
3460 Py_VISIT(lz->total);
3461 return 0;
3462}
3463
3464static PyObject *
3465accumulate_next(accumulateobject *lz)
3466{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003467 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003468
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003469 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003470 if (val == NULL)
3471 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003472
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003473 if (lz->total == NULL) {
3474 Py_INCREF(val);
3475 lz->total = val;
3476 return lz->total;
3477 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003478
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003479 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003480 newtotal = PyNumber_Add(lz->total, val);
3481 else
3482 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003483 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003484 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003485 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003486
Raymond Hettinger482ba772010-12-01 22:48:00 +00003487 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003488 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003489 return newtotal;
3490}
3491
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003492static PyObject *
3493accumulate_reduce(accumulateobject *lz)
3494{
Serhiy Storchakad5516252016-03-06 14:00:45 +02003495 if (lz->total == Py_None) {
3496 PyObject *it;
3497
3498 if (PyType_Ready(&chain_type) < 0)
3499 return NULL;
3500 if (PyType_Ready(&islice_type) < 0)
3501 return NULL;
3502 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3503 lz->total, lz->it);
3504 if (it == NULL)
3505 return NULL;
3506 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3507 it, lz->binop ? lz->binop : Py_None);
3508 if (it == NULL)
3509 return NULL;
3510 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3511 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003512 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3513 lz->it, lz->binop?lz->binop:Py_None,
3514 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003515}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003516
3517static PyObject *
3518accumulate_setstate(accumulateobject *lz, PyObject *state)
3519{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003520 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003521 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003522 Py_RETURN_NONE;
3523}
3524
3525static PyMethodDef accumulate_methods[] = {
3526 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3527 reduce_doc},
3528 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3529 setstate_doc},
3530 {NULL, NULL} /* sentinel */
3531};
3532
Raymond Hettinger482ba772010-12-01 22:48:00 +00003533PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003534"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003535\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003536Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003537
3538static PyTypeObject accumulate_type = {
3539 PyVarObject_HEAD_INIT(NULL, 0)
3540 "itertools.accumulate", /* tp_name */
3541 sizeof(accumulateobject), /* tp_basicsize */
3542 0, /* tp_itemsize */
3543 /* methods */
3544 (destructor)accumulate_dealloc, /* tp_dealloc */
3545 0, /* tp_print */
3546 0, /* tp_getattr */
3547 0, /* tp_setattr */
3548 0, /* tp_reserved */
3549 0, /* tp_repr */
3550 0, /* tp_as_number */
3551 0, /* tp_as_sequence */
3552 0, /* tp_as_mapping */
3553 0, /* tp_hash */
3554 0, /* tp_call */
3555 0, /* tp_str */
3556 PyObject_GenericGetAttr, /* tp_getattro */
3557 0, /* tp_setattro */
3558 0, /* tp_as_buffer */
3559 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3560 Py_TPFLAGS_BASETYPE, /* tp_flags */
3561 accumulate_doc, /* tp_doc */
3562 (traverseproc)accumulate_traverse, /* tp_traverse */
3563 0, /* tp_clear */
3564 0, /* tp_richcompare */
3565 0, /* tp_weaklistoffset */
3566 PyObject_SelfIter, /* tp_iter */
3567 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003568 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003569 0, /* tp_members */
3570 0, /* tp_getset */
3571 0, /* tp_base */
3572 0, /* tp_dict */
3573 0, /* tp_descr_get */
3574 0, /* tp_descr_set */
3575 0, /* tp_dictoffset */
3576 0, /* tp_init */
3577 0, /* tp_alloc */
3578 accumulate_new, /* tp_new */
3579 PyObject_GC_Del, /* tp_free */
3580};
3581
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003582
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003583/* compress object ************************************************************/
3584
3585/* Equivalent to:
3586
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003587 def compress(data, selectors):
3588 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3589 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003590*/
3591
3592typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003593 PyObject_HEAD
3594 PyObject *data;
3595 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003596} compressobject;
3597
3598static PyTypeObject compress_type;
3599
3600static PyObject *
3601compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3602{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003603 PyObject *seq1, *seq2;
3604 PyObject *data=NULL, *selectors=NULL;
3605 compressobject *lz;
3606 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003607
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003608 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3609 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003610
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003611 data = PyObject_GetIter(seq1);
3612 if (data == NULL)
3613 goto fail;
3614 selectors = PyObject_GetIter(seq2);
3615 if (selectors == NULL)
3616 goto fail;
3617
3618 /* create compressobject structure */
3619 lz = (compressobject *)type->tp_alloc(type, 0);
3620 if (lz == NULL)
3621 goto fail;
3622 lz->data = data;
3623 lz->selectors = selectors;
3624 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003625
3626fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003627 Py_XDECREF(data);
3628 Py_XDECREF(selectors);
3629 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003630}
3631
3632static void
3633compress_dealloc(compressobject *lz)
3634{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003635 PyObject_GC_UnTrack(lz);
3636 Py_XDECREF(lz->data);
3637 Py_XDECREF(lz->selectors);
3638 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003639}
3640
3641static int
3642compress_traverse(compressobject *lz, visitproc visit, void *arg)
3643{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003644 Py_VISIT(lz->data);
3645 Py_VISIT(lz->selectors);
3646 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003647}
3648
3649static PyObject *
3650compress_next(compressobject *lz)
3651{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003652 PyObject *data = lz->data, *selectors = lz->selectors;
3653 PyObject *datum, *selector;
3654 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3655 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3656 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003657
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003658 while (1) {
3659 /* Steps: get datum, get selector, evaluate selector.
3660 Order is important (to match the pure python version
3661 in terms of which input gets a chance to raise an
3662 exception first).
3663 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003664
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003665 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003666 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003667 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003668
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003669 selector = selectornext(selectors);
3670 if (selector == NULL) {
3671 Py_DECREF(datum);
3672 return NULL;
3673 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003674
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003675 ok = PyObject_IsTrue(selector);
3676 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003677 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003678 return datum;
3679 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003680 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003681 return NULL;
3682 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003683}
3684
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003685static PyObject *
3686compress_reduce(compressobject *lz)
3687{
3688 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3689 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003690}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003691
3692static PyMethodDef compress_methods[] = {
3693 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3694 reduce_doc},
3695 {NULL, NULL} /* sentinel */
3696};
3697
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003698PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003699"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003700\n\
3701Return data elements corresponding to true selector elements.\n\
3702Forms a shorter iterator from selected data elements using the\n\
3703selectors to choose the data elements.");
3704
3705static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003706 PyVarObject_HEAD_INIT(NULL, 0)
3707 "itertools.compress", /* tp_name */
3708 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003709 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003710 /* methods */
3711 (destructor)compress_dealloc, /* tp_dealloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003712 0, /* tp_print */
3713 0, /* tp_getattr */
3714 0, /* tp_setattr */
3715 0, /* tp_reserved */
3716 0, /* tp_repr */
3717 0, /* tp_as_number */
3718 0, /* tp_as_sequence */
3719 0, /* tp_as_mapping */
3720 0, /* tp_hash */
3721 0, /* tp_call */
3722 0, /* tp_str */
3723 PyObject_GenericGetAttr, /* tp_getattro */
3724 0, /* tp_setattro */
3725 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003726 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003727 Py_TPFLAGS_BASETYPE, /* tp_flags */
3728 compress_doc, /* tp_doc */
3729 (traverseproc)compress_traverse, /* tp_traverse */
3730 0, /* tp_clear */
3731 0, /* tp_richcompare */
3732 0, /* tp_weaklistoffset */
3733 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003734 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003735 compress_methods, /* tp_methods */
3736 0, /* tp_members */
3737 0, /* tp_getset */
3738 0, /* tp_base */
3739 0, /* tp_dict */
3740 0, /* tp_descr_get */
3741 0, /* tp_descr_set */
3742 0, /* tp_dictoffset */
3743 0, /* tp_init */
3744 0, /* tp_alloc */
3745 compress_new, /* tp_new */
3746 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003747};
3748
3749
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003750/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003751
3752typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003753 PyObject_HEAD
3754 PyObject *func;
3755 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003756} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003757
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003758static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003759
3760static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003761filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003762{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003763 PyObject *func, *seq;
3764 PyObject *it;
3765 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003766
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003767 if (type == &filterfalse_type &&
3768 !_PyArg_NoKeywords("filterfalse()", kwds))
3769 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003770
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003771 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3772 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003773
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003774 /* Get iterator. */
3775 it = PyObject_GetIter(seq);
3776 if (it == NULL)
3777 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003778
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003779 /* create filterfalseobject structure */
3780 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3781 if (lz == NULL) {
3782 Py_DECREF(it);
3783 return NULL;
3784 }
3785 Py_INCREF(func);
3786 lz->func = func;
3787 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003788
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003789 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003790}
3791
3792static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003793filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003794{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003795 PyObject_GC_UnTrack(lz);
3796 Py_XDECREF(lz->func);
3797 Py_XDECREF(lz->it);
3798 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003799}
3800
3801static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003802filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003803{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003804 Py_VISIT(lz->it);
3805 Py_VISIT(lz->func);
3806 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003807}
3808
3809static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003810filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003811{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003812 PyObject *item;
3813 PyObject *it = lz->it;
3814 long ok;
3815 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003816
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003817 iternext = *Py_TYPE(it)->tp_iternext;
3818 for (;;) {
3819 item = iternext(it);
3820 if (item == NULL)
3821 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003822
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003823 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3824 ok = PyObject_IsTrue(item);
3825 } else {
3826 PyObject *good;
Raymond Hettingera6a2d442015-08-16 14:38:07 -07003827 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003828 if (good == NULL) {
3829 Py_DECREF(item);
3830 return NULL;
3831 }
3832 ok = PyObject_IsTrue(good);
3833 Py_DECREF(good);
3834 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003835 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003836 return item;
3837 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003838 if (ok < 0)
3839 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003840 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003841}
3842
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003843static PyObject *
3844filterfalse_reduce(filterfalseobject *lz)
3845{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003846 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3847}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003848
3849static PyMethodDef filterfalse_methods[] = {
3850 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3851 reduce_doc},
3852 {NULL, NULL} /* sentinel */
3853};
3854
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003855PyDoc_STRVAR(filterfalse_doc,
3856"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003857\n\
3858Return those items of sequence for which function(item) is false.\n\
3859If function is None, return the items that are false.");
3860
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003861static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003862 PyVarObject_HEAD_INIT(NULL, 0)
3863 "itertools.filterfalse", /* tp_name */
3864 sizeof(filterfalseobject), /* tp_basicsize */
3865 0, /* tp_itemsize */
3866 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003867 (destructor)filterfalse_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003868 0, /* tp_print */
3869 0, /* tp_getattr */
3870 0, /* tp_setattr */
3871 0, /* tp_reserved */
3872 0, /* tp_repr */
3873 0, /* tp_as_number */
3874 0, /* tp_as_sequence */
3875 0, /* tp_as_mapping */
3876 0, /* tp_hash */
3877 0, /* tp_call */
3878 0, /* tp_str */
3879 PyObject_GenericGetAttr, /* tp_getattro */
3880 0, /* tp_setattro */
3881 0, /* tp_as_buffer */
3882 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3883 Py_TPFLAGS_BASETYPE, /* tp_flags */
3884 filterfalse_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003885 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003886 0, /* tp_clear */
3887 0, /* tp_richcompare */
3888 0, /* tp_weaklistoffset */
3889 PyObject_SelfIter, /* tp_iter */
3890 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003891 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003892 0, /* tp_members */
3893 0, /* tp_getset */
3894 0, /* tp_base */
3895 0, /* tp_dict */
3896 0, /* tp_descr_get */
3897 0, /* tp_descr_set */
3898 0, /* tp_dictoffset */
3899 0, /* tp_init */
3900 0, /* tp_alloc */
3901 filterfalse_new, /* tp_new */
3902 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003903};
3904
3905
3906/* count object ************************************************************/
3907
3908typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003909 PyObject_HEAD
3910 Py_ssize_t cnt;
3911 PyObject *long_cnt;
3912 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003913} countobject;
3914
Raymond Hettinger30729212009-02-12 06:28:27 +00003915/* Counting logic and invariants:
3916
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003917fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003918
Serhiy Storchaka483405b2015-02-17 10:14:30 +02003919 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003920 Advances with: cnt += 1
3921 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003922
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003923slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003924
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003925 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3926 All counting is done with python objects (no overflows or underflows).
3927 Advances with: long_cnt += long_step
3928 Step may be zero -- effectively a slow version of repeat(cnt).
3929 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003930*/
3931
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003932static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003933
3934static PyObject *
3935count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3936{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003937 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003938 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003939 Py_ssize_t cnt = 0;
3940 PyObject *long_cnt = NULL;
3941 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003942 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003943 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003944
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003945 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3946 kwlist, &long_cnt, &long_step))
3947 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003948
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003949 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3950 (long_step != NULL && !PyNumber_Check(long_step))) {
3951 PyErr_SetString(PyExc_TypeError, "a number is required");
3952 return NULL;
3953 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003954
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003955 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
3956 (long_step == NULL || PyLong_Check(long_step));
3957
3958 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003959 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003960 if (fast_mode) {
3961 assert(PyLong_Check(long_cnt));
3962 cnt = PyLong_AsSsize_t(long_cnt);
3963 if (cnt == -1 && PyErr_Occurred()) {
3964 PyErr_Clear();
3965 fast_mode = 0;
3966 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003967 }
3968 Py_INCREF(long_cnt);
3969 } else {
3970 cnt = 0;
3971 long_cnt = PyLong_FromLong(0);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003972 if (long_cnt == NULL) {
3973 return NULL;
3974 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003975 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003976
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003977 /* If not specified, step defaults to 1 */
3978 if (long_step == NULL) {
3979 long_step = PyLong_FromLong(1);
3980 if (long_step == NULL) {
3981 Py_DECREF(long_cnt);
3982 return NULL;
3983 }
3984 } else
3985 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003986
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003987 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003988
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003989 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003990 if (fast_mode) {
3991 assert(PyLong_Check(long_step));
3992 step = PyLong_AsLong(long_step);
3993 if (step != 1) {
3994 fast_mode = 0;
3995 if (step == -1 && PyErr_Occurred())
3996 PyErr_Clear();
3997 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003998 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003999
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004000 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004001 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004002 else
4003 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004004
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004005 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4006 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4007 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004008 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004009
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004010 /* create countobject structure */
4011 lz = (countobject *)type->tp_alloc(type, 0);
4012 if (lz == NULL) {
4013 Py_XDECREF(long_cnt);
4014 return NULL;
4015 }
4016 lz->cnt = cnt;
4017 lz->long_cnt = long_cnt;
4018 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004019
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004020 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004021}
4022
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004023static void
4024count_dealloc(countobject *lz)
4025{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004026 PyObject_GC_UnTrack(lz);
4027 Py_XDECREF(lz->long_cnt);
4028 Py_XDECREF(lz->long_step);
4029 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004030}
4031
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004032static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004033count_traverse(countobject *lz, visitproc visit, void *arg)
4034{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004035 Py_VISIT(lz->long_cnt);
4036 Py_VISIT(lz->long_step);
4037 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004038}
4039
4040static PyObject *
4041count_nextlong(countobject *lz)
4042{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004043 PyObject *long_cnt;
4044 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004045
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004046 long_cnt = lz->long_cnt;
4047 if (long_cnt == NULL) {
4048 /* Switch to slow_mode */
4049 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4050 if (long_cnt == NULL)
4051 return NULL;
4052 }
4053 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004054
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004055 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4056 if (stepped_up == NULL)
4057 return NULL;
4058 lz->long_cnt = stepped_up;
4059 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004060}
4061
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004062static PyObject *
4063count_next(countobject *lz)
4064{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004065 if (lz->cnt == PY_SSIZE_T_MAX)
4066 return count_nextlong(lz);
4067 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004068}
4069
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004070static PyObject *
4071count_repr(countobject *lz)
4072{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004073 if (lz->cnt != PY_SSIZE_T_MAX)
4074 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004075
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004076 if (PyLong_Check(lz->long_step)) {
4077 long step = PyLong_AsLong(lz->long_step);
4078 if (step == -1 && PyErr_Occurred()) {
4079 PyErr_Clear();
4080 }
4081 if (step == 1) {
4082 /* Don't display step when it is an integer equal to 1 */
4083 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
4084 }
4085 }
4086 return PyUnicode_FromFormat("count(%R, %R)",
4087 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004088}
4089
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004090static PyObject *
4091count_reduce(countobject *lz)
4092{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004093 if (lz->cnt == PY_SSIZE_T_MAX)
4094 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4095 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004096}
4097
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004098static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004099 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004100 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004101 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004102};
4103
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004104PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00004105 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004106\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004107Return a count object whose .__next__() method returns consecutive values.\n\
4108Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00004109 def count(firstval=0, step=1):\n\
Eli Benderskyc554f722013-09-03 06:37:19 -07004110 x = firstval\n\
4111 while 1:\n\
4112 yield x\n\
4113 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004114
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004115static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004116 PyVarObject_HEAD_INIT(NULL, 0)
4117 "itertools.count", /* tp_name */
4118 sizeof(countobject), /* tp_basicsize */
4119 0, /* tp_itemsize */
4120 /* methods */
4121 (destructor)count_dealloc, /* tp_dealloc */
4122 0, /* tp_print */
4123 0, /* tp_getattr */
4124 0, /* tp_setattr */
4125 0, /* tp_reserved */
4126 (reprfunc)count_repr, /* tp_repr */
4127 0, /* tp_as_number */
4128 0, /* tp_as_sequence */
4129 0, /* tp_as_mapping */
4130 0, /* tp_hash */
4131 0, /* tp_call */
4132 0, /* tp_str */
4133 PyObject_GenericGetAttr, /* tp_getattro */
4134 0, /* tp_setattro */
4135 0, /* tp_as_buffer */
4136 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004137 Py_TPFLAGS_BASETYPE, /* tp_flags */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004138 count_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004139 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004140 0, /* tp_clear */
4141 0, /* tp_richcompare */
4142 0, /* tp_weaklistoffset */
4143 PyObject_SelfIter, /* tp_iter */
4144 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004145 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004146 0, /* tp_members */
4147 0, /* tp_getset */
4148 0, /* tp_base */
4149 0, /* tp_dict */
4150 0, /* tp_descr_get */
4151 0, /* tp_descr_set */
4152 0, /* tp_dictoffset */
4153 0, /* tp_init */
4154 0, /* tp_alloc */
4155 count_new, /* tp_new */
4156 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004157};
4158
4159
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004160/* repeat object ************************************************************/
4161
4162typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004163 PyObject_HEAD
4164 PyObject *element;
4165 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004166} repeatobject;
4167
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004168static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004169
4170static PyObject *
4171repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4172{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004173 repeatobject *ro;
4174 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004175 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004176 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004177
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004178 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4179 &element, &cnt))
4180 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004181
Raymond Hettinger97d35552014-06-24 21:36:58 -07004182 if (kwds != NULL)
4183 n_kwds = PyDict_Size(kwds);
4184 /* Does user supply times argument? */
4185 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004186 cnt = 0;
4187
4188 ro = (repeatobject *)type->tp_alloc(type, 0);
4189 if (ro == NULL)
4190 return NULL;
4191 Py_INCREF(element);
4192 ro->element = element;
4193 ro->cnt = cnt;
4194 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004195}
4196
4197static void
4198repeat_dealloc(repeatobject *ro)
4199{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004200 PyObject_GC_UnTrack(ro);
4201 Py_XDECREF(ro->element);
4202 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004203}
4204
4205static int
4206repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4207{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004208 Py_VISIT(ro->element);
4209 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004210}
4211
4212static PyObject *
4213repeat_next(repeatobject *ro)
4214{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004215 if (ro->cnt == 0)
4216 return NULL;
4217 if (ro->cnt > 0)
4218 ro->cnt--;
4219 Py_INCREF(ro->element);
4220 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004221}
4222
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004223static PyObject *
4224repeat_repr(repeatobject *ro)
4225{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004226 if (ro->cnt == -1)
4227 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4228 else
4229 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4230}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004231
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004232static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004233repeat_len(repeatobject *ro)
4234{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004235 if (ro->cnt == -1) {
4236 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4237 return NULL;
4238 }
4239 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004240}
4241
Armin Rigof5b3e362006-02-11 21:32:43 +00004242PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004243
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004244static PyObject *
4245repeat_reduce(repeatobject *ro)
4246{
4247 /* unpickle this so that a new repeat iterator is constructed with an
4248 * object, then call __setstate__ on it to set cnt
4249 */
4250 if (ro->cnt >= 0)
4251 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4252 else
4253 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4254}
4255
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004256static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004257 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004258 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004259 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004260};
4261
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004262PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004263"repeat(object [,times]) -> create an iterator which returns the object\n\
4264for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004265endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004266
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004267static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004268 PyVarObject_HEAD_INIT(NULL, 0)
4269 "itertools.repeat", /* tp_name */
4270 sizeof(repeatobject), /* tp_basicsize */
4271 0, /* tp_itemsize */
4272 /* methods */
4273 (destructor)repeat_dealloc, /* tp_dealloc */
4274 0, /* tp_print */
4275 0, /* tp_getattr */
4276 0, /* tp_setattr */
4277 0, /* tp_reserved */
4278 (reprfunc)repeat_repr, /* tp_repr */
4279 0, /* tp_as_number */
4280 0, /* tp_as_sequence */
4281 0, /* tp_as_mapping */
4282 0, /* tp_hash */
4283 0, /* tp_call */
4284 0, /* tp_str */
4285 PyObject_GenericGetAttr, /* tp_getattro */
4286 0, /* tp_setattro */
4287 0, /* tp_as_buffer */
4288 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4289 Py_TPFLAGS_BASETYPE, /* tp_flags */
4290 repeat_doc, /* tp_doc */
4291 (traverseproc)repeat_traverse, /* tp_traverse */
4292 0, /* tp_clear */
4293 0, /* tp_richcompare */
4294 0, /* tp_weaklistoffset */
4295 PyObject_SelfIter, /* tp_iter */
4296 (iternextfunc)repeat_next, /* tp_iternext */
4297 repeat_methods, /* tp_methods */
4298 0, /* tp_members */
4299 0, /* tp_getset */
4300 0, /* tp_base */
4301 0, /* tp_dict */
4302 0, /* tp_descr_get */
4303 0, /* tp_descr_set */
4304 0, /* tp_dictoffset */
4305 0, /* tp_init */
4306 0, /* tp_alloc */
4307 repeat_new, /* tp_new */
4308 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004309};
4310
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004311/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004312
4313typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004314 PyObject_HEAD
4315 Py_ssize_t tuplesize;
4316 Py_ssize_t numactive;
4317 PyObject *ittuple; /* tuple of iterators */
4318 PyObject *result;
4319 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004320} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004321
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004322static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004323
4324static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004325zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004326{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004327 ziplongestobject *lz;
4328 Py_ssize_t i;
4329 PyObject *ittuple; /* tuple of iterators */
4330 PyObject *result;
4331 PyObject *fillvalue = Py_None;
4332 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004333
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004334 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
4335 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4336 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
4337 PyErr_SetString(PyExc_TypeError,
4338 "zip_longest() got an unexpected keyword argument");
4339 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004340 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004341 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004342
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004343 /* args must be a tuple */
4344 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00004345
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004346 /* obtain iterators */
4347 ittuple = PyTuple_New(tuplesize);
4348 if (ittuple == NULL)
4349 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004350 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004351 PyObject *item = PyTuple_GET_ITEM(args, i);
4352 PyObject *it = PyObject_GetIter(item);
4353 if (it == NULL) {
4354 if (PyErr_ExceptionMatches(PyExc_TypeError))
4355 PyErr_Format(PyExc_TypeError,
4356 "zip_longest argument #%zd must support iteration",
4357 i+1);
4358 Py_DECREF(ittuple);
4359 return NULL;
4360 }
4361 PyTuple_SET_ITEM(ittuple, i, it);
4362 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004363
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004364 /* create a result holder */
4365 result = PyTuple_New(tuplesize);
4366 if (result == NULL) {
4367 Py_DECREF(ittuple);
4368 return NULL;
4369 }
4370 for (i=0 ; i < tuplesize ; i++) {
4371 Py_INCREF(Py_None);
4372 PyTuple_SET_ITEM(result, i, Py_None);
4373 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004374
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004375 /* create ziplongestobject structure */
4376 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4377 if (lz == NULL) {
4378 Py_DECREF(ittuple);
4379 Py_DECREF(result);
4380 return NULL;
4381 }
4382 lz->ittuple = ittuple;
4383 lz->tuplesize = tuplesize;
4384 lz->numactive = tuplesize;
4385 lz->result = result;
4386 Py_INCREF(fillvalue);
4387 lz->fillvalue = fillvalue;
4388 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004389}
4390
4391static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004392zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004393{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004394 PyObject_GC_UnTrack(lz);
4395 Py_XDECREF(lz->ittuple);
4396 Py_XDECREF(lz->result);
4397 Py_XDECREF(lz->fillvalue);
4398 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004399}
4400
4401static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004402zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004403{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004404 Py_VISIT(lz->ittuple);
4405 Py_VISIT(lz->result);
4406 Py_VISIT(lz->fillvalue);
4407 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004408}
4409
4410static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004411zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004412{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004413 Py_ssize_t i;
4414 Py_ssize_t tuplesize = lz->tuplesize;
4415 PyObject *result = lz->result;
4416 PyObject *it;
4417 PyObject *item;
4418 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004419
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004420 if (tuplesize == 0)
4421 return NULL;
4422 if (lz->numactive == 0)
4423 return NULL;
4424 if (Py_REFCNT(result) == 1) {
4425 Py_INCREF(result);
4426 for (i=0 ; i < tuplesize ; i++) {
4427 it = PyTuple_GET_ITEM(lz->ittuple, i);
4428 if (it == NULL) {
4429 Py_INCREF(lz->fillvalue);
4430 item = lz->fillvalue;
4431 } else {
4432 item = PyIter_Next(it);
4433 if (item == NULL) {
4434 lz->numactive -= 1;
4435 if (lz->numactive == 0 || PyErr_Occurred()) {
4436 lz->numactive = 0;
4437 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004438 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004439 } else {
4440 Py_INCREF(lz->fillvalue);
4441 item = lz->fillvalue;
4442 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4443 Py_DECREF(it);
4444 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004445 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004446 }
4447 olditem = PyTuple_GET_ITEM(result, i);
4448 PyTuple_SET_ITEM(result, i, item);
4449 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004450 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004451 } else {
4452 result = PyTuple_New(tuplesize);
4453 if (result == NULL)
4454 return NULL;
4455 for (i=0 ; i < tuplesize ; i++) {
4456 it = PyTuple_GET_ITEM(lz->ittuple, i);
4457 if (it == NULL) {
4458 Py_INCREF(lz->fillvalue);
4459 item = lz->fillvalue;
4460 } else {
4461 item = PyIter_Next(it);
4462 if (item == NULL) {
4463 lz->numactive -= 1;
4464 if (lz->numactive == 0 || PyErr_Occurred()) {
4465 lz->numactive = 0;
4466 Py_DECREF(result);
4467 return NULL;
4468 } else {
4469 Py_INCREF(lz->fillvalue);
4470 item = lz->fillvalue;
4471 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4472 Py_DECREF(it);
4473 }
4474 }
4475 }
4476 PyTuple_SET_ITEM(result, i, item);
4477 }
4478 }
4479 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004480}
4481
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004482static PyObject *
4483zip_longest_reduce(ziplongestobject *lz)
4484{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004485
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004486 /* Create a new tuple with empty sequences where appropriate to pickle.
4487 * Then use setstate to set the fillvalue
4488 */
4489 int i;
4490 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004491
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004492 if (args == NULL)
4493 return NULL;
4494 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4495 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4496 if (elem == NULL) {
4497 elem = PyTuple_New(0);
4498 if (elem == NULL) {
4499 Py_DECREF(args);
4500 return NULL;
4501 }
4502 } else
4503 Py_INCREF(elem);
4504 PyTuple_SET_ITEM(args, i, elem);
4505 }
4506 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4507}
4508
4509static PyObject *
4510zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4511{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004512 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004513 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004514 Py_RETURN_NONE;
4515}
4516
4517static PyMethodDef zip_longest_methods[] = {
4518 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4519 reduce_doc},
4520 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4521 setstate_doc},
4522 {NULL, NULL} /* sentinel */
4523};
4524
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004525PyDoc_STRVAR(zip_longest_doc,
4526"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004527\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004528Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004529the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004530method continues until the longest iterable in the argument sequence\n\
4531is exhausted and then it raises StopIteration. When the shorter iterables\n\
4532are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4533defaults to None or can be specified by a keyword argument.\n\
4534");
4535
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004536static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004537 PyVarObject_HEAD_INIT(NULL, 0)
4538 "itertools.zip_longest", /* tp_name */
4539 sizeof(ziplongestobject), /* tp_basicsize */
4540 0, /* tp_itemsize */
4541 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004542 (destructor)zip_longest_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004543 0, /* tp_print */
4544 0, /* tp_getattr */
4545 0, /* tp_setattr */
4546 0, /* tp_reserved */
4547 0, /* tp_repr */
4548 0, /* tp_as_number */
4549 0, /* tp_as_sequence */
4550 0, /* tp_as_mapping */
4551 0, /* tp_hash */
4552 0, /* tp_call */
4553 0, /* tp_str */
4554 PyObject_GenericGetAttr, /* tp_getattro */
4555 0, /* tp_setattro */
4556 0, /* tp_as_buffer */
4557 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4558 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004559 zip_longest_doc, /* tp_doc */
4560 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004561 0, /* tp_clear */
4562 0, /* tp_richcompare */
4563 0, /* tp_weaklistoffset */
4564 PyObject_SelfIter, /* tp_iter */
4565 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004566 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004567 0, /* tp_members */
4568 0, /* tp_getset */
4569 0, /* tp_base */
4570 0, /* tp_dict */
4571 0, /* tp_descr_get */
4572 0, /* tp_descr_set */
4573 0, /* tp_dictoffset */
4574 0, /* tp_init */
4575 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004576 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004577 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004578};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004579
4580/* module level code ********************************************************/
4581
4582PyDoc_STRVAR(module_doc,
4583"Functional tools for creating and using iterators.\n\
4584\n\
4585Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004586count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004587cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004588repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004589\n\
4590Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004591accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004592chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger8df58f72013-09-09 01:29:40 -05004593chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004594compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4595dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4596groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004597filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004598islice(seq, [start,] stop [, step]) --> elements from\n\
4599 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004600starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004601tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004602takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004603zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004604\n\
4605Combinatoric generators:\n\
4606product(p, q, ... [repeat=1]) --> cartesian product\n\
4607permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004608combinations(p, r)\n\
4609combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004610");
4611
4612
Raymond Hettingerad983e72003-11-12 14:32:26 +00004613static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004614 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4615 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004616};
4617
Martin v. Löwis1a214512008-06-11 05:26:20 +00004618
4619static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004620 PyModuleDef_HEAD_INIT,
4621 "itertools",
4622 module_doc,
4623 -1,
4624 module_methods,
4625 NULL,
4626 NULL,
4627 NULL,
4628 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004629};
4630
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004631PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004632PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004633{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004634 int i;
4635 PyObject *m;
4636 char *name;
4637 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004638 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004639 &combinations_type,
4640 &cwr_type,
4641 &cycle_type,
4642 &dropwhile_type,
4643 &takewhile_type,
4644 &islice_type,
4645 &starmap_type,
4646 &chain_type,
4647 &compress_type,
4648 &filterfalse_type,
4649 &count_type,
4650 &ziplongest_type,
4651 &permutations_type,
4652 &product_type,
4653 &repeat_type,
4654 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004655 &_grouper_type,
4656 &tee_type,
4657 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004658 NULL
4659 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004660
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004661 Py_TYPE(&teedataobject_type) = &PyType_Type;
4662 m = PyModule_Create(&itertoolsmodule);
4663 if (m == NULL)
4664 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004665
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004666 for (i=0 ; typelist[i] != NULL ; i++) {
4667 if (PyType_Ready(typelist[i]) < 0)
4668 return NULL;
4669 name = strchr(typelist[i]->tp_name, '.');
4670 assert (name != NULL);
4671 Py_INCREF(typelist[i]);
4672 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4673 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004674
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004675 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004676}