blob: 919023636498826a393490bfd47c8187f397a115 [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;
Serhiy Storchakac247caf2017-09-24 13:36:11 +030020 const void *currgrouper; /* borrowed reference */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000021} groupbyobject;
22
23static PyTypeObject groupby_type;
24static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26static PyObject *
27groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000029 static char *kwargs[] = {"iterable", "key", NULL};
30 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000032
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000033 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
34 &it, &keyfunc))
35 return NULL;
36
37 gbo = (groupbyobject *)type->tp_alloc(type, 0);
38 if (gbo == NULL)
39 return NULL;
40 gbo->tgtkey = NULL;
41 gbo->currkey = NULL;
42 gbo->currvalue = NULL;
43 gbo->keyfunc = keyfunc;
44 Py_INCREF(keyfunc);
45 gbo->it = PyObject_GetIter(it);
46 if (gbo->it == NULL) {
47 Py_DECREF(gbo);
48 return NULL;
49 }
50 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000051}
52
53static void
54groupby_dealloc(groupbyobject *gbo)
55{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000056 PyObject_GC_UnTrack(gbo);
57 Py_XDECREF(gbo->it);
58 Py_XDECREF(gbo->keyfunc);
59 Py_XDECREF(gbo->tgtkey);
60 Py_XDECREF(gbo->currkey);
61 Py_XDECREF(gbo->currvalue);
62 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000063}
64
65static int
66groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
67{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000068 Py_VISIT(gbo->it);
69 Py_VISIT(gbo->keyfunc);
70 Py_VISIT(gbo->tgtkey);
71 Py_VISIT(gbo->currkey);
72 Py_VISIT(gbo->currvalue);
73 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000074}
75
Serhiy Storchakac740e4f2017-09-26 21:47:56 +030076Py_LOCAL_INLINE(int)
77groupby_step(groupbyobject *gbo)
78{
79 PyObject *newvalue, *newkey, *oldvalue;
80
81 newvalue = PyIter_Next(gbo->it);
82 if (newvalue == NULL)
83 return -1;
84
85 if (gbo->keyfunc == Py_None) {
86 newkey = newvalue;
87 Py_INCREF(newvalue);
88 } else {
89 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL);
90 if (newkey == NULL) {
91 Py_DECREF(newvalue);
92 return -1;
93 }
94 }
95
96 oldvalue = gbo->currvalue;
97 gbo->currvalue = newvalue;
98 Py_XSETREF(gbo->currkey, newkey);
99 Py_XDECREF(oldvalue);
100 return 0;
101}
102
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000103static PyObject *
104groupby_next(groupbyobject *gbo)
105{
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300106 PyObject *r, *grouper;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000107
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300108 gbo->currgrouper = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000109 /* skip to next iteration group */
110 for (;;) {
111 if (gbo->currkey == NULL)
112 /* pass */;
113 else if (gbo->tgtkey == NULL)
114 break;
115 else {
116 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000117
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700118 rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000119 if (rcmp == -1)
120 return NULL;
121 else if (rcmp == 0)
122 break;
123 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000124
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300125 if (groupby_step(gbo) < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000126 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000127 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000128 Py_INCREF(gbo->currkey);
Serhiy Storchakaec397562016-04-06 09:50:03 +0300129 Py_XSETREF(gbo->tgtkey, gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000130
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000131 grouper = _grouper_create(gbo, gbo->tgtkey);
132 if (grouper == NULL)
133 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000134
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000135 r = PyTuple_Pack(2, gbo->currkey, grouper);
136 Py_DECREF(grouper);
137 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000138}
139
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000140static PyObject *
141groupby_reduce(groupbyobject *lz)
142{
143 /* reduce as a 'new' call with an optional 'setstate' if groupby
144 * has started
145 */
146 PyObject *value;
147 if (lz->tgtkey && lz->currkey && lz->currvalue)
148 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
149 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
150 else
151 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
152 lz->it, lz->keyfunc);
153
154 return value;
155}
156
157PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
158
159static PyObject *
160groupby_setstate(groupbyobject *lz, PyObject *state)
161{
162 PyObject *currkey, *currvalue, *tgtkey;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300163 if (!PyTuple_Check(state)) {
164 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000165 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300166 }
167 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey)) {
168 return NULL;
169 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200170 Py_INCREF(currkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300171 Py_XSETREF(lz->currkey, currkey);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200172 Py_INCREF(currvalue);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300173 Py_XSETREF(lz->currvalue, currvalue);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200174 Py_INCREF(tgtkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300175 Py_XSETREF(lz->tgtkey, tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000176 Py_RETURN_NONE;
177}
178
179PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
180
181static PyMethodDef groupby_methods[] = {
182 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
183 reduce_doc},
184 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
185 setstate_doc},
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700186 {NULL, NULL} /* sentinel */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000187};
188
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000189PyDoc_STRVAR(groupby_doc,
Raymond Hettinger49392c62017-09-25 01:21:06 -0700190"groupby(iterable, key=None) -> make an iterator that returns consecutive\n\
191keys and groups from the iterable. If the key function is not specified or\n\
192is None, the element itself is used for grouping.\n");
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000193
194static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000195 PyVarObject_HEAD_INIT(NULL, 0)
196 "itertools.groupby", /* tp_name */
197 sizeof(groupbyobject), /* tp_basicsize */
198 0, /* tp_itemsize */
199 /* methods */
200 (destructor)groupby_dealloc, /* tp_dealloc */
201 0, /* tp_print */
202 0, /* tp_getattr */
203 0, /* tp_setattr */
204 0, /* tp_reserved */
205 0, /* tp_repr */
206 0, /* tp_as_number */
207 0, /* tp_as_sequence */
208 0, /* tp_as_mapping */
209 0, /* tp_hash */
210 0, /* tp_call */
211 0, /* tp_str */
212 PyObject_GenericGetAttr, /* tp_getattro */
213 0, /* tp_setattro */
214 0, /* tp_as_buffer */
215 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
216 Py_TPFLAGS_BASETYPE, /* tp_flags */
217 groupby_doc, /* tp_doc */
218 (traverseproc)groupby_traverse, /* tp_traverse */
219 0, /* tp_clear */
220 0, /* tp_richcompare */
221 0, /* tp_weaklistoffset */
222 PyObject_SelfIter, /* tp_iter */
223 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000224 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000225 0, /* tp_members */
226 0, /* tp_getset */
227 0, /* tp_base */
228 0, /* tp_dict */
229 0, /* tp_descr_get */
230 0, /* tp_descr_set */
231 0, /* tp_dictoffset */
232 0, /* tp_init */
233 0, /* tp_alloc */
234 groupby_new, /* tp_new */
235 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000236};
237
238
239/* _grouper object (internal) ************************************************/
240
241typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000242 PyObject_HEAD
243 PyObject *parent;
244 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000245} _grouperobject;
246
247static PyTypeObject _grouper_type;
248
249static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000250_grouper_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
251{
252 PyObject *parent, *tgtkey;
253
254 if (!PyArg_ParseTuple(args, "O!O", &groupby_type, &parent, &tgtkey))
255 return NULL;
256
257 return _grouper_create((groupbyobject*) parent, tgtkey);
258}
259
260static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000261_grouper_create(groupbyobject *parent, PyObject *tgtkey)
262{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000263 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000264
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000265 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
266 if (igo == NULL)
267 return NULL;
268 igo->parent = (PyObject *)parent;
269 Py_INCREF(parent);
270 igo->tgtkey = tgtkey;
271 Py_INCREF(tgtkey);
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300272 parent->currgrouper = igo; /* borrowed reference */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000273
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000274 PyObject_GC_Track(igo);
275 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000276}
277
278static void
279_grouper_dealloc(_grouperobject *igo)
280{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000281 PyObject_GC_UnTrack(igo);
282 Py_DECREF(igo->parent);
283 Py_DECREF(igo->tgtkey);
284 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000285}
286
287static int
288_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
289{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000290 Py_VISIT(igo->parent);
291 Py_VISIT(igo->tgtkey);
292 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000293}
294
295static PyObject *
296_grouper_next(_grouperobject *igo)
297{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000298 groupbyobject *gbo = (groupbyobject *)igo->parent;
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300299 PyObject *r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000300 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000301
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300302 if (gbo->currgrouper != igo)
303 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000304 if (gbo->currvalue == NULL) {
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300305 if (groupby_step(gbo) < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000306 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000307 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000308
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000309 assert(gbo->currkey != NULL);
310 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
311 if (rcmp <= 0)
312 /* got any error or current group is end */
313 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000314
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000315 r = gbo->currvalue;
316 gbo->currvalue = NULL;
317 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000318
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000319 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000320}
321
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000322static PyObject *
323_grouper_reduce(_grouperobject *lz)
324{
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300325 if (((groupbyobject *)lz->parent)->currgrouper != lz) {
326 return Py_BuildValue("N(())", _PyObject_GetBuiltin("iter"));
327 }
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700328 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000329}
330
331static PyMethodDef _grouper_methods[] = {
332 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
333 reduce_doc},
334 {NULL, NULL} /* sentinel */
335};
336
337
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000338static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000339 PyVarObject_HEAD_INIT(NULL, 0)
340 "itertools._grouper", /* tp_name */
341 sizeof(_grouperobject), /* tp_basicsize */
342 0, /* tp_itemsize */
343 /* methods */
344 (destructor)_grouper_dealloc, /* tp_dealloc */
345 0, /* tp_print */
346 0, /* tp_getattr */
347 0, /* tp_setattr */
348 0, /* tp_reserved */
349 0, /* tp_repr */
350 0, /* tp_as_number */
351 0, /* tp_as_sequence */
352 0, /* tp_as_mapping */
353 0, /* tp_hash */
354 0, /* tp_call */
355 0, /* tp_str */
356 PyObject_GenericGetAttr, /* tp_getattro */
357 0, /* tp_setattro */
358 0, /* tp_as_buffer */
359 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
360 0, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700361 (traverseproc)_grouper_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000362 0, /* tp_clear */
363 0, /* tp_richcompare */
364 0, /* tp_weaklistoffset */
365 PyObject_SelfIter, /* tp_iter */
366 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000367 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000368 0, /* tp_members */
369 0, /* tp_getset */
370 0, /* tp_base */
371 0, /* tp_dict */
372 0, /* tp_descr_get */
373 0, /* tp_descr_set */
374 0, /* tp_dictoffset */
375 0, /* tp_init */
376 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000377 _grouper_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000378 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000379};
380
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000381
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700382/* tee object and with supporting function and objects ***********************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000383
Raymond Hettingerad983e72003-11-12 14:32:26 +0000384/* The teedataobject pre-allocates space for LINKCELLS number of objects.
385 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000386 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000387 the other structure members including PyHEAD overhead. The larger the
388 value, the less memory overhead per object and the less time spent
389 allocating/deallocating new links. The smaller the number, the less
390 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000391*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000392#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000393
394typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000395 PyObject_HEAD
396 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700397 int numread; /* 0 <= numread <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000398 PyObject *nextlink;
399 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000400} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000401
402typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000403 PyObject_HEAD
404 teedataobject *dataobj;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700405 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000406 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000407} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000408
Raymond Hettingerad983e72003-11-12 14:32:26 +0000409static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000410
411static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000412teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000413{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000414 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000416 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
417 if (tdo == NULL)
418 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000419
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000420 tdo->numread = 0;
421 tdo->nextlink = NULL;
422 Py_INCREF(it);
423 tdo->it = it;
424 PyObject_GC_Track(tdo);
425 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000426}
427
428static PyObject *
429teedataobject_jumplink(teedataobject *tdo)
430{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000431 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000432 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000433 Py_XINCREF(tdo->nextlink);
434 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000435}
436
437static PyObject *
438teedataobject_getitem(teedataobject *tdo, int i)
439{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000440 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000441
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000442 assert(i < LINKCELLS);
443 if (i < tdo->numread)
444 value = tdo->values[i];
445 else {
446 /* this is the lead iterator, so fetch more data */
447 assert(i == tdo->numread);
448 value = PyIter_Next(tdo->it);
449 if (value == NULL)
450 return NULL;
451 tdo->numread++;
452 tdo->values[i] = value;
453 }
454 Py_INCREF(value);
455 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000456}
457
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000458static int
459teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
460{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000461 int i;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700462
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000463 Py_VISIT(tdo->it);
464 for (i = 0; i < tdo->numread; i++)
465 Py_VISIT(tdo->values[i]);
466 Py_VISIT(tdo->nextlink);
467 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000468}
469
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200470static void
471teedataobject_safe_decref(PyObject *obj)
472{
473 while (obj && Py_TYPE(obj) == &teedataobject_type &&
474 Py_REFCNT(obj) == 1) {
475 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
476 ((teedataobject *)obj)->nextlink = NULL;
477 Py_DECREF(obj);
478 obj = nextlink;
479 }
480 Py_XDECREF(obj);
481}
482
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000483static int
484teedataobject_clear(teedataobject *tdo)
485{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000486 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200487 PyObject *tmp;
488
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000489 Py_CLEAR(tdo->it);
490 for (i=0 ; i<tdo->numread ; i++)
491 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200492 tmp = tdo->nextlink;
493 tdo->nextlink = NULL;
494 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000495 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000496}
497
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000498static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000499teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000500{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000501 PyObject_GC_UnTrack(tdo);
502 teedataobject_clear(tdo);
503 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000504}
505
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000506static PyObject *
507teedataobject_reduce(teedataobject *tdo)
508{
509 int i;
510 /* create a temporary list of already iterated values */
511 PyObject *values = PyList_New(tdo->numread);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700512
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000513 if (!values)
514 return NULL;
515 for (i=0 ; i<tdo->numread ; i++) {
516 Py_INCREF(tdo->values[i]);
517 PyList_SET_ITEM(values, i, tdo->values[i]);
518 }
519 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
520 values,
521 tdo->nextlink ? tdo->nextlink : Py_None);
522}
523
524static PyTypeObject teedataobject_type;
525
526static PyObject *
527teedataobject_new(PyTypeObject *type, PyObject *args, PyObject *kw)
528{
529 teedataobject *tdo;
530 PyObject *it, *values, *next;
531 Py_ssize_t i, len;
532
533 assert(type == &teedataobject_type);
534 if (!PyArg_ParseTuple(args, "OO!O", &it, &PyList_Type, &values, &next))
535 return NULL;
536
537 tdo = (teedataobject *)teedataobject_newinternal(it);
538 if (!tdo)
539 return NULL;
540
541 len = PyList_GET_SIZE(values);
542 if (len > LINKCELLS)
543 goto err;
544 for (i=0; i<len; i++) {
545 tdo->values[i] = PyList_GET_ITEM(values, i);
546 Py_INCREF(tdo->values[i]);
547 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200548 /* len <= LINKCELLS < INT_MAX */
549 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000550
551 if (len == LINKCELLS) {
552 if (next != Py_None) {
553 if (Py_TYPE(next) != &teedataobject_type)
554 goto err;
555 assert(tdo->nextlink == NULL);
556 Py_INCREF(next);
557 tdo->nextlink = next;
558 }
559 } else {
560 if (next != Py_None)
561 goto err; /* shouldn't have a next if we are not full */
562 }
563 return (PyObject*)tdo;
564
565err:
566 Py_XDECREF(tdo);
567 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
568 return NULL;
569}
570
571static PyMethodDef teedataobject_methods[] = {
572 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
573 reduce_doc},
574 {NULL, NULL} /* sentinel */
575};
576
Raymond Hettingerad983e72003-11-12 14:32:26 +0000577PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000578
Raymond Hettingerad983e72003-11-12 14:32:26 +0000579static PyTypeObject teedataobject_type = {
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700580 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000581 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000582 sizeof(teedataobject), /* tp_basicsize */
583 0, /* tp_itemsize */
584 /* methods */
585 (destructor)teedataobject_dealloc, /* tp_dealloc */
586 0, /* tp_print */
587 0, /* tp_getattr */
588 0, /* tp_setattr */
589 0, /* tp_reserved */
590 0, /* tp_repr */
591 0, /* tp_as_number */
592 0, /* tp_as_sequence */
593 0, /* tp_as_mapping */
594 0, /* tp_hash */
595 0, /* tp_call */
596 0, /* tp_str */
597 PyObject_GenericGetAttr, /* tp_getattro */
598 0, /* tp_setattro */
599 0, /* tp_as_buffer */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700600 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000601 teedataobject_doc, /* tp_doc */
602 (traverseproc)teedataobject_traverse, /* tp_traverse */
603 (inquiry)teedataobject_clear, /* tp_clear */
604 0, /* tp_richcompare */
605 0, /* tp_weaklistoffset */
606 0, /* tp_iter */
607 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000608 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000609 0, /* tp_members */
610 0, /* tp_getset */
611 0, /* tp_base */
612 0, /* tp_dict */
613 0, /* tp_descr_get */
614 0, /* tp_descr_set */
615 0, /* tp_dictoffset */
616 0, /* tp_init */
617 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000618 teedataobject_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000619 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000620};
621
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000622
623static PyTypeObject tee_type;
624
625static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000626tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000627{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000628 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000629
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000630 if (to->index >= LINKCELLS) {
631 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200632 if (link == NULL)
633 return NULL;
Serhiy Storchaka57a01d32016-04-10 18:05:40 +0300634 Py_SETREF(to->dataobj, (teedataobject *)link);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000635 to->index = 0;
636 }
637 value = teedataobject_getitem(to->dataobj, to->index);
638 if (value == NULL)
639 return NULL;
640 to->index++;
641 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000642}
643
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000644static int
645tee_traverse(teeobject *to, visitproc visit, void *arg)
646{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000647 Py_VISIT((PyObject *)to->dataobj);
648 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000649}
650
Raymond Hettingerad983e72003-11-12 14:32:26 +0000651static PyObject *
652tee_copy(teeobject *to)
653{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000654 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000655
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000656 newto = PyObject_GC_New(teeobject, &tee_type);
657 if (newto == NULL)
658 return NULL;
659 Py_INCREF(to->dataobj);
660 newto->dataobj = to->dataobj;
661 newto->index = to->index;
662 newto->weakreflist = NULL;
663 PyObject_GC_Track(newto);
664 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000665}
666
667PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
668
669static PyObject *
670tee_fromiterable(PyObject *iterable)
671{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000672 teeobject *to;
673 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000674
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000675 it = PyObject_GetIter(iterable);
676 if (it == NULL)
677 return NULL;
678 if (PyObject_TypeCheck(it, &tee_type)) {
679 to = (teeobject *)tee_copy((teeobject *)it);
680 goto done;
681 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000682
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000683 to = PyObject_GC_New(teeobject, &tee_type);
684 if (to == NULL)
685 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000686 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000687 if (!to->dataobj) {
688 PyObject_GC_Del(to);
689 to = NULL;
690 goto done;
691 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000692
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000693 to->index = 0;
694 to->weakreflist = NULL;
695 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000696done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000697 Py_XDECREF(it);
698 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000699}
700
701static PyObject *
702tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
703{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000704 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000705
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000706 if (!PyArg_UnpackTuple(args, "_tee", 1, 1, &iterable))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000707 return NULL;
708 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000709}
710
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000711static int
712tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000713{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000714 if (to->weakreflist != NULL)
715 PyObject_ClearWeakRefs((PyObject *) to);
716 Py_CLEAR(to->dataobj);
717 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000718}
719
720static void
721tee_dealloc(teeobject *to)
722{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000723 PyObject_GC_UnTrack(to);
724 tee_clear(to);
725 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000726}
727
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000728static PyObject *
729tee_reduce(teeobject *to)
730{
731 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
732}
733
734static PyObject *
735tee_setstate(teeobject *to, PyObject *state)
736{
737 teedataobject *tdo;
738 int index;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300739 if (!PyTuple_Check(state)) {
740 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000741 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300742 }
743 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
744 return NULL;
745 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000746 if (index < 0 || index > LINKCELLS) {
747 PyErr_SetString(PyExc_ValueError, "Index out of range");
748 return NULL;
749 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200750 Py_INCREF(tdo);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300751 Py_XSETREF(to->dataobj, tdo);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000752 to->index = index;
753 Py_RETURN_NONE;
754}
755
Raymond Hettingerad983e72003-11-12 14:32:26 +0000756PyDoc_STRVAR(teeobject_doc,
757"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000758
Raymond Hettingerad983e72003-11-12 14:32:26 +0000759static PyMethodDef tee_methods[] = {
Raymond Hettingera6a2d442015-08-16 14:38:07 -0700760 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
761 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
762 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000763 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000764};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000765
766static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000767 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000768 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000769 sizeof(teeobject), /* tp_basicsize */
770 0, /* tp_itemsize */
771 /* methods */
772 (destructor)tee_dealloc, /* tp_dealloc */
773 0, /* tp_print */
774 0, /* tp_getattr */
775 0, /* tp_setattr */
776 0, /* tp_reserved */
777 0, /* tp_repr */
778 0, /* tp_as_number */
779 0, /* tp_as_sequence */
780 0, /* tp_as_mapping */
781 0, /* tp_hash */
782 0, /* tp_call */
783 0, /* tp_str */
784 0, /* tp_getattro */
785 0, /* tp_setattro */
786 0, /* tp_as_buffer */
787 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
788 teeobject_doc, /* tp_doc */
789 (traverseproc)tee_traverse, /* tp_traverse */
790 (inquiry)tee_clear, /* tp_clear */
791 0, /* tp_richcompare */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700792 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000793 PyObject_SelfIter, /* tp_iter */
794 (iternextfunc)tee_next, /* tp_iternext */
795 tee_methods, /* tp_methods */
796 0, /* tp_members */
797 0, /* tp_getset */
798 0, /* tp_base */
799 0, /* tp_dict */
800 0, /* tp_descr_get */
801 0, /* tp_descr_set */
802 0, /* tp_dictoffset */
803 0, /* tp_init */
804 0, /* tp_alloc */
805 tee_new, /* tp_new */
806 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000807};
808
Raymond Hettingerad983e72003-11-12 14:32:26 +0000809static PyObject *
810tee(PyObject *self, PyObject *args)
811{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000812 Py_ssize_t i, n=2;
813 PyObject *it, *iterable, *copyable, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200814 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000815
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000816 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
817 return NULL;
818 if (n < 0) {
819 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
820 return NULL;
821 }
822 result = PyTuple_New(n);
823 if (result == NULL)
824 return NULL;
825 if (n == 0)
826 return result;
827 it = PyObject_GetIter(iterable);
828 if (it == NULL) {
829 Py_DECREF(result);
830 return NULL;
831 }
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200832 if (!_PyObject_HasAttrId(it, &PyId___copy__)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000833 copyable = tee_fromiterable(it);
834 Py_DECREF(it);
835 if (copyable == NULL) {
836 Py_DECREF(result);
837 return NULL;
838 }
839 } else
840 copyable = it;
841 PyTuple_SET_ITEM(result, 0, copyable);
842 for (i=1 ; i<n ; i++) {
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200843
844 copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000845 if (copyable == NULL) {
846 Py_DECREF(result);
847 return NULL;
848 }
849 PyTuple_SET_ITEM(result, i, copyable);
850 }
851 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000852}
853
854PyDoc_STRVAR(tee_doc,
855"tee(iterable, n=2) --> tuple of n independent iterators.");
856
857
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700858/* cycle object **************************************************************/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000859
860typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000861 PyObject_HEAD
862 PyObject *it;
863 PyObject *saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700864 Py_ssize_t index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000865 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000866} cycleobject;
867
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000868static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000869
870static PyObject *
871cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
872{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000873 PyObject *it;
874 PyObject *iterable;
875 PyObject *saved;
876 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000877
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +0300878 if (type == &cycle_type && !_PyArg_NoKeywords("cycle", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000879 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000880
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000881 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
882 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000883
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000884 /* Get iterator. */
885 it = PyObject_GetIter(iterable);
886 if (it == NULL)
887 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000888
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000889 saved = PyList_New(0);
890 if (saved == NULL) {
891 Py_DECREF(it);
892 return NULL;
893 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000894
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000895 /* create cycleobject structure */
896 lz = (cycleobject *)type->tp_alloc(type, 0);
897 if (lz == NULL) {
898 Py_DECREF(it);
899 Py_DECREF(saved);
900 return NULL;
901 }
902 lz->it = it;
903 lz->saved = saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700904 lz->index = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000905 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000906
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000907 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000908}
909
910static void
911cycle_dealloc(cycleobject *lz)
912{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000913 PyObject_GC_UnTrack(lz);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000914 Py_XDECREF(lz->it);
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700915 Py_XDECREF(lz->saved);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000916 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000917}
918
919static int
920cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
921{
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700922 if (lz->it)
923 Py_VISIT(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000924 Py_VISIT(lz->saved);
925 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000926}
927
928static PyObject *
929cycle_next(cycleobject *lz)
930{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000931 PyObject *item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000932
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700933 if (lz->it != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000934 item = PyIter_Next(lz->it);
935 if (item != NULL) {
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700936 if (lz->firstpass)
937 return item;
938 if (PyList_Append(lz->saved, item)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000939 Py_DECREF(item);
940 return NULL;
941 }
942 return item;
943 }
Raymond Hettinger98958fe2015-08-15 15:09:30 -0700944 /* Note: StopIteration is already cleared by PyIter_Next() */
945 if (PyErr_Occurred())
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000946 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700947 Py_CLEAR(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000948 }
Serhiy Storchakafff9a312017-03-21 08:53:25 +0200949 if (PyList_GET_SIZE(lz->saved) == 0)
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700950 return NULL;
951 item = PyList_GET_ITEM(lz->saved, lz->index);
952 lz->index++;
Serhiy Storchakafff9a312017-03-21 08:53:25 +0200953 if (lz->index >= PyList_GET_SIZE(lz->saved))
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700954 lz->index = 0;
955 Py_INCREF(item);
956 return item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000957}
958
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000959static PyObject *
960cycle_reduce(cycleobject *lz)
961{
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700962 /* Create a new cycle with the iterator tuple, then set the saved state */
963 if (lz->it == NULL) {
964 PyObject *it = PyObject_GetIter(lz->saved);
965 if (it == NULL)
966 return NULL;
967 if (lz->index != 0) {
968 _Py_IDENTIFIER(__setstate__);
969 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
970 "n", lz->index);
971 if (res == NULL) {
972 Py_DECREF(it);
973 return NULL;
974 }
975 Py_DECREF(res);
976 }
977 return Py_BuildValue("O(N)(Oi)", Py_TYPE(lz), it, lz->saved, 1);
978 }
979 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), lz->it, lz->saved,
980 lz->firstpass);
Raymond Hettingerc39786d2015-08-15 15:16:12 -0700981}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000982
983static PyObject *
984cycle_setstate(cycleobject *lz, PyObject *state)
985{
986 PyObject *saved=NULL;
987 int firstpass;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300988 if (!PyTuple_Check(state)) {
989 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000990 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300991 }
992 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
993 return NULL;
994 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700995 Py_INCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300996 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000997 lz->firstpass = firstpass != 0;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700998 lz->index = 0;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000999 Py_RETURN_NONE;
1000}
1001
1002static PyMethodDef cycle_methods[] = {
1003 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1004 reduce_doc},
1005 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1006 setstate_doc},
1007 {NULL, NULL} /* sentinel */
1008};
1009
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001010PyDoc_STRVAR(cycle_doc,
1011"cycle(iterable) --> cycle object\n\
1012\n\
1013Return elements from the iterable until it is exhausted.\n\
1014Then repeat the sequence indefinitely.");
1015
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001016static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001017 PyVarObject_HEAD_INIT(NULL, 0)
1018 "itertools.cycle", /* tp_name */
1019 sizeof(cycleobject), /* tp_basicsize */
1020 0, /* tp_itemsize */
1021 /* methods */
1022 (destructor)cycle_dealloc, /* tp_dealloc */
1023 0, /* tp_print */
1024 0, /* tp_getattr */
1025 0, /* tp_setattr */
1026 0, /* tp_reserved */
1027 0, /* tp_repr */
1028 0, /* tp_as_number */
1029 0, /* tp_as_sequence */
1030 0, /* tp_as_mapping */
1031 0, /* tp_hash */
1032 0, /* tp_call */
1033 0, /* tp_str */
1034 PyObject_GenericGetAttr, /* tp_getattro */
1035 0, /* tp_setattro */
1036 0, /* tp_as_buffer */
1037 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1038 Py_TPFLAGS_BASETYPE, /* tp_flags */
1039 cycle_doc, /* tp_doc */
1040 (traverseproc)cycle_traverse, /* tp_traverse */
1041 0, /* tp_clear */
1042 0, /* tp_richcompare */
1043 0, /* tp_weaklistoffset */
1044 PyObject_SelfIter, /* tp_iter */
1045 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001046 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001047 0, /* tp_members */
1048 0, /* tp_getset */
1049 0, /* tp_base */
1050 0, /* tp_dict */
1051 0, /* tp_descr_get */
1052 0, /* tp_descr_set */
1053 0, /* tp_dictoffset */
1054 0, /* tp_init */
1055 0, /* tp_alloc */
1056 cycle_new, /* tp_new */
1057 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001058};
1059
1060
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001061/* dropwhile object **********************************************************/
1062
1063typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001064 PyObject_HEAD
1065 PyObject *func;
1066 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001067 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001068} dropwhileobject;
1069
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001070static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001071
1072static PyObject *
1073dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1074{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001075 PyObject *func, *seq;
1076 PyObject *it;
1077 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001078
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001079 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001080 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001081
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001082 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
1083 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001084
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001085 /* Get iterator. */
1086 it = PyObject_GetIter(seq);
1087 if (it == NULL)
1088 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001089
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001090 /* create dropwhileobject structure */
1091 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1092 if (lz == NULL) {
1093 Py_DECREF(it);
1094 return NULL;
1095 }
1096 Py_INCREF(func);
1097 lz->func = func;
1098 lz->it = it;
1099 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001100
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001101 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001102}
1103
1104static void
1105dropwhile_dealloc(dropwhileobject *lz)
1106{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001107 PyObject_GC_UnTrack(lz);
1108 Py_XDECREF(lz->func);
1109 Py_XDECREF(lz->it);
1110 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001111}
1112
1113static int
1114dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1115{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001116 Py_VISIT(lz->it);
1117 Py_VISIT(lz->func);
1118 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001119}
1120
1121static PyObject *
1122dropwhile_next(dropwhileobject *lz)
1123{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001124 PyObject *item, *good;
1125 PyObject *it = lz->it;
1126 long ok;
1127 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001128
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001129 iternext = *Py_TYPE(it)->tp_iternext;
1130 for (;;) {
1131 item = iternext(it);
1132 if (item == NULL)
1133 return NULL;
1134 if (lz->start == 1)
1135 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001136
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01001137 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001138 if (good == NULL) {
1139 Py_DECREF(item);
1140 return NULL;
1141 }
1142 ok = PyObject_IsTrue(good);
1143 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001144 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001145 lz->start = 1;
1146 return item;
1147 }
1148 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001149 if (ok < 0)
1150 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001151 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001152}
1153
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001154static PyObject *
1155dropwhile_reduce(dropwhileobject *lz)
1156{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001157 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 +00001158}
1159
1160static PyObject *
1161dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1162{
1163 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001164 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001165 return NULL;
1166 lz->start = start;
1167 Py_RETURN_NONE;
1168}
1169
1170static PyMethodDef dropwhile_methods[] = {
1171 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1172 reduce_doc},
1173 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1174 setstate_doc},
1175 {NULL, NULL} /* sentinel */
1176};
1177
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001178PyDoc_STRVAR(dropwhile_doc,
1179"dropwhile(predicate, iterable) --> dropwhile object\n\
1180\n\
1181Drop items from the iterable while predicate(item) is true.\n\
1182Afterwards, return every element until the iterable is exhausted.");
1183
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001184static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001185 PyVarObject_HEAD_INIT(NULL, 0)
1186 "itertools.dropwhile", /* tp_name */
1187 sizeof(dropwhileobject), /* tp_basicsize */
1188 0, /* tp_itemsize */
1189 /* methods */
1190 (destructor)dropwhile_dealloc, /* tp_dealloc */
1191 0, /* tp_print */
1192 0, /* tp_getattr */
1193 0, /* tp_setattr */
1194 0, /* tp_reserved */
1195 0, /* tp_repr */
1196 0, /* tp_as_number */
1197 0, /* tp_as_sequence */
1198 0, /* tp_as_mapping */
1199 0, /* tp_hash */
1200 0, /* tp_call */
1201 0, /* tp_str */
1202 PyObject_GenericGetAttr, /* tp_getattro */
1203 0, /* tp_setattro */
1204 0, /* tp_as_buffer */
1205 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1206 Py_TPFLAGS_BASETYPE, /* tp_flags */
1207 dropwhile_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001208 (traverseproc)dropwhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001209 0, /* tp_clear */
1210 0, /* tp_richcompare */
1211 0, /* tp_weaklistoffset */
1212 PyObject_SelfIter, /* tp_iter */
1213 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001214 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001215 0, /* tp_members */
1216 0, /* tp_getset */
1217 0, /* tp_base */
1218 0, /* tp_dict */
1219 0, /* tp_descr_get */
1220 0, /* tp_descr_set */
1221 0, /* tp_dictoffset */
1222 0, /* tp_init */
1223 0, /* tp_alloc */
1224 dropwhile_new, /* tp_new */
1225 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001226};
1227
1228
1229/* takewhile object **********************************************************/
1230
1231typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001232 PyObject_HEAD
1233 PyObject *func;
1234 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001235 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001236} takewhileobject;
1237
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001238static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001239
1240static PyObject *
1241takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1242{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001243 PyObject *func, *seq;
1244 PyObject *it;
1245 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001246
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001247 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001248 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001249
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001250 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1251 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001252
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001253 /* Get iterator. */
1254 it = PyObject_GetIter(seq);
1255 if (it == NULL)
1256 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001257
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001258 /* create takewhileobject structure */
1259 lz = (takewhileobject *)type->tp_alloc(type, 0);
1260 if (lz == NULL) {
1261 Py_DECREF(it);
1262 return NULL;
1263 }
1264 Py_INCREF(func);
1265 lz->func = func;
1266 lz->it = it;
1267 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001268
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001269 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001270}
1271
1272static void
1273takewhile_dealloc(takewhileobject *lz)
1274{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001275 PyObject_GC_UnTrack(lz);
1276 Py_XDECREF(lz->func);
1277 Py_XDECREF(lz->it);
1278 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001279}
1280
1281static int
1282takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1283{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001284 Py_VISIT(lz->it);
1285 Py_VISIT(lz->func);
1286 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001287}
1288
1289static PyObject *
1290takewhile_next(takewhileobject *lz)
1291{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001292 PyObject *item, *good;
1293 PyObject *it = lz->it;
1294 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001295
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001296 if (lz->stop == 1)
1297 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001298
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001299 item = (*Py_TYPE(it)->tp_iternext)(it);
1300 if (item == NULL)
1301 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001302
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01001303 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001304 if (good == NULL) {
1305 Py_DECREF(item);
1306 return NULL;
1307 }
1308 ok = PyObject_IsTrue(good);
1309 Py_DECREF(good);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001310 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001311 return item;
1312 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001313 if (ok == 0)
1314 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001315 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001316}
1317
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001318static PyObject *
1319takewhile_reduce(takewhileobject *lz)
1320{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001321 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 +00001322}
1323
1324static PyObject *
1325takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1326{
1327 int stop = PyObject_IsTrue(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001328
Antoine Pitrou721738f2012-08-15 23:20:39 +02001329 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001330 return NULL;
1331 lz->stop = stop;
1332 Py_RETURN_NONE;
1333}
1334
1335static PyMethodDef takewhile_reduce_methods[] = {
1336 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1337 reduce_doc},
1338 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1339 setstate_doc},
1340 {NULL, NULL} /* sentinel */
1341};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001342PyDoc_STRVAR(takewhile_doc,
1343"takewhile(predicate, iterable) --> takewhile object\n\
1344\n\
1345Return successive entries from an iterable as long as the \n\
1346predicate evaluates to true for each entry.");
1347
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001348static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001349 PyVarObject_HEAD_INIT(NULL, 0)
1350 "itertools.takewhile", /* tp_name */
1351 sizeof(takewhileobject), /* tp_basicsize */
1352 0, /* tp_itemsize */
1353 /* methods */
1354 (destructor)takewhile_dealloc, /* tp_dealloc */
1355 0, /* tp_print */
1356 0, /* tp_getattr */
1357 0, /* tp_setattr */
1358 0, /* tp_reserved */
1359 0, /* tp_repr */
1360 0, /* tp_as_number */
1361 0, /* tp_as_sequence */
1362 0, /* tp_as_mapping */
1363 0, /* tp_hash */
1364 0, /* tp_call */
1365 0, /* tp_str */
1366 PyObject_GenericGetAttr, /* tp_getattro */
1367 0, /* tp_setattro */
1368 0, /* tp_as_buffer */
1369 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1370 Py_TPFLAGS_BASETYPE, /* tp_flags */
1371 takewhile_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001372 (traverseproc)takewhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001373 0, /* tp_clear */
1374 0, /* tp_richcompare */
1375 0, /* tp_weaklistoffset */
1376 PyObject_SelfIter, /* tp_iter */
1377 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001378 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001379 0, /* tp_members */
1380 0, /* tp_getset */
1381 0, /* tp_base */
1382 0, /* tp_dict */
1383 0, /* tp_descr_get */
1384 0, /* tp_descr_set */
1385 0, /* tp_dictoffset */
1386 0, /* tp_init */
1387 0, /* tp_alloc */
1388 takewhile_new, /* tp_new */
1389 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001390};
1391
1392
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001393/* islice object *************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001394
1395typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001396 PyObject_HEAD
1397 PyObject *it;
1398 Py_ssize_t next;
1399 Py_ssize_t stop;
1400 Py_ssize_t step;
1401 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001402} isliceobject;
1403
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001404static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001405
1406static PyObject *
1407islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1408{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001409 PyObject *seq;
1410 Py_ssize_t start=0, stop=-1, step=1;
1411 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1412 Py_ssize_t numargs;
1413 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001414
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001415 if (type == &islice_type && !_PyArg_NoKeywords("islice", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001416 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001417
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001418 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1419 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001420
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001421 numargs = PyTuple_Size(args);
1422 if (numargs == 2) {
1423 if (a1 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001424 stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001425 if (stop == -1) {
1426 if (PyErr_Occurred())
1427 PyErr_Clear();
1428 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001429 "Stop argument for islice() must be None or "
1430 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001431 return NULL;
1432 }
1433 }
1434 } else {
1435 if (a1 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001436 start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001437 if (start == -1 && PyErr_Occurred())
1438 PyErr_Clear();
1439 if (a2 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001440 stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001441 if (stop == -1) {
1442 if (PyErr_Occurred())
1443 PyErr_Clear();
1444 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001445 "Stop argument for islice() must be None or "
1446 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001447 return NULL;
1448 }
1449 }
1450 }
1451 if (start<0 || stop<-1) {
1452 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001453 "Indices for islice() must be None or "
1454 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001455 return NULL;
1456 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001458 if (a3 != NULL) {
1459 if (a3 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001460 step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001461 if (step == -1 && PyErr_Occurred())
1462 PyErr_Clear();
1463 }
1464 if (step<1) {
1465 PyErr_SetString(PyExc_ValueError,
1466 "Step for islice() must be a positive integer or None.");
1467 return NULL;
1468 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001470 /* Get iterator. */
1471 it = PyObject_GetIter(seq);
1472 if (it == NULL)
1473 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001474
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001475 /* create isliceobject structure */
1476 lz = (isliceobject *)type->tp_alloc(type, 0);
1477 if (lz == NULL) {
1478 Py_DECREF(it);
1479 return NULL;
1480 }
1481 lz->it = it;
1482 lz->next = start;
1483 lz->stop = stop;
1484 lz->step = step;
1485 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001486
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001487 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001488}
1489
1490static void
1491islice_dealloc(isliceobject *lz)
1492{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001493 PyObject_GC_UnTrack(lz);
1494 Py_XDECREF(lz->it);
1495 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001496}
1497
1498static int
1499islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1500{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001501 Py_VISIT(lz->it);
1502 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001503}
1504
1505static PyObject *
1506islice_next(isliceobject *lz)
1507{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001508 PyObject *item;
1509 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001510 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001511 Py_ssize_t oldnext;
1512 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001513
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001514 if (it == NULL)
1515 return NULL;
1516
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001517 iternext = *Py_TYPE(it)->tp_iternext;
1518 while (lz->cnt < lz->next) {
1519 item = iternext(it);
1520 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001521 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001522 Py_DECREF(item);
1523 lz->cnt++;
1524 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001525 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001526 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001527 item = iternext(it);
1528 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001529 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001530 lz->cnt++;
1531 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001532 /* The (size_t) cast below avoids the danger of undefined
1533 behaviour from signed integer overflow. */
1534 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001535 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1536 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001537 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001538
1539empty:
1540 Py_CLEAR(lz->it);
1541 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001542}
1543
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001544static PyObject *
1545islice_reduce(isliceobject *lz)
1546{
1547 /* When unpickled, generate a new object with the same bounds,
1548 * then 'setstate' with the next and count
1549 */
1550 PyObject *stop;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001551
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001552 if (lz->it == NULL) {
1553 PyObject *empty_list;
1554 PyObject *empty_it;
1555 empty_list = PyList_New(0);
1556 if (empty_list == NULL)
1557 return NULL;
1558 empty_it = PyObject_GetIter(empty_list);
1559 Py_DECREF(empty_list);
1560 if (empty_it == NULL)
1561 return NULL;
1562 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1563 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001564 if (lz->stop == -1) {
1565 stop = Py_None;
1566 Py_INCREF(stop);
1567 } else {
1568 stop = PyLong_FromSsize_t(lz->stop);
1569 if (stop == NULL)
1570 return NULL;
1571 }
1572 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1573 lz->it, lz->next, stop, lz->step,
1574 lz->cnt);
1575}
1576
1577static PyObject *
1578islice_setstate(isliceobject *lz, PyObject *state)
1579{
1580 Py_ssize_t cnt = PyLong_AsSsize_t(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001581
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001582 if (cnt == -1 && PyErr_Occurred())
1583 return NULL;
1584 lz->cnt = cnt;
1585 Py_RETURN_NONE;
1586}
1587
1588static PyMethodDef islice_methods[] = {
1589 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1590 reduce_doc},
1591 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1592 setstate_doc},
1593 {NULL, NULL} /* sentinel */
1594};
1595
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001596PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001597"islice(iterable, stop) --> islice object\n\
1598islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001599\n\
1600Return an iterator whose next() method returns selected values from an\n\
1601iterable. If start is specified, will skip all preceding elements;\n\
1602otherwise, start defaults to zero. Step defaults to one. If\n\
1603specified as another value, step determines how many values are \n\
1604skipped between successive calls. Works like a slice() on a list\n\
1605but returns an iterator.");
1606
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001607static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001608 PyVarObject_HEAD_INIT(NULL, 0)
1609 "itertools.islice", /* tp_name */
1610 sizeof(isliceobject), /* tp_basicsize */
1611 0, /* tp_itemsize */
1612 /* methods */
1613 (destructor)islice_dealloc, /* tp_dealloc */
1614 0, /* tp_print */
1615 0, /* tp_getattr */
1616 0, /* tp_setattr */
1617 0, /* tp_reserved */
1618 0, /* tp_repr */
1619 0, /* tp_as_number */
1620 0, /* tp_as_sequence */
1621 0, /* tp_as_mapping */
1622 0, /* tp_hash */
1623 0, /* tp_call */
1624 0, /* tp_str */
1625 PyObject_GenericGetAttr, /* tp_getattro */
1626 0, /* tp_setattro */
1627 0, /* tp_as_buffer */
1628 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1629 Py_TPFLAGS_BASETYPE, /* tp_flags */
1630 islice_doc, /* tp_doc */
1631 (traverseproc)islice_traverse, /* tp_traverse */
1632 0, /* tp_clear */
1633 0, /* tp_richcompare */
1634 0, /* tp_weaklistoffset */
1635 PyObject_SelfIter, /* tp_iter */
1636 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001637 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001638 0, /* tp_members */
1639 0, /* tp_getset */
1640 0, /* tp_base */
1641 0, /* tp_dict */
1642 0, /* tp_descr_get */
1643 0, /* tp_descr_set */
1644 0, /* tp_dictoffset */
1645 0, /* tp_init */
1646 0, /* tp_alloc */
1647 islice_new, /* tp_new */
1648 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001649};
1650
1651
1652/* starmap object ************************************************************/
1653
1654typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001655 PyObject_HEAD
1656 PyObject *func;
1657 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001658} starmapobject;
1659
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001660static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001661
1662static PyObject *
1663starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1664{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001665 PyObject *func, *seq;
1666 PyObject *it;
1667 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001668
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001669 if (type == &starmap_type && !_PyArg_NoKeywords("starmap", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001670 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001671
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001672 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1673 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001674
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001675 /* Get iterator. */
1676 it = PyObject_GetIter(seq);
1677 if (it == NULL)
1678 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001679
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001680 /* create starmapobject structure */
1681 lz = (starmapobject *)type->tp_alloc(type, 0);
1682 if (lz == NULL) {
1683 Py_DECREF(it);
1684 return NULL;
1685 }
1686 Py_INCREF(func);
1687 lz->func = func;
1688 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001689
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001690 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001691}
1692
1693static void
1694starmap_dealloc(starmapobject *lz)
1695{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001696 PyObject_GC_UnTrack(lz);
1697 Py_XDECREF(lz->func);
1698 Py_XDECREF(lz->it);
1699 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001700}
1701
1702static int
1703starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1704{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001705 Py_VISIT(lz->it);
1706 Py_VISIT(lz->func);
1707 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001708}
1709
1710static PyObject *
1711starmap_next(starmapobject *lz)
1712{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001713 PyObject *args;
1714 PyObject *result;
1715 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001716
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001717 args = (*Py_TYPE(it)->tp_iternext)(it);
1718 if (args == NULL)
1719 return NULL;
1720 if (!PyTuple_CheckExact(args)) {
1721 PyObject *newargs = PySequence_Tuple(args);
1722 Py_DECREF(args);
1723 if (newargs == NULL)
1724 return NULL;
1725 args = newargs;
1726 }
1727 result = PyObject_Call(lz->func, args, NULL);
1728 Py_DECREF(args);
1729 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001730}
1731
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001732static PyObject *
1733starmap_reduce(starmapobject *lz)
1734{
1735 /* Just pickle the iterator */
1736 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1737}
1738
1739static PyMethodDef starmap_methods[] = {
1740 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1741 reduce_doc},
1742 {NULL, NULL} /* sentinel */
1743};
1744
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001745PyDoc_STRVAR(starmap_doc,
1746"starmap(function, sequence) --> starmap object\n\
1747\n\
1748Return an iterator whose values are returned from the function evaluated\n\
Serhiy Storchakad65c9492015-11-02 14:10:23 +02001749with an argument tuple taken from the given sequence.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001750
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001751static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001752 PyVarObject_HEAD_INIT(NULL, 0)
1753 "itertools.starmap", /* tp_name */
1754 sizeof(starmapobject), /* tp_basicsize */
1755 0, /* tp_itemsize */
1756 /* methods */
1757 (destructor)starmap_dealloc, /* tp_dealloc */
1758 0, /* tp_print */
1759 0, /* tp_getattr */
1760 0, /* tp_setattr */
1761 0, /* tp_reserved */
1762 0, /* tp_repr */
1763 0, /* tp_as_number */
1764 0, /* tp_as_sequence */
1765 0, /* tp_as_mapping */
1766 0, /* tp_hash */
1767 0, /* tp_call */
1768 0, /* tp_str */
1769 PyObject_GenericGetAttr, /* tp_getattro */
1770 0, /* tp_setattro */
1771 0, /* tp_as_buffer */
1772 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1773 Py_TPFLAGS_BASETYPE, /* tp_flags */
1774 starmap_doc, /* tp_doc */
1775 (traverseproc)starmap_traverse, /* tp_traverse */
1776 0, /* tp_clear */
1777 0, /* tp_richcompare */
1778 0, /* tp_weaklistoffset */
1779 PyObject_SelfIter, /* tp_iter */
1780 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001781 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001782 0, /* tp_members */
1783 0, /* tp_getset */
1784 0, /* tp_base */
1785 0, /* tp_dict */
1786 0, /* tp_descr_get */
1787 0, /* tp_descr_set */
1788 0, /* tp_dictoffset */
1789 0, /* tp_init */
1790 0, /* tp_alloc */
1791 starmap_new, /* tp_new */
1792 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001793};
1794
1795
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001796/* chain object **************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001797
1798typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001799 PyObject_HEAD
1800 PyObject *source; /* Iterator over input iterables */
1801 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001802} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001803
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001804static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001805
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001806static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001807chain_new_internal(PyTypeObject *type, PyObject *source)
1808{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001809 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001810
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001811 lz = (chainobject *)type->tp_alloc(type, 0);
1812 if (lz == NULL) {
1813 Py_DECREF(source);
1814 return NULL;
1815 }
1816
1817 lz->source = source;
1818 lz->active = NULL;
1819 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001820}
1821
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001822static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001823chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001824{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001825 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001826
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001827 if (type == &chain_type && !_PyArg_NoKeywords("chain", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001828 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001829
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001830 source = PyObject_GetIter(args);
1831 if (source == NULL)
1832 return NULL;
1833
1834 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001835}
1836
1837static PyObject *
1838chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1839{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001840 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001841
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001842 source = PyObject_GetIter(arg);
1843 if (source == NULL)
1844 return NULL;
1845
1846 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001847}
1848
1849static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001850chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001851{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001852 PyObject_GC_UnTrack(lz);
1853 Py_XDECREF(lz->active);
1854 Py_XDECREF(lz->source);
1855 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001856}
1857
Raymond Hettinger2012f172003-02-07 05:32:58 +00001858static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001859chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001860{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001861 Py_VISIT(lz->source);
1862 Py_VISIT(lz->active);
1863 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001864}
1865
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001866static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001867chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001868{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001869 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001870
T. Wouters5466d4a2017-03-30 09:58:35 -07001871 /* lz->source is the iterator of iterables. If it's NULL, we've already
1872 * consumed them all. lz->active is the current iterator. If it's NULL,
1873 * we should grab a new one from lz->source. */
1874 while (lz->source != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001875 if (lz->active == NULL) {
T. Wouters5466d4a2017-03-30 09:58:35 -07001876 PyObject *iterable = PyIter_Next(lz->source);
1877 if (iterable == NULL) {
1878 Py_CLEAR(lz->source);
1879 return NULL; /* no more input sources */
1880 }
1881 lz->active = PyObject_GetIter(iterable);
1882 Py_DECREF(iterable);
1883 if (lz->active == NULL) {
1884 Py_CLEAR(lz->source);
1885 return NULL; /* input not iterable */
1886 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001887 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001888 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1889 if (item != NULL)
1890 return item;
1891 if (PyErr_Occurred()) {
1892 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1893 PyErr_Clear();
1894 else
1895 return NULL; /* input raised an exception */
1896 }
1897 /* lz->active is consumed, try with the next iterable. */
1898 Py_CLEAR(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001899 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001900 /* Everything had been consumed already. */
1901 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001902}
1903
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001904static PyObject *
1905chain_reduce(chainobject *lz)
1906{
1907 if (lz->source) {
1908 /* we can't pickle function objects (itertools.from_iterable) so
1909 * we must use setstate to replace the iterable. One day we
1910 * will fix pickling of functions
1911 */
1912 if (lz->active) {
1913 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1914 } else {
1915 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1916 }
1917 } else {
1918 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1919 }
1920 return NULL;
1921}
1922
1923static PyObject *
1924chain_setstate(chainobject *lz, PyObject *state)
1925{
1926 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001927
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001928 if (!PyTuple_Check(state)) {
1929 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001930 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001931 }
1932 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
1933 return NULL;
1934 }
1935 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
1936 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
1937 return NULL;
1938 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001939
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02001940 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001941 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02001942 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001943 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001944 Py_RETURN_NONE;
1945}
1946
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001947PyDoc_STRVAR(chain_doc,
1948"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001949\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001950Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001951first iterable until it is exhausted, then elements from the next\n\
1952iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001953
Christian Heimesf16baeb2008-02-29 14:57:44 +00001954PyDoc_STRVAR(chain_from_iterable_doc,
1955"chain.from_iterable(iterable) --> chain object\n\
1956\n\
Martin Pantereb995702016-07-28 01:11:04 +00001957Alternate chain() constructor taking a single iterable argument\n\
Christian Heimesf16baeb2008-02-29 14:57:44 +00001958that evaluates lazily.");
1959
1960static PyMethodDef chain_methods[] = {
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001961 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1962 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001963 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1964 reduce_doc},
1965 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1966 setstate_doc},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001967 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001968};
1969
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001970static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001971 PyVarObject_HEAD_INIT(NULL, 0)
1972 "itertools.chain", /* tp_name */
1973 sizeof(chainobject), /* tp_basicsize */
1974 0, /* tp_itemsize */
1975 /* methods */
1976 (destructor)chain_dealloc, /* tp_dealloc */
1977 0, /* tp_print */
1978 0, /* tp_getattr */
1979 0, /* tp_setattr */
1980 0, /* tp_reserved */
1981 0, /* tp_repr */
1982 0, /* tp_as_number */
1983 0, /* tp_as_sequence */
1984 0, /* tp_as_mapping */
1985 0, /* tp_hash */
1986 0, /* tp_call */
1987 0, /* tp_str */
1988 PyObject_GenericGetAttr, /* tp_getattro */
1989 0, /* tp_setattro */
1990 0, /* tp_as_buffer */
1991 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1992 Py_TPFLAGS_BASETYPE, /* tp_flags */
1993 chain_doc, /* tp_doc */
1994 (traverseproc)chain_traverse, /* tp_traverse */
1995 0, /* tp_clear */
1996 0, /* tp_richcompare */
1997 0, /* tp_weaklistoffset */
1998 PyObject_SelfIter, /* tp_iter */
1999 (iternextfunc)chain_next, /* tp_iternext */
2000 chain_methods, /* tp_methods */
2001 0, /* tp_members */
2002 0, /* tp_getset */
2003 0, /* tp_base */
2004 0, /* tp_dict */
2005 0, /* tp_descr_get */
2006 0, /* tp_descr_set */
2007 0, /* tp_dictoffset */
2008 0, /* tp_init */
2009 0, /* tp_alloc */
2010 chain_new, /* tp_new */
2011 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002012};
2013
2014
Christian Heimesc3f30c42008-02-22 16:37:40 +00002015/* product object ************************************************************/
2016
2017typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002018 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002019 PyObject *pools; /* tuple of pool tuples */
2020 Py_ssize_t *indices; /* one index per pool */
2021 PyObject *result; /* most recently returned result tuple */
2022 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002023} productobject;
2024
2025static PyTypeObject product_type;
2026
2027static PyObject *
2028product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2029{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002030 productobject *lz;
2031 Py_ssize_t nargs, npools, repeat=1;
2032 PyObject *pools = NULL;
2033 Py_ssize_t *indices = NULL;
2034 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002035
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002036 if (kwds != NULL) {
2037 char *kwlist[] = {"repeat", 0};
2038 PyObject *tmpargs = PyTuple_New(0);
2039 if (tmpargs == NULL)
2040 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002041 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2042 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002043 Py_DECREF(tmpargs);
2044 return NULL;
2045 }
2046 Py_DECREF(tmpargs);
2047 if (repeat < 0) {
2048 PyErr_SetString(PyExc_ValueError,
2049 "repeat argument cannot be negative");
2050 return NULL;
2051 }
2052 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002053
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002054 assert(PyTuple_CheckExact(args));
2055 if (repeat == 0) {
2056 nargs = 0;
2057 } else {
2058 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002059 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002060 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2061 return NULL;
2062 }
2063 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002064 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002065
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002066 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002067 if (indices == NULL) {
2068 PyErr_NoMemory();
2069 goto error;
2070 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002071
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002072 pools = PyTuple_New(npools);
2073 if (pools == NULL)
2074 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002075
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002076 for (i=0; i < nargs ; ++i) {
2077 PyObject *item = PyTuple_GET_ITEM(args, i);
2078 PyObject *pool = PySequence_Tuple(item);
2079 if (pool == NULL)
2080 goto error;
2081 PyTuple_SET_ITEM(pools, i, pool);
2082 indices[i] = 0;
2083 }
2084 for ( ; i < npools; ++i) {
2085 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2086 Py_INCREF(pool);
2087 PyTuple_SET_ITEM(pools, i, pool);
2088 indices[i] = 0;
2089 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002090
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002091 /* create productobject structure */
2092 lz = (productobject *)type->tp_alloc(type, 0);
2093 if (lz == NULL)
2094 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002095
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002096 lz->pools = pools;
2097 lz->indices = indices;
2098 lz->result = NULL;
2099 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002100
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002101 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002102
2103error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002104 if (indices != NULL)
2105 PyMem_Free(indices);
2106 Py_XDECREF(pools);
2107 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002108}
2109
2110static void
2111product_dealloc(productobject *lz)
2112{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002113 PyObject_GC_UnTrack(lz);
2114 Py_XDECREF(lz->pools);
2115 Py_XDECREF(lz->result);
2116 if (lz->indices != NULL)
2117 PyMem_Free(lz->indices);
2118 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002119}
2120
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002121static PyObject *
2122product_sizeof(productobject *lz, void *unused)
2123{
2124 Py_ssize_t res;
2125
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002126 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002127 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2128 return PyLong_FromSsize_t(res);
2129}
2130
2131PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2132
Christian Heimesc3f30c42008-02-22 16:37:40 +00002133static int
2134product_traverse(productobject *lz, visitproc visit, void *arg)
2135{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002136 Py_VISIT(lz->pools);
2137 Py_VISIT(lz->result);
2138 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002139}
2140
2141static PyObject *
2142product_next(productobject *lz)
2143{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002144 PyObject *pool;
2145 PyObject *elem;
2146 PyObject *oldelem;
2147 PyObject *pools = lz->pools;
2148 PyObject *result = lz->result;
2149 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2150 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002151
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002152 if (lz->stopped)
2153 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002154
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002155 if (result == NULL) {
2156 /* On the first pass, return an initial tuple filled with the
2157 first element from each pool. */
2158 result = PyTuple_New(npools);
2159 if (result == NULL)
2160 goto empty;
2161 lz->result = result;
2162 for (i=0; i < npools; i++) {
2163 pool = PyTuple_GET_ITEM(pools, i);
2164 if (PyTuple_GET_SIZE(pool) == 0)
2165 goto empty;
2166 elem = PyTuple_GET_ITEM(pool, 0);
2167 Py_INCREF(elem);
2168 PyTuple_SET_ITEM(result, i, elem);
2169 }
2170 } else {
2171 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002172
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002173 /* Copy the previous result tuple or re-use it if available */
2174 if (Py_REFCNT(result) > 1) {
2175 PyObject *old_result = result;
2176 result = PyTuple_New(npools);
2177 if (result == NULL)
2178 goto empty;
2179 lz->result = result;
2180 for (i=0; i < npools; i++) {
2181 elem = PyTuple_GET_ITEM(old_result, i);
2182 Py_INCREF(elem);
2183 PyTuple_SET_ITEM(result, i, elem);
2184 }
2185 Py_DECREF(old_result);
2186 }
2187 /* Now, we've got the only copy so we can update it in-place */
2188 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002189
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002190 /* Update the pool indices right-to-left. Only advance to the
2191 next pool when the previous one rolls-over */
2192 for (i=npools-1 ; i >= 0 ; i--) {
2193 pool = PyTuple_GET_ITEM(pools, i);
2194 indices[i]++;
2195 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2196 /* Roll-over and advance to next pool */
2197 indices[i] = 0;
2198 elem = PyTuple_GET_ITEM(pool, 0);
2199 Py_INCREF(elem);
2200 oldelem = PyTuple_GET_ITEM(result, i);
2201 PyTuple_SET_ITEM(result, i, elem);
2202 Py_DECREF(oldelem);
2203 } else {
2204 /* No rollover. Just increment and stop here. */
2205 elem = PyTuple_GET_ITEM(pool, indices[i]);
2206 Py_INCREF(elem);
2207 oldelem = PyTuple_GET_ITEM(result, i);
2208 PyTuple_SET_ITEM(result, i, elem);
2209 Py_DECREF(oldelem);
2210 break;
2211 }
2212 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002213
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002214 /* If i is negative, then the indices have all rolled-over
2215 and we're done. */
2216 if (i < 0)
2217 goto empty;
2218 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002219
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002220 Py_INCREF(result);
2221 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002222
2223empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002224 lz->stopped = 1;
2225 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002226}
2227
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002228static PyObject *
2229product_reduce(productobject *lz)
2230{
2231 if (lz->stopped) {
2232 return Py_BuildValue("O(())", Py_TYPE(lz));
2233 } else if (lz->result == NULL) {
2234 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2235 } else {
2236 PyObject *indices;
2237 Py_ssize_t n, i;
2238
2239 /* we must pickle the indices use them for setstate, and
2240 * additionally indicate that the iterator has started
2241 */
2242 n = PyTuple_GET_SIZE(lz->pools);
2243 indices = PyTuple_New(n);
2244 if (indices == NULL)
2245 return NULL;
2246 for (i=0; i<n; i++){
2247 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2248 if (!index) {
2249 Py_DECREF(indices);
2250 return NULL;
2251 }
2252 PyTuple_SET_ITEM(indices, i, index);
2253 }
2254 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2255 }
2256}
2257
2258static PyObject *
2259product_setstate(productobject *lz, PyObject *state)
2260{
2261 PyObject *result;
2262 Py_ssize_t n, i;
2263
2264 n = PyTuple_GET_SIZE(lz->pools);
2265 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2266 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2267 return NULL;
2268 }
2269 for (i=0; i<n; i++)
2270 {
2271 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2272 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002273 PyObject* pool;
2274 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002275 if (index < 0 && PyErr_Occurred())
2276 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002277 pool = PyTuple_GET_ITEM(lz->pools, i);
2278 poolsize = PyTuple_GET_SIZE(pool);
2279 if (poolsize == 0) {
2280 lz->stopped = 1;
2281 Py_RETURN_NONE;
2282 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002283 /* clamp the index */
2284 if (index < 0)
2285 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002286 else if (index > poolsize-1)
2287 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002288 lz->indices[i] = index;
2289 }
2290
2291 result = PyTuple_New(n);
2292 if (!result)
2293 return NULL;
2294 for (i=0; i<n; i++) {
2295 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2296 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2297 Py_INCREF(element);
2298 PyTuple_SET_ITEM(result, i, element);
2299 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002300 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002301 Py_RETURN_NONE;
2302}
2303
2304static PyMethodDef product_methods[] = {
2305 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2306 reduce_doc},
2307 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2308 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002309 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2310 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002311 {NULL, NULL} /* sentinel */
2312};
2313
Christian Heimesc3f30c42008-02-22 16:37:40 +00002314PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002315"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002316\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002317Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002318For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2319The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2320cycle in a manner similar to an odometer (with the rightmost element changing\n\
2321on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002322To compute the product of an iterable with itself, specify the number\n\
2323of repetitions with the optional repeat keyword argument. For example,\n\
2324product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002325product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2326product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2327
2328static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002329 PyVarObject_HEAD_INIT(NULL, 0)
2330 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002331 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002332 0, /* tp_itemsize */
2333 /* methods */
2334 (destructor)product_dealloc, /* tp_dealloc */
2335 0, /* tp_print */
2336 0, /* tp_getattr */
2337 0, /* tp_setattr */
2338 0, /* tp_reserved */
2339 0, /* tp_repr */
2340 0, /* tp_as_number */
2341 0, /* tp_as_sequence */
2342 0, /* tp_as_mapping */
2343 0, /* tp_hash */
2344 0, /* tp_call */
2345 0, /* tp_str */
2346 PyObject_GenericGetAttr, /* tp_getattro */
2347 0, /* tp_setattro */
2348 0, /* tp_as_buffer */
2349 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2350 Py_TPFLAGS_BASETYPE, /* tp_flags */
2351 product_doc, /* tp_doc */
2352 (traverseproc)product_traverse, /* tp_traverse */
2353 0, /* tp_clear */
2354 0, /* tp_richcompare */
2355 0, /* tp_weaklistoffset */
2356 PyObject_SelfIter, /* tp_iter */
2357 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002358 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002359 0, /* tp_members */
2360 0, /* tp_getset */
2361 0, /* tp_base */
2362 0, /* tp_dict */
2363 0, /* tp_descr_get */
2364 0, /* tp_descr_set */
2365 0, /* tp_dictoffset */
2366 0, /* tp_init */
2367 0, /* tp_alloc */
2368 product_new, /* tp_new */
2369 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002370};
2371
2372
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002373/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002374
2375typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002376 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002377 PyObject *pool; /* input converted to a tuple */
2378 Py_ssize_t *indices; /* one index per result element */
2379 PyObject *result; /* most recently returned result tuple */
2380 Py_ssize_t r; /* size of result tuple */
2381 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002382} combinationsobject;
2383
2384static PyTypeObject combinations_type;
2385
2386static PyObject *
2387combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2388{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002389 combinationsobject *co;
2390 Py_ssize_t n;
2391 Py_ssize_t r;
2392 PyObject *pool = NULL;
2393 PyObject *iterable = NULL;
2394 Py_ssize_t *indices = NULL;
2395 Py_ssize_t i;
2396 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002397
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002398 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2399 &iterable, &r))
2400 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002401
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002402 pool = PySequence_Tuple(iterable);
2403 if (pool == NULL)
2404 goto error;
2405 n = PyTuple_GET_SIZE(pool);
2406 if (r < 0) {
2407 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2408 goto error;
2409 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002410
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002411 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002412 if (indices == NULL) {
2413 PyErr_NoMemory();
2414 goto error;
2415 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002416
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002417 for (i=0 ; i<r ; i++)
2418 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002419
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002420 /* create combinationsobject structure */
2421 co = (combinationsobject *)type->tp_alloc(type, 0);
2422 if (co == NULL)
2423 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002424
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002425 co->pool = pool;
2426 co->indices = indices;
2427 co->result = NULL;
2428 co->r = r;
2429 co->stopped = r > n ? 1 : 0;
2430
2431 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002432
2433error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002434 if (indices != NULL)
2435 PyMem_Free(indices);
2436 Py_XDECREF(pool);
2437 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002438}
2439
2440static void
2441combinations_dealloc(combinationsobject *co)
2442{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002443 PyObject_GC_UnTrack(co);
2444 Py_XDECREF(co->pool);
2445 Py_XDECREF(co->result);
2446 if (co->indices != NULL)
2447 PyMem_Free(co->indices);
2448 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002449}
2450
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002451static PyObject *
2452combinations_sizeof(combinationsobject *co, void *unused)
2453{
2454 Py_ssize_t res;
2455
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002456 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002457 res += co->r * sizeof(Py_ssize_t);
2458 return PyLong_FromSsize_t(res);
2459}
2460
Christian Heimes380f7f22008-02-28 11:19:05 +00002461static int
2462combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2463{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002464 Py_VISIT(co->pool);
2465 Py_VISIT(co->result);
2466 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002467}
2468
2469static PyObject *
2470combinations_next(combinationsobject *co)
2471{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002472 PyObject *elem;
2473 PyObject *oldelem;
2474 PyObject *pool = co->pool;
2475 Py_ssize_t *indices = co->indices;
2476 PyObject *result = co->result;
2477 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2478 Py_ssize_t r = co->r;
2479 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002480
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002481 if (co->stopped)
2482 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002483
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002484 if (result == NULL) {
2485 /* On the first pass, initialize result tuple using the indices */
2486 result = PyTuple_New(r);
2487 if (result == NULL)
2488 goto empty;
2489 co->result = result;
2490 for (i=0; i<r ; i++) {
2491 index = indices[i];
2492 elem = PyTuple_GET_ITEM(pool, index);
2493 Py_INCREF(elem);
2494 PyTuple_SET_ITEM(result, i, elem);
2495 }
2496 } else {
2497 /* Copy the previous result tuple or re-use it if available */
2498 if (Py_REFCNT(result) > 1) {
2499 PyObject *old_result = result;
2500 result = PyTuple_New(r);
2501 if (result == NULL)
2502 goto empty;
2503 co->result = result;
2504 for (i=0; i<r ; i++) {
2505 elem = PyTuple_GET_ITEM(old_result, i);
2506 Py_INCREF(elem);
2507 PyTuple_SET_ITEM(result, i, elem);
2508 }
2509 Py_DECREF(old_result);
2510 }
2511 /* Now, we've got the only copy so we can update it in-place
2512 * CPython's empty tuple is a singleton and cached in
2513 * PyTuple's freelist.
2514 */
2515 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002516
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002517 /* Scan indices right-to-left until finding one that is not
2518 at its maximum (i + n - r). */
2519 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2520 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002521
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002522 /* If i is negative, then the indices are all at
2523 their maximum value and we're done. */
2524 if (i < 0)
2525 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002526
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002527 /* Increment the current index which we know is not at its
2528 maximum. Then move back to the right setting each index
2529 to its lowest possible value (one higher than the index
2530 to its left -- this maintains the sort order invariant). */
2531 indices[i]++;
2532 for (j=i+1 ; j<r ; j++)
2533 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002534
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002535 /* Update the result tuple for the new indices
2536 starting with i, the leftmost index that changed */
2537 for ( ; i<r ; i++) {
2538 index = indices[i];
2539 elem = PyTuple_GET_ITEM(pool, index);
2540 Py_INCREF(elem);
2541 oldelem = PyTuple_GET_ITEM(result, i);
2542 PyTuple_SET_ITEM(result, i, elem);
2543 Py_DECREF(oldelem);
2544 }
2545 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002546
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002547 Py_INCREF(result);
2548 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002549
2550empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002551 co->stopped = 1;
2552 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002553}
2554
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002555static PyObject *
2556combinations_reduce(combinationsobject *lz)
2557{
2558 if (lz->result == NULL) {
2559 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2560 } else if (lz->stopped) {
2561 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2562 } else {
2563 PyObject *indices;
2564 Py_ssize_t i;
2565
2566 /* we must pickle the indices and use them for setstate */
2567 indices = PyTuple_New(lz->r);
2568 if (!indices)
2569 return NULL;
2570 for (i=0; i<lz->r; i++)
2571 {
2572 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2573 if (!index) {
2574 Py_DECREF(indices);
2575 return NULL;
2576 }
2577 PyTuple_SET_ITEM(indices, i, index);
2578 }
2579
2580 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2581 }
2582}
2583
2584static PyObject *
2585combinations_setstate(combinationsobject *lz, PyObject *state)
2586{
2587 PyObject *result;
2588 Py_ssize_t i;
2589 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2590
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002591 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002592 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2593 return NULL;
2594 }
2595
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002596 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002597 Py_ssize_t max;
2598 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2599 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002600
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002601 if (index == -1 && PyErr_Occurred())
2602 return NULL; /* not an integer */
2603 max = i + n - lz->r;
2604 /* clamp the index (beware of negative max) */
2605 if (index > max)
2606 index = max;
2607 if (index < 0)
2608 index = 0;
2609 lz->indices[i] = index;
2610 }
2611
2612 result = PyTuple_New(lz->r);
2613 if (result == NULL)
2614 return NULL;
2615 for (i=0; i<lz->r; i++) {
2616 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2617 Py_INCREF(element);
2618 PyTuple_SET_ITEM(result, i, element);
2619 }
2620
Serhiy Storchaka48842712016-04-06 09:45:48 +03002621 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002622 Py_RETURN_NONE;
2623}
2624
2625static PyMethodDef combinations_methods[] = {
2626 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2627 reduce_doc},
2628 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2629 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002630 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2631 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002632 {NULL, NULL} /* sentinel */
2633};
2634
Christian Heimes380f7f22008-02-28 11:19:05 +00002635PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002636"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002637\n\
2638Return successive r-length combinations of elements in the iterable.\n\n\
2639combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2640
2641static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002642 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002643 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002644 sizeof(combinationsobject), /* tp_basicsize */
2645 0, /* tp_itemsize */
2646 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002647 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002648 0, /* tp_print */
2649 0, /* tp_getattr */
2650 0, /* tp_setattr */
2651 0, /* tp_reserved */
2652 0, /* tp_repr */
2653 0, /* tp_as_number */
2654 0, /* tp_as_sequence */
2655 0, /* tp_as_mapping */
2656 0, /* tp_hash */
2657 0, /* tp_call */
2658 0, /* tp_str */
2659 PyObject_GenericGetAttr, /* tp_getattro */
2660 0, /* tp_setattro */
2661 0, /* tp_as_buffer */
2662 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2663 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002664 combinations_doc, /* tp_doc */
2665 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002666 0, /* tp_clear */
2667 0, /* tp_richcompare */
2668 0, /* tp_weaklistoffset */
2669 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002670 (iternextfunc)combinations_next, /* tp_iternext */
2671 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002672 0, /* tp_members */
2673 0, /* tp_getset */
2674 0, /* tp_base */
2675 0, /* tp_dict */
2676 0, /* tp_descr_get */
2677 0, /* tp_descr_set */
2678 0, /* tp_dictoffset */
2679 0, /* tp_init */
2680 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002681 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002682 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002683};
2684
2685
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002686/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002687
2688/* Equivalent to:
2689
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002690 def combinations_with_replacement(iterable, r):
2691 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2692 # number items returned: (n+r-1)! / r! / (n-1)!
2693 pool = tuple(iterable)
2694 n = len(pool)
2695 indices = [0] * r
2696 yield tuple(pool[i] for i in indices)
2697 while 1:
2698 for i in reversed(range(r)):
2699 if indices[i] != n - 1:
2700 break
2701 else:
2702 return
2703 indices[i:] = [indices[i] + 1] * (r - i)
2704 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002705
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002706 def combinations_with_replacement2(iterable, r):
2707 'Alternate version that filters from product()'
2708 pool = tuple(iterable)
2709 n = len(pool)
2710 for indices in product(range(n), repeat=r):
2711 if sorted(indices) == list(indices):
2712 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002713*/
2714typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002715 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002716 PyObject *pool; /* input converted to a tuple */
2717 Py_ssize_t *indices; /* one index per result element */
2718 PyObject *result; /* most recently returned result tuple */
2719 Py_ssize_t r; /* size of result tuple */
2720 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002721} cwrobject;
2722
2723static PyTypeObject cwr_type;
2724
2725static PyObject *
2726cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2727{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002728 cwrobject *co;
2729 Py_ssize_t n;
2730 Py_ssize_t r;
2731 PyObject *pool = NULL;
2732 PyObject *iterable = NULL;
2733 Py_ssize_t *indices = NULL;
2734 Py_ssize_t i;
2735 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002736
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002737 if (!PyArg_ParseTupleAndKeywords(args, kwds,
2738 "On:combinations_with_replacement",
2739 kwargs, &iterable, &r))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002740 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002741
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002742 pool = PySequence_Tuple(iterable);
2743 if (pool == NULL)
2744 goto error;
2745 n = PyTuple_GET_SIZE(pool);
2746 if (r < 0) {
2747 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2748 goto error;
2749 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002750
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002751 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002752 if (indices == NULL) {
2753 PyErr_NoMemory();
2754 goto error;
2755 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002756
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002757 for (i=0 ; i<r ; i++)
2758 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002759
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002760 /* create cwrobject structure */
2761 co = (cwrobject *)type->tp_alloc(type, 0);
2762 if (co == NULL)
2763 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002764
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002765 co->pool = pool;
2766 co->indices = indices;
2767 co->result = NULL;
2768 co->r = r;
2769 co->stopped = !n && r;
2770
2771 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002772
2773error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002774 if (indices != NULL)
2775 PyMem_Free(indices);
2776 Py_XDECREF(pool);
2777 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002778}
2779
2780static void
2781cwr_dealloc(cwrobject *co)
2782{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002783 PyObject_GC_UnTrack(co);
2784 Py_XDECREF(co->pool);
2785 Py_XDECREF(co->result);
2786 if (co->indices != NULL)
2787 PyMem_Free(co->indices);
2788 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002789}
2790
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002791static PyObject *
2792cwr_sizeof(cwrobject *co, void *unused)
2793{
2794 Py_ssize_t res;
2795
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002796 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002797 res += co->r * sizeof(Py_ssize_t);
2798 return PyLong_FromSsize_t(res);
2799}
2800
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002801static int
2802cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2803{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002804 Py_VISIT(co->pool);
2805 Py_VISIT(co->result);
2806 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002807}
2808
2809static PyObject *
2810cwr_next(cwrobject *co)
2811{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002812 PyObject *elem;
2813 PyObject *oldelem;
2814 PyObject *pool = co->pool;
2815 Py_ssize_t *indices = co->indices;
2816 PyObject *result = co->result;
2817 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2818 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002819 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002820
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002821 if (co->stopped)
2822 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002823
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002824 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002825 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002826 result = PyTuple_New(r);
2827 if (result == NULL)
2828 goto empty;
2829 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002830 if (n > 0) {
2831 elem = PyTuple_GET_ITEM(pool, 0);
2832 for (i=0; i<r ; i++) {
2833 assert(indices[i] == 0);
2834 Py_INCREF(elem);
2835 PyTuple_SET_ITEM(result, i, elem);
2836 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002837 }
2838 } else {
2839 /* Copy the previous result tuple or re-use it if available */
2840 if (Py_REFCNT(result) > 1) {
2841 PyObject *old_result = result;
2842 result = PyTuple_New(r);
2843 if (result == NULL)
2844 goto empty;
2845 co->result = result;
2846 for (i=0; i<r ; i++) {
2847 elem = PyTuple_GET_ITEM(old_result, i);
2848 Py_INCREF(elem);
2849 PyTuple_SET_ITEM(result, i, elem);
2850 }
2851 Py_DECREF(old_result);
2852 }
2853 /* Now, we've got the only copy so we can update it in-place CPython's
2854 empty tuple is a singleton and cached in PyTuple's freelist. */
2855 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002856
Tim Peters9edb1682013-09-03 11:49:31 -05002857 /* Scan indices right-to-left until finding one that is not
2858 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002859 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2860 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002861
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002862 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002863 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002864 if (i < 0)
2865 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002866
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002867 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002868 maximum. Then set all to the right to the same value. */
2869 index = indices[i] + 1;
2870 assert(index < n);
2871 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002872 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002873 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002874 Py_INCREF(elem);
2875 oldelem = PyTuple_GET_ITEM(result, i);
2876 PyTuple_SET_ITEM(result, i, elem);
2877 Py_DECREF(oldelem);
2878 }
2879 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002880
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002881 Py_INCREF(result);
2882 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002883
2884empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002885 co->stopped = 1;
2886 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002887}
2888
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002889static PyObject *
2890cwr_reduce(cwrobject *lz)
2891{
2892 if (lz->result == NULL) {
2893 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2894 } else if (lz->stopped) {
2895 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2896 } else {
2897 PyObject *indices;
2898 Py_ssize_t i;
2899
2900 /* we must pickle the indices and use them for setstate */
2901 indices = PyTuple_New(lz->r);
2902 if (!indices)
2903 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002904 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002905 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2906 if (!index) {
2907 Py_DECREF(indices);
2908 return NULL;
2909 }
2910 PyTuple_SET_ITEM(indices, i, index);
2911 }
2912
2913 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2914 }
2915}
2916
2917static PyObject *
2918cwr_setstate(cwrobject *lz, PyObject *state)
2919{
2920 PyObject *result;
2921 Py_ssize_t n, i;
2922
2923 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2924 {
2925 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2926 return NULL;
2927 }
2928
2929 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002930 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002931 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2932 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002933
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002934 if (index < 0 && PyErr_Occurred())
2935 return NULL; /* not an integer */
2936 /* clamp the index */
2937 if (index < 0)
2938 index = 0;
2939 else if (index > n-1)
2940 index = n-1;
2941 lz->indices[i] = index;
2942 }
2943 result = PyTuple_New(lz->r);
2944 if (result == NULL)
2945 return NULL;
2946 for (i=0; i<lz->r; i++) {
2947 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2948 Py_INCREF(element);
2949 PyTuple_SET_ITEM(result, i, element);
2950 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002951 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002952 Py_RETURN_NONE;
2953}
2954
2955static PyMethodDef cwr_methods[] = {
2956 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2957 reduce_doc},
2958 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2959 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002960 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
2961 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002962 {NULL, NULL} /* sentinel */
2963};
2964
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002965PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002966"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002967\n\
2968Return successive r-length combinations of elements in the iterable\n\
2969allowing individual elements to have successive repeats.\n\
2970combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2971
2972static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002973 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002974 "itertools.combinations_with_replacement", /* tp_name */
2975 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002976 0, /* tp_itemsize */
2977 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002978 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002979 0, /* tp_print */
2980 0, /* tp_getattr */
2981 0, /* tp_setattr */
2982 0, /* tp_reserved */
2983 0, /* tp_repr */
2984 0, /* tp_as_number */
2985 0, /* tp_as_sequence */
2986 0, /* tp_as_mapping */
2987 0, /* tp_hash */
2988 0, /* tp_call */
2989 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002990 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002991 0, /* tp_setattro */
2992 0, /* tp_as_buffer */
2993 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002994 Py_TPFLAGS_BASETYPE, /* tp_flags */
2995 cwr_doc, /* tp_doc */
2996 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002997 0, /* tp_clear */
2998 0, /* tp_richcompare */
2999 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003000 PyObject_SelfIter, /* tp_iter */
3001 (iternextfunc)cwr_next, /* tp_iternext */
3002 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003003 0, /* tp_members */
3004 0, /* tp_getset */
3005 0, /* tp_base */
3006 0, /* tp_dict */
3007 0, /* tp_descr_get */
3008 0, /* tp_descr_set */
3009 0, /* tp_dictoffset */
3010 0, /* tp_init */
3011 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003012 cwr_new, /* tp_new */
3013 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003014};
3015
3016
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003017/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003018
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003019def permutations(iterable, r=None):
3020 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
3021 pool = tuple(iterable)
3022 n = len(pool)
3023 r = n if r is None else r
3024 indices = range(n)
3025 cycles = range(n-r+1, n+1)[::-1]
3026 yield tuple(pool[i] for i in indices[:r])
3027 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003028 for i in reversed(range(r)):
3029 cycles[i] -= 1
3030 if cycles[i] == 0:
3031 indices[i:] = indices[i+1:] + indices[i:i+1]
3032 cycles[i] = n - i
3033 else:
3034 j = cycles[i]
3035 indices[i], indices[-j] = indices[-j], indices[i]
3036 yield tuple(pool[i] for i in indices[:r])
3037 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003038 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003039 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003040*/
3041
3042typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003043 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003044 PyObject *pool; /* input converted to a tuple */
3045 Py_ssize_t *indices; /* one index per element in the pool */
3046 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3047 PyObject *result; /* most recently returned result tuple */
3048 Py_ssize_t r; /* size of result tuple */
3049 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003050} permutationsobject;
3051
3052static PyTypeObject permutations_type;
3053
3054static PyObject *
3055permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3056{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003057 permutationsobject *po;
3058 Py_ssize_t n;
3059 Py_ssize_t r;
3060 PyObject *robj = Py_None;
3061 PyObject *pool = NULL;
3062 PyObject *iterable = NULL;
3063 Py_ssize_t *indices = NULL;
3064 Py_ssize_t *cycles = NULL;
3065 Py_ssize_t i;
3066 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003067
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003068 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
3069 &iterable, &robj))
3070 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003071
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003072 pool = PySequence_Tuple(iterable);
3073 if (pool == NULL)
3074 goto error;
3075 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003076
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003077 r = n;
3078 if (robj != Py_None) {
3079 if (!PyLong_Check(robj)) {
3080 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3081 goto error;
3082 }
3083 r = PyLong_AsSsize_t(robj);
3084 if (r == -1 && PyErr_Occurred())
3085 goto error;
3086 }
3087 if (r < 0) {
3088 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3089 goto error;
3090 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003091
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003092 indices = PyMem_New(Py_ssize_t, n);
3093 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003094 if (indices == NULL || cycles == NULL) {
3095 PyErr_NoMemory();
3096 goto error;
3097 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003098
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003099 for (i=0 ; i<n ; i++)
3100 indices[i] = i;
3101 for (i=0 ; i<r ; i++)
3102 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003103
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003104 /* create permutationsobject structure */
3105 po = (permutationsobject *)type->tp_alloc(type, 0);
3106 if (po == NULL)
3107 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003108
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003109 po->pool = pool;
3110 po->indices = indices;
3111 po->cycles = cycles;
3112 po->result = NULL;
3113 po->r = r;
3114 po->stopped = r > n ? 1 : 0;
3115
3116 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003117
3118error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003119 if (indices != NULL)
3120 PyMem_Free(indices);
3121 if (cycles != NULL)
3122 PyMem_Free(cycles);
3123 Py_XDECREF(pool);
3124 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003125}
3126
3127static void
3128permutations_dealloc(permutationsobject *po)
3129{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003130 PyObject_GC_UnTrack(po);
3131 Py_XDECREF(po->pool);
3132 Py_XDECREF(po->result);
3133 PyMem_Free(po->indices);
3134 PyMem_Free(po->cycles);
3135 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003136}
3137
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003138static PyObject *
3139permutations_sizeof(permutationsobject *po, void *unused)
3140{
3141 Py_ssize_t res;
3142
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003143 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003144 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3145 res += po->r * sizeof(Py_ssize_t);
3146 return PyLong_FromSsize_t(res);
3147}
3148
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003149static int
3150permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3151{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003152 Py_VISIT(po->pool);
3153 Py_VISIT(po->result);
3154 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003155}
3156
3157static PyObject *
3158permutations_next(permutationsobject *po)
3159{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003160 PyObject *elem;
3161 PyObject *oldelem;
3162 PyObject *pool = po->pool;
3163 Py_ssize_t *indices = po->indices;
3164 Py_ssize_t *cycles = po->cycles;
3165 PyObject *result = po->result;
3166 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3167 Py_ssize_t r = po->r;
3168 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003169
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003170 if (po->stopped)
3171 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003172
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003173 if (result == NULL) {
3174 /* On the first pass, initialize result tuple using the indices */
3175 result = PyTuple_New(r);
3176 if (result == NULL)
3177 goto empty;
3178 po->result = result;
3179 for (i=0; i<r ; i++) {
3180 index = indices[i];
3181 elem = PyTuple_GET_ITEM(pool, index);
3182 Py_INCREF(elem);
3183 PyTuple_SET_ITEM(result, i, elem);
3184 }
3185 } else {
3186 if (n == 0)
3187 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003188
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003189 /* Copy the previous result tuple or re-use it if available */
3190 if (Py_REFCNT(result) > 1) {
3191 PyObject *old_result = result;
3192 result = PyTuple_New(r);
3193 if (result == NULL)
3194 goto empty;
3195 po->result = result;
3196 for (i=0; i<r ; i++) {
3197 elem = PyTuple_GET_ITEM(old_result, i);
3198 Py_INCREF(elem);
3199 PyTuple_SET_ITEM(result, i, elem);
3200 }
3201 Py_DECREF(old_result);
3202 }
3203 /* Now, we've got the only copy so we can update it in-place */
3204 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003205
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003206 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3207 for (i=r-1 ; i>=0 ; i--) {
3208 cycles[i] -= 1;
3209 if (cycles[i] == 0) {
3210 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3211 index = indices[i];
3212 for (j=i ; j<n-1 ; j++)
3213 indices[j] = indices[j+1];
3214 indices[n-1] = index;
3215 cycles[i] = n - i;
3216 } else {
3217 j = cycles[i];
3218 index = indices[i];
3219 indices[i] = indices[n-j];
3220 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003221
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003222 for (k=i; k<r ; k++) {
3223 /* start with i, the leftmost element that changed */
3224 /* yield tuple(pool[k] for k in indices[:r]) */
3225 index = indices[k];
3226 elem = PyTuple_GET_ITEM(pool, index);
3227 Py_INCREF(elem);
3228 oldelem = PyTuple_GET_ITEM(result, k);
3229 PyTuple_SET_ITEM(result, k, elem);
3230 Py_DECREF(oldelem);
3231 }
3232 break;
3233 }
3234 }
3235 /* If i is negative, then the cycles have all
3236 rolled-over and we're done. */
3237 if (i < 0)
3238 goto empty;
3239 }
3240 Py_INCREF(result);
3241 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003242
3243empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003244 po->stopped = 1;
3245 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003246}
3247
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003248static PyObject *
3249permutations_reduce(permutationsobject *po)
3250{
3251 if (po->result == NULL) {
3252 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3253 } else if (po->stopped) {
3254 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3255 } else {
3256 PyObject *indices=NULL, *cycles=NULL;
3257 Py_ssize_t n, i;
3258
3259 /* we must pickle the indices and cycles and use them for setstate */
3260 n = PyTuple_GET_SIZE(po->pool);
3261 indices = PyTuple_New(n);
3262 if (indices == NULL)
3263 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003264 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003265 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3266 if (!index)
3267 goto err;
3268 PyTuple_SET_ITEM(indices, i, index);
3269 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003270
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003271 cycles = PyTuple_New(po->r);
3272 if (cycles == NULL)
3273 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003274 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003275 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3276 if (!index)
3277 goto err;
3278 PyTuple_SET_ITEM(cycles, i, index);
3279 }
3280 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3281 po->pool, po->r,
3282 indices, cycles);
3283 err:
3284 Py_XDECREF(indices);
3285 Py_XDECREF(cycles);
3286 return NULL;
3287 }
3288}
3289
3290static PyObject *
3291permutations_setstate(permutationsobject *po, PyObject *state)
3292{
3293 PyObject *indices, *cycles, *result;
3294 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003295
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003296 if (!PyTuple_Check(state)) {
3297 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3298 return NULL;
3299 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003300 if (!PyArg_ParseTuple(state, "O!O!",
3301 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003302 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003303 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003304 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003305
3306 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003307 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003308 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3309 return NULL;
3310 }
3311
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003312 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003313 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3314 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3315 if (index < 0 && PyErr_Occurred())
3316 return NULL; /* not an integer */
3317 /* clamp the index */
3318 if (index < 0)
3319 index = 0;
3320 else if (index > n-1)
3321 index = n-1;
3322 po->indices[i] = index;
3323 }
3324
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003325 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003326 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3327 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3328 if (index < 0 && PyErr_Occurred())
3329 return NULL; /* not an integer */
3330 if (index < 1)
3331 index = 1;
3332 else if (index > n-i)
3333 index = n-i;
3334 po->cycles[i] = index;
3335 }
3336 result = PyTuple_New(po->r);
3337 if (result == NULL)
3338 return NULL;
3339 for (i=0; i<po->r; i++) {
3340 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3341 Py_INCREF(element);
3342 PyTuple_SET_ITEM(result, i, element);
3343 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003344 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003345 Py_RETURN_NONE;
3346}
3347
3348static PyMethodDef permuations_methods[] = {
3349 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3350 reduce_doc},
3351 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3352 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003353 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3354 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003355 {NULL, NULL} /* sentinel */
3356};
3357
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003358PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003359"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003360\n\
3361Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003362permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003363
3364static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003365 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003366 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003367 sizeof(permutationsobject), /* tp_basicsize */
3368 0, /* tp_itemsize */
3369 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003370 (destructor)permutations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003371 0, /* tp_print */
3372 0, /* tp_getattr */
3373 0, /* tp_setattr */
3374 0, /* tp_reserved */
3375 0, /* tp_repr */
3376 0, /* tp_as_number */
3377 0, /* tp_as_sequence */
3378 0, /* tp_as_mapping */
3379 0, /* tp_hash */
3380 0, /* tp_call */
3381 0, /* tp_str */
3382 PyObject_GenericGetAttr, /* tp_getattro */
3383 0, /* tp_setattro */
3384 0, /* tp_as_buffer */
3385 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3386 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003387 permutations_doc, /* tp_doc */
3388 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003389 0, /* tp_clear */
3390 0, /* tp_richcompare */
3391 0, /* tp_weaklistoffset */
3392 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003393 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003394 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003395 0, /* tp_members */
3396 0, /* tp_getset */
3397 0, /* tp_base */
3398 0, /* tp_dict */
3399 0, /* tp_descr_get */
3400 0, /* tp_descr_set */
3401 0, /* tp_dictoffset */
3402 0, /* tp_init */
3403 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003404 permutations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003405 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003406};
3407
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003408/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003409
3410typedef struct {
3411 PyObject_HEAD
3412 PyObject *total;
3413 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003414 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003415} accumulateobject;
3416
3417static PyTypeObject accumulate_type;
3418
3419static PyObject *
3420accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3421{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003422 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003423 PyObject *iterable;
3424 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003425 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003426 accumulateobject *lz;
3427
Raymond Hettinger5d446132011-03-27 18:52:10 -07003428 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3429 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003430 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003431
3432 /* Get iterator. */
3433 it = PyObject_GetIter(iterable);
3434 if (it == NULL)
3435 return NULL;
3436
Raymond Hettinger482ba772010-12-01 22:48:00 +00003437 /* create accumulateobject structure */
3438 lz = (accumulateobject *)type->tp_alloc(type, 0);
3439 if (lz == NULL) {
3440 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003441 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003442 }
3443
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003444 if (binop != Py_None) {
3445 Py_XINCREF(binop);
3446 lz->binop = binop;
3447 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003448 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003449 lz->it = it;
3450 return (PyObject *)lz;
3451}
3452
3453static void
3454accumulate_dealloc(accumulateobject *lz)
3455{
3456 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003457 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003458 Py_XDECREF(lz->total);
3459 Py_XDECREF(lz->it);
3460 Py_TYPE(lz)->tp_free(lz);
3461}
3462
3463static int
3464accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3465{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003466 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003467 Py_VISIT(lz->it);
3468 Py_VISIT(lz->total);
3469 return 0;
3470}
3471
3472static PyObject *
3473accumulate_next(accumulateobject *lz)
3474{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003475 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003476
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003477 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003478 if (val == NULL)
3479 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003480
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003481 if (lz->total == NULL) {
3482 Py_INCREF(val);
3483 lz->total = val;
3484 return lz->total;
3485 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003486
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003487 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003488 newtotal = PyNumber_Add(lz->total, val);
3489 else
3490 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003491 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003492 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003493 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003494
Raymond Hettinger482ba772010-12-01 22:48:00 +00003495 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003496 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003497 return newtotal;
3498}
3499
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003500static PyObject *
3501accumulate_reduce(accumulateobject *lz)
3502{
Serhiy Storchakad5516252016-03-06 14:00:45 +02003503 if (lz->total == Py_None) {
3504 PyObject *it;
3505
3506 if (PyType_Ready(&chain_type) < 0)
3507 return NULL;
3508 if (PyType_Ready(&islice_type) < 0)
3509 return NULL;
3510 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3511 lz->total, lz->it);
3512 if (it == NULL)
3513 return NULL;
3514 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3515 it, lz->binop ? lz->binop : Py_None);
3516 if (it == NULL)
3517 return NULL;
3518 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3519 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003520 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3521 lz->it, lz->binop?lz->binop:Py_None,
3522 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003523}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003524
3525static PyObject *
3526accumulate_setstate(accumulateobject *lz, PyObject *state)
3527{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003528 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003529 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003530 Py_RETURN_NONE;
3531}
3532
3533static PyMethodDef accumulate_methods[] = {
3534 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3535 reduce_doc},
3536 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3537 setstate_doc},
3538 {NULL, NULL} /* sentinel */
3539};
3540
Raymond Hettinger482ba772010-12-01 22:48:00 +00003541PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003542"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003543\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003544Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003545
3546static PyTypeObject accumulate_type = {
3547 PyVarObject_HEAD_INIT(NULL, 0)
3548 "itertools.accumulate", /* tp_name */
3549 sizeof(accumulateobject), /* tp_basicsize */
3550 0, /* tp_itemsize */
3551 /* methods */
3552 (destructor)accumulate_dealloc, /* tp_dealloc */
3553 0, /* tp_print */
3554 0, /* tp_getattr */
3555 0, /* tp_setattr */
3556 0, /* tp_reserved */
3557 0, /* tp_repr */
3558 0, /* tp_as_number */
3559 0, /* tp_as_sequence */
3560 0, /* tp_as_mapping */
3561 0, /* tp_hash */
3562 0, /* tp_call */
3563 0, /* tp_str */
3564 PyObject_GenericGetAttr, /* tp_getattro */
3565 0, /* tp_setattro */
3566 0, /* tp_as_buffer */
3567 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3568 Py_TPFLAGS_BASETYPE, /* tp_flags */
3569 accumulate_doc, /* tp_doc */
3570 (traverseproc)accumulate_traverse, /* tp_traverse */
3571 0, /* tp_clear */
3572 0, /* tp_richcompare */
3573 0, /* tp_weaklistoffset */
3574 PyObject_SelfIter, /* tp_iter */
3575 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003576 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003577 0, /* tp_members */
3578 0, /* tp_getset */
3579 0, /* tp_base */
3580 0, /* tp_dict */
3581 0, /* tp_descr_get */
3582 0, /* tp_descr_set */
3583 0, /* tp_dictoffset */
3584 0, /* tp_init */
3585 0, /* tp_alloc */
3586 accumulate_new, /* tp_new */
3587 PyObject_GC_Del, /* tp_free */
3588};
3589
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003590
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003591/* compress object ************************************************************/
3592
3593/* Equivalent to:
3594
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003595 def compress(data, selectors):
3596 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3597 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003598*/
3599
3600typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003601 PyObject_HEAD
3602 PyObject *data;
3603 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003604} compressobject;
3605
3606static PyTypeObject compress_type;
3607
3608static PyObject *
3609compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3610{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003611 PyObject *seq1, *seq2;
3612 PyObject *data=NULL, *selectors=NULL;
3613 compressobject *lz;
3614 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003615
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003616 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3617 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003618
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003619 data = PyObject_GetIter(seq1);
3620 if (data == NULL)
3621 goto fail;
3622 selectors = PyObject_GetIter(seq2);
3623 if (selectors == NULL)
3624 goto fail;
3625
3626 /* create compressobject structure */
3627 lz = (compressobject *)type->tp_alloc(type, 0);
3628 if (lz == NULL)
3629 goto fail;
3630 lz->data = data;
3631 lz->selectors = selectors;
3632 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003633
3634fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003635 Py_XDECREF(data);
3636 Py_XDECREF(selectors);
3637 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003638}
3639
3640static void
3641compress_dealloc(compressobject *lz)
3642{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003643 PyObject_GC_UnTrack(lz);
3644 Py_XDECREF(lz->data);
3645 Py_XDECREF(lz->selectors);
3646 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003647}
3648
3649static int
3650compress_traverse(compressobject *lz, visitproc visit, void *arg)
3651{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003652 Py_VISIT(lz->data);
3653 Py_VISIT(lz->selectors);
3654 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003655}
3656
3657static PyObject *
3658compress_next(compressobject *lz)
3659{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003660 PyObject *data = lz->data, *selectors = lz->selectors;
3661 PyObject *datum, *selector;
3662 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3663 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3664 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003665
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003666 while (1) {
3667 /* Steps: get datum, get selector, evaluate selector.
3668 Order is important (to match the pure python version
3669 in terms of which input gets a chance to raise an
3670 exception first).
3671 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003672
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003673 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003674 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003675 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003676
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003677 selector = selectornext(selectors);
3678 if (selector == NULL) {
3679 Py_DECREF(datum);
3680 return NULL;
3681 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003682
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003683 ok = PyObject_IsTrue(selector);
3684 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003685 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003686 return datum;
3687 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003688 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003689 return NULL;
3690 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003691}
3692
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003693static PyObject *
3694compress_reduce(compressobject *lz)
3695{
3696 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3697 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003698}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003699
3700static PyMethodDef compress_methods[] = {
3701 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3702 reduce_doc},
3703 {NULL, NULL} /* sentinel */
3704};
3705
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003706PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003707"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003708\n\
3709Return data elements corresponding to true selector elements.\n\
3710Forms a shorter iterator from selected data elements using the\n\
3711selectors to choose the data elements.");
3712
3713static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003714 PyVarObject_HEAD_INIT(NULL, 0)
3715 "itertools.compress", /* tp_name */
3716 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003717 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003718 /* methods */
3719 (destructor)compress_dealloc, /* tp_dealloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003720 0, /* tp_print */
3721 0, /* tp_getattr */
3722 0, /* tp_setattr */
3723 0, /* tp_reserved */
3724 0, /* tp_repr */
3725 0, /* tp_as_number */
3726 0, /* tp_as_sequence */
3727 0, /* tp_as_mapping */
3728 0, /* tp_hash */
3729 0, /* tp_call */
3730 0, /* tp_str */
3731 PyObject_GenericGetAttr, /* tp_getattro */
3732 0, /* tp_setattro */
3733 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003734 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003735 Py_TPFLAGS_BASETYPE, /* tp_flags */
3736 compress_doc, /* tp_doc */
3737 (traverseproc)compress_traverse, /* tp_traverse */
3738 0, /* tp_clear */
3739 0, /* tp_richcompare */
3740 0, /* tp_weaklistoffset */
3741 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003742 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003743 compress_methods, /* tp_methods */
3744 0, /* tp_members */
3745 0, /* tp_getset */
3746 0, /* tp_base */
3747 0, /* tp_dict */
3748 0, /* tp_descr_get */
3749 0, /* tp_descr_set */
3750 0, /* tp_dictoffset */
3751 0, /* tp_init */
3752 0, /* tp_alloc */
3753 compress_new, /* tp_new */
3754 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003755};
3756
3757
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003758/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003759
3760typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003761 PyObject_HEAD
3762 PyObject *func;
3763 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003764} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003765
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003766static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003767
3768static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003769filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003770{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003771 PyObject *func, *seq;
3772 PyObject *it;
3773 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003774
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003775 if (type == &filterfalse_type &&
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03003776 !_PyArg_NoKeywords("filterfalse", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003777 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003778
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003779 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3780 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003781
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003782 /* Get iterator. */
3783 it = PyObject_GetIter(seq);
3784 if (it == NULL)
3785 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003786
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003787 /* create filterfalseobject structure */
3788 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3789 if (lz == NULL) {
3790 Py_DECREF(it);
3791 return NULL;
3792 }
3793 Py_INCREF(func);
3794 lz->func = func;
3795 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003796
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003797 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003798}
3799
3800static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003801filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003802{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003803 PyObject_GC_UnTrack(lz);
3804 Py_XDECREF(lz->func);
3805 Py_XDECREF(lz->it);
3806 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003807}
3808
3809static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003810filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003811{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003812 Py_VISIT(lz->it);
3813 Py_VISIT(lz->func);
3814 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003815}
3816
3817static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003818filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003819{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003820 PyObject *item;
3821 PyObject *it = lz->it;
3822 long ok;
3823 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003824
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003825 iternext = *Py_TYPE(it)->tp_iternext;
3826 for (;;) {
3827 item = iternext(it);
3828 if (item == NULL)
3829 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003830
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003831 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3832 ok = PyObject_IsTrue(item);
3833 } else {
3834 PyObject *good;
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01003835 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003836 if (good == NULL) {
3837 Py_DECREF(item);
3838 return NULL;
3839 }
3840 ok = PyObject_IsTrue(good);
3841 Py_DECREF(good);
3842 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003843 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003844 return item;
3845 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003846 if (ok < 0)
3847 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003848 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003849}
3850
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003851static PyObject *
3852filterfalse_reduce(filterfalseobject *lz)
3853{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003854 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3855}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003856
3857static PyMethodDef filterfalse_methods[] = {
3858 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3859 reduce_doc},
3860 {NULL, NULL} /* sentinel */
3861};
3862
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003863PyDoc_STRVAR(filterfalse_doc,
3864"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003865\n\
3866Return those items of sequence for which function(item) is false.\n\
3867If function is None, return the items that are false.");
3868
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003869static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003870 PyVarObject_HEAD_INIT(NULL, 0)
3871 "itertools.filterfalse", /* tp_name */
3872 sizeof(filterfalseobject), /* tp_basicsize */
3873 0, /* tp_itemsize */
3874 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003875 (destructor)filterfalse_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003876 0, /* tp_print */
3877 0, /* tp_getattr */
3878 0, /* tp_setattr */
3879 0, /* tp_reserved */
3880 0, /* tp_repr */
3881 0, /* tp_as_number */
3882 0, /* tp_as_sequence */
3883 0, /* tp_as_mapping */
3884 0, /* tp_hash */
3885 0, /* tp_call */
3886 0, /* tp_str */
3887 PyObject_GenericGetAttr, /* tp_getattro */
3888 0, /* tp_setattro */
3889 0, /* tp_as_buffer */
3890 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3891 Py_TPFLAGS_BASETYPE, /* tp_flags */
3892 filterfalse_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003893 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003894 0, /* tp_clear */
3895 0, /* tp_richcompare */
3896 0, /* tp_weaklistoffset */
3897 PyObject_SelfIter, /* tp_iter */
3898 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003899 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003900 0, /* tp_members */
3901 0, /* tp_getset */
3902 0, /* tp_base */
3903 0, /* tp_dict */
3904 0, /* tp_descr_get */
3905 0, /* tp_descr_set */
3906 0, /* tp_dictoffset */
3907 0, /* tp_init */
3908 0, /* tp_alloc */
3909 filterfalse_new, /* tp_new */
3910 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003911};
3912
3913
3914/* count object ************************************************************/
3915
3916typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003917 PyObject_HEAD
3918 Py_ssize_t cnt;
3919 PyObject *long_cnt;
3920 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003921} countobject;
3922
Raymond Hettinger30729212009-02-12 06:28:27 +00003923/* Counting logic and invariants:
3924
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003925fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003926
Serhiy Storchaka483405b2015-02-17 10:14:30 +02003927 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003928 Advances with: cnt += 1
3929 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003930
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003931slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003932
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003933 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3934 All counting is done with python objects (no overflows or underflows).
3935 Advances with: long_cnt += long_step
3936 Step may be zero -- effectively a slow version of repeat(cnt).
3937 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003938*/
3939
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003940static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003941
3942static PyObject *
3943count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3944{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003945 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003946 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003947 Py_ssize_t cnt = 0;
3948 PyObject *long_cnt = NULL;
3949 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003950 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003951 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003952
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003953 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3954 kwlist, &long_cnt, &long_step))
3955 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003956
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003957 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3958 (long_step != NULL && !PyNumber_Check(long_step))) {
3959 PyErr_SetString(PyExc_TypeError, "a number is required");
3960 return NULL;
3961 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003962
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003963 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
3964 (long_step == NULL || PyLong_Check(long_step));
3965
3966 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003967 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003968 if (fast_mode) {
3969 assert(PyLong_Check(long_cnt));
3970 cnt = PyLong_AsSsize_t(long_cnt);
3971 if (cnt == -1 && PyErr_Occurred()) {
3972 PyErr_Clear();
3973 fast_mode = 0;
3974 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003975 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003976 } else {
3977 cnt = 0;
Serhiy Storchakaba85d692017-03-30 09:09:41 +03003978 long_cnt = _PyLong_Zero;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003979 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03003980 Py_INCREF(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003981
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003982 /* If not specified, step defaults to 1 */
Serhiy Storchakaba85d692017-03-30 09:09:41 +03003983 if (long_step == NULL)
3984 long_step = _PyLong_One;
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)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004074 return PyUnicode_FromFormat("%s(%zd)",
4075 _PyType_Name(Py_TYPE(lz)), lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004076
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004077 if (PyLong_Check(lz->long_step)) {
4078 long step = PyLong_AsLong(lz->long_step);
4079 if (step == -1 && PyErr_Occurred()) {
4080 PyErr_Clear();
4081 }
4082 if (step == 1) {
4083 /* Don't display step when it is an integer equal to 1 */
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004084 return PyUnicode_FromFormat("%s(%R)",
4085 _PyType_Name(Py_TYPE(lz)),
4086 lz->long_cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004087 }
4088 }
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004089 return PyUnicode_FromFormat("%s(%R, %R)",
4090 _PyType_Name(Py_TYPE(lz)),
4091 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004092}
4093
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004094static PyObject *
4095count_reduce(countobject *lz)
4096{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004097 if (lz->cnt == PY_SSIZE_T_MAX)
4098 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4099 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004100}
4101
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004102static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004103 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004104 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004105 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004106};
4107
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004108PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00004109 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004110\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004111Return a count object whose .__next__() method returns consecutive values.\n\
4112Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00004113 def count(firstval=0, step=1):\n\
Eli Benderskyc554f722013-09-03 06:37:19 -07004114 x = firstval\n\
4115 while 1:\n\
4116 yield x\n\
4117 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004118
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004119static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004120 PyVarObject_HEAD_INIT(NULL, 0)
4121 "itertools.count", /* tp_name */
4122 sizeof(countobject), /* tp_basicsize */
4123 0, /* tp_itemsize */
4124 /* methods */
4125 (destructor)count_dealloc, /* tp_dealloc */
4126 0, /* tp_print */
4127 0, /* tp_getattr */
4128 0, /* tp_setattr */
4129 0, /* tp_reserved */
4130 (reprfunc)count_repr, /* tp_repr */
4131 0, /* tp_as_number */
4132 0, /* tp_as_sequence */
4133 0, /* tp_as_mapping */
4134 0, /* tp_hash */
4135 0, /* tp_call */
4136 0, /* tp_str */
4137 PyObject_GenericGetAttr, /* tp_getattro */
4138 0, /* tp_setattro */
4139 0, /* tp_as_buffer */
4140 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004141 Py_TPFLAGS_BASETYPE, /* tp_flags */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004142 count_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004143 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004144 0, /* tp_clear */
4145 0, /* tp_richcompare */
4146 0, /* tp_weaklistoffset */
4147 PyObject_SelfIter, /* tp_iter */
4148 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004149 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004150 0, /* tp_members */
4151 0, /* tp_getset */
4152 0, /* tp_base */
4153 0, /* tp_dict */
4154 0, /* tp_descr_get */
4155 0, /* tp_descr_set */
4156 0, /* tp_dictoffset */
4157 0, /* tp_init */
4158 0, /* tp_alloc */
4159 count_new, /* tp_new */
4160 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004161};
4162
4163
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004164/* repeat object ************************************************************/
4165
4166typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004167 PyObject_HEAD
4168 PyObject *element;
4169 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004170} repeatobject;
4171
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004172static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004173
4174static PyObject *
4175repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4176{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004177 repeatobject *ro;
4178 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004179 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004180 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004181
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004182 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4183 &element, &cnt))
4184 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004185
Raymond Hettinger97d35552014-06-24 21:36:58 -07004186 if (kwds != NULL)
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004187 n_kwds = PyDict_GET_SIZE(kwds);
Raymond Hettinger97d35552014-06-24 21:36:58 -07004188 /* Does user supply times argument? */
4189 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004190 cnt = 0;
4191
4192 ro = (repeatobject *)type->tp_alloc(type, 0);
4193 if (ro == NULL)
4194 return NULL;
4195 Py_INCREF(element);
4196 ro->element = element;
4197 ro->cnt = cnt;
4198 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004199}
4200
4201static void
4202repeat_dealloc(repeatobject *ro)
4203{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004204 PyObject_GC_UnTrack(ro);
4205 Py_XDECREF(ro->element);
4206 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004207}
4208
4209static int
4210repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4211{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004212 Py_VISIT(ro->element);
4213 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004214}
4215
4216static PyObject *
4217repeat_next(repeatobject *ro)
4218{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004219 if (ro->cnt == 0)
4220 return NULL;
4221 if (ro->cnt > 0)
4222 ro->cnt--;
4223 Py_INCREF(ro->element);
4224 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004225}
4226
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004227static PyObject *
4228repeat_repr(repeatobject *ro)
4229{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004230 if (ro->cnt == -1)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004231 return PyUnicode_FromFormat("%s(%R)",
4232 _PyType_Name(Py_TYPE(ro)), ro->element);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004233 else
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004234 return PyUnicode_FromFormat("%s(%R, %zd)",
4235 _PyType_Name(Py_TYPE(ro)), ro->element,
4236 ro->cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004237}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004238
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004239static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004240repeat_len(repeatobject *ro)
4241{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004242 if (ro->cnt == -1) {
4243 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4244 return NULL;
4245 }
4246 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004247}
4248
Armin Rigof5b3e362006-02-11 21:32:43 +00004249PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004250
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004251static PyObject *
4252repeat_reduce(repeatobject *ro)
4253{
4254 /* unpickle this so that a new repeat iterator is constructed with an
4255 * object, then call __setstate__ on it to set cnt
4256 */
4257 if (ro->cnt >= 0)
4258 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4259 else
4260 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4261}
4262
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004263static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004264 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004265 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004266 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004267};
4268
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004269PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004270"repeat(object [,times]) -> create an iterator which returns the object\n\
4271for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004272endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004273
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004274static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004275 PyVarObject_HEAD_INIT(NULL, 0)
4276 "itertools.repeat", /* tp_name */
4277 sizeof(repeatobject), /* tp_basicsize */
4278 0, /* tp_itemsize */
4279 /* methods */
4280 (destructor)repeat_dealloc, /* tp_dealloc */
4281 0, /* tp_print */
4282 0, /* tp_getattr */
4283 0, /* tp_setattr */
4284 0, /* tp_reserved */
4285 (reprfunc)repeat_repr, /* tp_repr */
4286 0, /* tp_as_number */
4287 0, /* tp_as_sequence */
4288 0, /* tp_as_mapping */
4289 0, /* tp_hash */
4290 0, /* tp_call */
4291 0, /* tp_str */
4292 PyObject_GenericGetAttr, /* tp_getattro */
4293 0, /* tp_setattro */
4294 0, /* tp_as_buffer */
4295 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4296 Py_TPFLAGS_BASETYPE, /* tp_flags */
4297 repeat_doc, /* tp_doc */
4298 (traverseproc)repeat_traverse, /* tp_traverse */
4299 0, /* tp_clear */
4300 0, /* tp_richcompare */
4301 0, /* tp_weaklistoffset */
4302 PyObject_SelfIter, /* tp_iter */
4303 (iternextfunc)repeat_next, /* tp_iternext */
4304 repeat_methods, /* tp_methods */
4305 0, /* tp_members */
4306 0, /* tp_getset */
4307 0, /* tp_base */
4308 0, /* tp_dict */
4309 0, /* tp_descr_get */
4310 0, /* tp_descr_set */
4311 0, /* tp_dictoffset */
4312 0, /* tp_init */
4313 0, /* tp_alloc */
4314 repeat_new, /* tp_new */
4315 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004316};
4317
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004318/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004319
4320typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004321 PyObject_HEAD
4322 Py_ssize_t tuplesize;
4323 Py_ssize_t numactive;
4324 PyObject *ittuple; /* tuple of iterators */
4325 PyObject *result;
4326 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004327} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004328
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004329static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004330
4331static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004332zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004333{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004334 ziplongestobject *lz;
4335 Py_ssize_t i;
4336 PyObject *ittuple; /* tuple of iterators */
4337 PyObject *result;
4338 PyObject *fillvalue = Py_None;
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004339 Py_ssize_t tuplesize;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004340
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004341 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004342 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004343 if (fillvalue == NULL || PyDict_GET_SIZE(kwds) > 1) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004344 PyErr_SetString(PyExc_TypeError,
4345 "zip_longest() got an unexpected keyword argument");
4346 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004347 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004348 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004349
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004350 /* args must be a tuple */
4351 assert(PyTuple_Check(args));
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004352 tuplesize = PyTuple_GET_SIZE(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004353
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004354 /* obtain iterators */
4355 ittuple = PyTuple_New(tuplesize);
4356 if (ittuple == NULL)
4357 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004358 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004359 PyObject *item = PyTuple_GET_ITEM(args, i);
4360 PyObject *it = PyObject_GetIter(item);
4361 if (it == NULL) {
4362 if (PyErr_ExceptionMatches(PyExc_TypeError))
4363 PyErr_Format(PyExc_TypeError,
4364 "zip_longest argument #%zd must support iteration",
4365 i+1);
4366 Py_DECREF(ittuple);
4367 return NULL;
4368 }
4369 PyTuple_SET_ITEM(ittuple, i, it);
4370 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004371
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004372 /* create a result holder */
4373 result = PyTuple_New(tuplesize);
4374 if (result == NULL) {
4375 Py_DECREF(ittuple);
4376 return NULL;
4377 }
4378 for (i=0 ; i < tuplesize ; i++) {
4379 Py_INCREF(Py_None);
4380 PyTuple_SET_ITEM(result, i, Py_None);
4381 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004382
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004383 /* create ziplongestobject structure */
4384 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4385 if (lz == NULL) {
4386 Py_DECREF(ittuple);
4387 Py_DECREF(result);
4388 return NULL;
4389 }
4390 lz->ittuple = ittuple;
4391 lz->tuplesize = tuplesize;
4392 lz->numactive = tuplesize;
4393 lz->result = result;
4394 Py_INCREF(fillvalue);
4395 lz->fillvalue = fillvalue;
4396 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004397}
4398
4399static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004400zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004401{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004402 PyObject_GC_UnTrack(lz);
4403 Py_XDECREF(lz->ittuple);
4404 Py_XDECREF(lz->result);
4405 Py_XDECREF(lz->fillvalue);
4406 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004407}
4408
4409static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004410zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004411{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004412 Py_VISIT(lz->ittuple);
4413 Py_VISIT(lz->result);
4414 Py_VISIT(lz->fillvalue);
4415 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004416}
4417
4418static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004419zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004420{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004421 Py_ssize_t i;
4422 Py_ssize_t tuplesize = lz->tuplesize;
4423 PyObject *result = lz->result;
4424 PyObject *it;
4425 PyObject *item;
4426 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004427
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004428 if (tuplesize == 0)
4429 return NULL;
4430 if (lz->numactive == 0)
4431 return NULL;
4432 if (Py_REFCNT(result) == 1) {
4433 Py_INCREF(result);
4434 for (i=0 ; i < tuplesize ; i++) {
4435 it = PyTuple_GET_ITEM(lz->ittuple, i);
4436 if (it == NULL) {
4437 Py_INCREF(lz->fillvalue);
4438 item = lz->fillvalue;
4439 } else {
4440 item = PyIter_Next(it);
4441 if (item == NULL) {
4442 lz->numactive -= 1;
4443 if (lz->numactive == 0 || PyErr_Occurred()) {
4444 lz->numactive = 0;
4445 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004446 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004447 } else {
4448 Py_INCREF(lz->fillvalue);
4449 item = lz->fillvalue;
4450 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4451 Py_DECREF(it);
4452 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004453 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004454 }
4455 olditem = PyTuple_GET_ITEM(result, i);
4456 PyTuple_SET_ITEM(result, i, item);
4457 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004458 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004459 } else {
4460 result = PyTuple_New(tuplesize);
4461 if (result == NULL)
4462 return NULL;
4463 for (i=0 ; i < tuplesize ; i++) {
4464 it = PyTuple_GET_ITEM(lz->ittuple, i);
4465 if (it == NULL) {
4466 Py_INCREF(lz->fillvalue);
4467 item = lz->fillvalue;
4468 } else {
4469 item = PyIter_Next(it);
4470 if (item == NULL) {
4471 lz->numactive -= 1;
4472 if (lz->numactive == 0 || PyErr_Occurred()) {
4473 lz->numactive = 0;
4474 Py_DECREF(result);
4475 return NULL;
4476 } else {
4477 Py_INCREF(lz->fillvalue);
4478 item = lz->fillvalue;
4479 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4480 Py_DECREF(it);
4481 }
4482 }
4483 }
4484 PyTuple_SET_ITEM(result, i, item);
4485 }
4486 }
4487 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004488}
4489
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004490static PyObject *
4491zip_longest_reduce(ziplongestobject *lz)
4492{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004493
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004494 /* Create a new tuple with empty sequences where appropriate to pickle.
4495 * Then use setstate to set the fillvalue
4496 */
4497 int i;
4498 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004499
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004500 if (args == NULL)
4501 return NULL;
4502 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4503 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4504 if (elem == NULL) {
4505 elem = PyTuple_New(0);
4506 if (elem == NULL) {
4507 Py_DECREF(args);
4508 return NULL;
4509 }
4510 } else
4511 Py_INCREF(elem);
4512 PyTuple_SET_ITEM(args, i, elem);
4513 }
4514 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4515}
4516
4517static PyObject *
4518zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4519{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004520 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004521 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004522 Py_RETURN_NONE;
4523}
4524
4525static PyMethodDef zip_longest_methods[] = {
4526 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4527 reduce_doc},
4528 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4529 setstate_doc},
4530 {NULL, NULL} /* sentinel */
4531};
4532
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004533PyDoc_STRVAR(zip_longest_doc,
4534"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004535\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004536Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004537the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004538method continues until the longest iterable in the argument sequence\n\
4539is exhausted and then it raises StopIteration. When the shorter iterables\n\
4540are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4541defaults to None or can be specified by a keyword argument.\n\
4542");
4543
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004544static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004545 PyVarObject_HEAD_INIT(NULL, 0)
4546 "itertools.zip_longest", /* tp_name */
4547 sizeof(ziplongestobject), /* tp_basicsize */
4548 0, /* tp_itemsize */
4549 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004550 (destructor)zip_longest_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004551 0, /* tp_print */
4552 0, /* tp_getattr */
4553 0, /* tp_setattr */
4554 0, /* tp_reserved */
4555 0, /* tp_repr */
4556 0, /* tp_as_number */
4557 0, /* tp_as_sequence */
4558 0, /* tp_as_mapping */
4559 0, /* tp_hash */
4560 0, /* tp_call */
4561 0, /* tp_str */
4562 PyObject_GenericGetAttr, /* tp_getattro */
4563 0, /* tp_setattro */
4564 0, /* tp_as_buffer */
4565 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4566 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004567 zip_longest_doc, /* tp_doc */
4568 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004569 0, /* tp_clear */
4570 0, /* tp_richcompare */
4571 0, /* tp_weaklistoffset */
4572 PyObject_SelfIter, /* tp_iter */
4573 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004574 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004575 0, /* tp_members */
4576 0, /* tp_getset */
4577 0, /* tp_base */
4578 0, /* tp_dict */
4579 0, /* tp_descr_get */
4580 0, /* tp_descr_set */
4581 0, /* tp_dictoffset */
4582 0, /* tp_init */
4583 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004584 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004585 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004586};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004587
4588/* module level code ********************************************************/
4589
4590PyDoc_STRVAR(module_doc,
4591"Functional tools for creating and using iterators.\n\
4592\n\
4593Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004594count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004595cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004596repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004597\n\
4598Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004599accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004600chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger8df58f72013-09-09 01:29:40 -05004601chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004602compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4603dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4604groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004605filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004606islice(seq, [start,] stop [, step]) --> elements from\n\
4607 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004608starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004609tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004610takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004611zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004612\n\
4613Combinatoric generators:\n\
4614product(p, q, ... [repeat=1]) --> cartesian product\n\
4615permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004616combinations(p, r)\n\
4617combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004618");
4619
4620
Raymond Hettingerad983e72003-11-12 14:32:26 +00004621static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004622 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4623 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004624};
4625
Martin v. Löwis1a214512008-06-11 05:26:20 +00004626
4627static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004628 PyModuleDef_HEAD_INIT,
4629 "itertools",
4630 module_doc,
4631 -1,
4632 module_methods,
4633 NULL,
4634 NULL,
4635 NULL,
4636 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004637};
4638
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004639PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004640PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004641{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004642 int i;
4643 PyObject *m;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004644 const char *name;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004645 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004646 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004647 &combinations_type,
4648 &cwr_type,
4649 &cycle_type,
4650 &dropwhile_type,
4651 &takewhile_type,
4652 &islice_type,
4653 &starmap_type,
4654 &chain_type,
4655 &compress_type,
4656 &filterfalse_type,
4657 &count_type,
4658 &ziplongest_type,
4659 &permutations_type,
4660 &product_type,
4661 &repeat_type,
4662 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004663 &_grouper_type,
4664 &tee_type,
4665 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004666 NULL
4667 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004668
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004669 Py_TYPE(&teedataobject_type) = &PyType_Type;
4670 m = PyModule_Create(&itertoolsmodule);
4671 if (m == NULL)
4672 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004673
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004674 for (i=0 ; typelist[i] != NULL ; i++) {
4675 if (PyType_Ready(typelist[i]) < 0)
4676 return NULL;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004677 name = _PyType_Name(typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004678 Py_INCREF(typelist[i]);
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004679 PyModule_AddObject(m, name, (PyObject *)typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004680 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004681
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004682 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004683}