blob: 8a36755bfa72657d06b1833c2f1803d7f3ba468d [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 *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530141groupby_reduce(groupbyobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000142{
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 *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530323_grouper_reduce(_grouperobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000324{
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 *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530507teedataobject_reduce(teedataobject *tdo, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000508{
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 *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530652tee_copy(teeobject *to, PyObject *Py_UNUSED(ignored))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000653{
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)) {
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530679 to = (teeobject *)tee_copy((teeobject *)it, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000680 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 *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530729tee_reduce(teeobject *to, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000730{
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;
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200813 PyObject *it, *iterable, *copyable, *copyfunc, *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 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200832
Serhiy Storchakaf320be72018-01-25 10:49:40 +0200833 if (_PyObject_LookupAttrId(it, &PyId___copy__, &copyfunc) < 0) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200834 Py_DECREF(it);
835 Py_DECREF(result);
836 return NULL;
837 }
Serhiy Storchakaf320be72018-01-25 10:49:40 +0200838 if (copyfunc != NULL) {
839 copyable = it;
840 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200841 else {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000842 copyable = tee_fromiterable(it);
843 Py_DECREF(it);
844 if (copyable == NULL) {
845 Py_DECREF(result);
846 return NULL;
847 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200848 copyfunc = _PyObject_GetAttrId(copyable, &PyId___copy__);
849 if (copyfunc == NULL) {
850 Py_DECREF(copyable);
851 Py_DECREF(result);
852 return NULL;
853 }
854 }
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200855
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200856 PyTuple_SET_ITEM(result, 0, copyable);
857 for (i = 1; i < n; i++) {
858 copyable = _PyObject_CallNoArg(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000859 if (copyable == NULL) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200860 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000861 Py_DECREF(result);
862 return NULL;
863 }
864 PyTuple_SET_ITEM(result, i, copyable);
865 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200866 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000867 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000868}
869
870PyDoc_STRVAR(tee_doc,
871"tee(iterable, n=2) --> tuple of n independent iterators.");
872
873
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700874/* cycle object **************************************************************/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000875
876typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000877 PyObject_HEAD
878 PyObject *it;
879 PyObject *saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700880 Py_ssize_t index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000881 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000882} cycleobject;
883
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000884static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000885
886static PyObject *
887cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
888{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000889 PyObject *it;
890 PyObject *iterable;
891 PyObject *saved;
892 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000893
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +0300894 if (type == &cycle_type && !_PyArg_NoKeywords("cycle", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000895 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000896
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000897 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
898 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000899
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000900 /* Get iterator. */
901 it = PyObject_GetIter(iterable);
902 if (it == NULL)
903 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000904
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000905 saved = PyList_New(0);
906 if (saved == NULL) {
907 Py_DECREF(it);
908 return NULL;
909 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000910
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000911 /* create cycleobject structure */
912 lz = (cycleobject *)type->tp_alloc(type, 0);
913 if (lz == NULL) {
914 Py_DECREF(it);
915 Py_DECREF(saved);
916 return NULL;
917 }
918 lz->it = it;
919 lz->saved = saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700920 lz->index = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000921 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000922
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000923 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000924}
925
926static void
927cycle_dealloc(cycleobject *lz)
928{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000929 PyObject_GC_UnTrack(lz);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000930 Py_XDECREF(lz->it);
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700931 Py_XDECREF(lz->saved);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000932 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000933}
934
935static int
936cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
937{
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700938 if (lz->it)
939 Py_VISIT(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000940 Py_VISIT(lz->saved);
941 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000942}
943
944static PyObject *
945cycle_next(cycleobject *lz)
946{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000947 PyObject *item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000948
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700949 if (lz->it != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000950 item = PyIter_Next(lz->it);
951 if (item != NULL) {
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700952 if (lz->firstpass)
953 return item;
954 if (PyList_Append(lz->saved, item)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000955 Py_DECREF(item);
956 return NULL;
957 }
958 return item;
959 }
Raymond Hettinger98958fe2015-08-15 15:09:30 -0700960 /* Note: StopIteration is already cleared by PyIter_Next() */
961 if (PyErr_Occurred())
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000962 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700963 Py_CLEAR(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000964 }
Serhiy Storchakafff9a312017-03-21 08:53:25 +0200965 if (PyList_GET_SIZE(lz->saved) == 0)
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700966 return NULL;
967 item = PyList_GET_ITEM(lz->saved, lz->index);
968 lz->index++;
Serhiy Storchakafff9a312017-03-21 08:53:25 +0200969 if (lz->index >= PyList_GET_SIZE(lz->saved))
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700970 lz->index = 0;
971 Py_INCREF(item);
972 return item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000973}
974
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000975static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530976cycle_reduce(cycleobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000977{
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700978 /* Create a new cycle with the iterator tuple, then set the saved state */
979 if (lz->it == NULL) {
980 PyObject *it = PyObject_GetIter(lz->saved);
981 if (it == NULL)
982 return NULL;
983 if (lz->index != 0) {
984 _Py_IDENTIFIER(__setstate__);
985 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
986 "n", lz->index);
987 if (res == NULL) {
988 Py_DECREF(it);
989 return NULL;
990 }
991 Py_DECREF(res);
992 }
993 return Py_BuildValue("O(N)(Oi)", Py_TYPE(lz), it, lz->saved, 1);
994 }
995 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), lz->it, lz->saved,
996 lz->firstpass);
Raymond Hettingerc39786d2015-08-15 15:16:12 -0700997}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000998
999static PyObject *
1000cycle_setstate(cycleobject *lz, PyObject *state)
1001{
1002 PyObject *saved=NULL;
1003 int firstpass;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001004 if (!PyTuple_Check(state)) {
1005 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001006 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001007 }
1008 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
1009 return NULL;
1010 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001011 Py_INCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001012 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001013 lz->firstpass = firstpass != 0;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001014 lz->index = 0;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001015 Py_RETURN_NONE;
1016}
1017
1018static PyMethodDef cycle_methods[] = {
1019 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1020 reduce_doc},
1021 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1022 setstate_doc},
1023 {NULL, NULL} /* sentinel */
1024};
1025
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001026PyDoc_STRVAR(cycle_doc,
1027"cycle(iterable) --> cycle object\n\
1028\n\
1029Return elements from the iterable until it is exhausted.\n\
1030Then repeat the sequence indefinitely.");
1031
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001032static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001033 PyVarObject_HEAD_INIT(NULL, 0)
1034 "itertools.cycle", /* tp_name */
1035 sizeof(cycleobject), /* tp_basicsize */
1036 0, /* tp_itemsize */
1037 /* methods */
1038 (destructor)cycle_dealloc, /* tp_dealloc */
1039 0, /* tp_print */
1040 0, /* tp_getattr */
1041 0, /* tp_setattr */
1042 0, /* tp_reserved */
1043 0, /* tp_repr */
1044 0, /* tp_as_number */
1045 0, /* tp_as_sequence */
1046 0, /* tp_as_mapping */
1047 0, /* tp_hash */
1048 0, /* tp_call */
1049 0, /* tp_str */
1050 PyObject_GenericGetAttr, /* tp_getattro */
1051 0, /* tp_setattro */
1052 0, /* tp_as_buffer */
1053 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1054 Py_TPFLAGS_BASETYPE, /* tp_flags */
1055 cycle_doc, /* tp_doc */
1056 (traverseproc)cycle_traverse, /* tp_traverse */
1057 0, /* tp_clear */
1058 0, /* tp_richcompare */
1059 0, /* tp_weaklistoffset */
1060 PyObject_SelfIter, /* tp_iter */
1061 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001062 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001063 0, /* tp_members */
1064 0, /* tp_getset */
1065 0, /* tp_base */
1066 0, /* tp_dict */
1067 0, /* tp_descr_get */
1068 0, /* tp_descr_set */
1069 0, /* tp_dictoffset */
1070 0, /* tp_init */
1071 0, /* tp_alloc */
1072 cycle_new, /* tp_new */
1073 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001074};
1075
1076
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001077/* dropwhile object **********************************************************/
1078
1079typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001080 PyObject_HEAD
1081 PyObject *func;
1082 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001083 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001084} dropwhileobject;
1085
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001086static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001087
1088static PyObject *
1089dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1090{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001091 PyObject *func, *seq;
1092 PyObject *it;
1093 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001094
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001095 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001096 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001097
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001098 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
1099 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001100
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001101 /* Get iterator. */
1102 it = PyObject_GetIter(seq);
1103 if (it == NULL)
1104 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001105
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001106 /* create dropwhileobject structure */
1107 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1108 if (lz == NULL) {
1109 Py_DECREF(it);
1110 return NULL;
1111 }
1112 Py_INCREF(func);
1113 lz->func = func;
1114 lz->it = it;
1115 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001117 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001118}
1119
1120static void
1121dropwhile_dealloc(dropwhileobject *lz)
1122{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001123 PyObject_GC_UnTrack(lz);
1124 Py_XDECREF(lz->func);
1125 Py_XDECREF(lz->it);
1126 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001127}
1128
1129static int
1130dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1131{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001132 Py_VISIT(lz->it);
1133 Py_VISIT(lz->func);
1134 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001135}
1136
1137static PyObject *
1138dropwhile_next(dropwhileobject *lz)
1139{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001140 PyObject *item, *good;
1141 PyObject *it = lz->it;
1142 long ok;
1143 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001144
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001145 iternext = *Py_TYPE(it)->tp_iternext;
1146 for (;;) {
1147 item = iternext(it);
1148 if (item == NULL)
1149 return NULL;
1150 if (lz->start == 1)
1151 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001152
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01001153 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001154 if (good == NULL) {
1155 Py_DECREF(item);
1156 return NULL;
1157 }
1158 ok = PyObject_IsTrue(good);
1159 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001160 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001161 lz->start = 1;
1162 return item;
1163 }
1164 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001165 if (ok < 0)
1166 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001167 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001168}
1169
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001170static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301171dropwhile_reduce(dropwhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001172{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001173 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 +00001174}
1175
1176static PyObject *
1177dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1178{
1179 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001180 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001181 return NULL;
1182 lz->start = start;
1183 Py_RETURN_NONE;
1184}
1185
1186static PyMethodDef dropwhile_methods[] = {
1187 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1188 reduce_doc},
1189 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1190 setstate_doc},
1191 {NULL, NULL} /* sentinel */
1192};
1193
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001194PyDoc_STRVAR(dropwhile_doc,
1195"dropwhile(predicate, iterable) --> dropwhile object\n\
1196\n\
1197Drop items from the iterable while predicate(item) is true.\n\
1198Afterwards, return every element until the iterable is exhausted.");
1199
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001200static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001201 PyVarObject_HEAD_INIT(NULL, 0)
1202 "itertools.dropwhile", /* tp_name */
1203 sizeof(dropwhileobject), /* tp_basicsize */
1204 0, /* tp_itemsize */
1205 /* methods */
1206 (destructor)dropwhile_dealloc, /* tp_dealloc */
1207 0, /* tp_print */
1208 0, /* tp_getattr */
1209 0, /* tp_setattr */
1210 0, /* tp_reserved */
1211 0, /* tp_repr */
1212 0, /* tp_as_number */
1213 0, /* tp_as_sequence */
1214 0, /* tp_as_mapping */
1215 0, /* tp_hash */
1216 0, /* tp_call */
1217 0, /* tp_str */
1218 PyObject_GenericGetAttr, /* tp_getattro */
1219 0, /* tp_setattro */
1220 0, /* tp_as_buffer */
1221 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1222 Py_TPFLAGS_BASETYPE, /* tp_flags */
1223 dropwhile_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001224 (traverseproc)dropwhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001225 0, /* tp_clear */
1226 0, /* tp_richcompare */
1227 0, /* tp_weaklistoffset */
1228 PyObject_SelfIter, /* tp_iter */
1229 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001230 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001231 0, /* tp_members */
1232 0, /* tp_getset */
1233 0, /* tp_base */
1234 0, /* tp_dict */
1235 0, /* tp_descr_get */
1236 0, /* tp_descr_set */
1237 0, /* tp_dictoffset */
1238 0, /* tp_init */
1239 0, /* tp_alloc */
1240 dropwhile_new, /* tp_new */
1241 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001242};
1243
1244
1245/* takewhile object **********************************************************/
1246
1247typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001248 PyObject_HEAD
1249 PyObject *func;
1250 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001251 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001252} takewhileobject;
1253
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001254static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001255
1256static PyObject *
1257takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1258{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001259 PyObject *func, *seq;
1260 PyObject *it;
1261 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001262
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001263 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001264 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001265
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001266 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1267 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001268
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001269 /* Get iterator. */
1270 it = PyObject_GetIter(seq);
1271 if (it == NULL)
1272 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001273
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001274 /* create takewhileobject structure */
1275 lz = (takewhileobject *)type->tp_alloc(type, 0);
1276 if (lz == NULL) {
1277 Py_DECREF(it);
1278 return NULL;
1279 }
1280 Py_INCREF(func);
1281 lz->func = func;
1282 lz->it = it;
1283 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001284
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001285 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001286}
1287
1288static void
1289takewhile_dealloc(takewhileobject *lz)
1290{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001291 PyObject_GC_UnTrack(lz);
1292 Py_XDECREF(lz->func);
1293 Py_XDECREF(lz->it);
1294 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001295}
1296
1297static int
1298takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1299{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001300 Py_VISIT(lz->it);
1301 Py_VISIT(lz->func);
1302 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001303}
1304
1305static PyObject *
1306takewhile_next(takewhileobject *lz)
1307{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001308 PyObject *item, *good;
1309 PyObject *it = lz->it;
1310 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001311
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001312 if (lz->stop == 1)
1313 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001314
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001315 item = (*Py_TYPE(it)->tp_iternext)(it);
1316 if (item == NULL)
1317 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001318
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01001319 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001320 if (good == NULL) {
1321 Py_DECREF(item);
1322 return NULL;
1323 }
1324 ok = PyObject_IsTrue(good);
1325 Py_DECREF(good);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001326 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001327 return item;
1328 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001329 if (ok == 0)
1330 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001331 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001332}
1333
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001334static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301335takewhile_reduce(takewhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001336{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001337 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 +00001338}
1339
1340static PyObject *
1341takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1342{
1343 int stop = PyObject_IsTrue(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001344
Antoine Pitrou721738f2012-08-15 23:20:39 +02001345 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001346 return NULL;
1347 lz->stop = stop;
1348 Py_RETURN_NONE;
1349}
1350
1351static PyMethodDef takewhile_reduce_methods[] = {
1352 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1353 reduce_doc},
1354 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1355 setstate_doc},
1356 {NULL, NULL} /* sentinel */
1357};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001358PyDoc_STRVAR(takewhile_doc,
1359"takewhile(predicate, iterable) --> takewhile object\n\
1360\n\
oldkaa0735f2018-02-02 16:52:55 +08001361Return successive entries from an iterable as long as the\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001362predicate evaluates to true for each entry.");
1363
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001364static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001365 PyVarObject_HEAD_INIT(NULL, 0)
1366 "itertools.takewhile", /* tp_name */
1367 sizeof(takewhileobject), /* tp_basicsize */
1368 0, /* tp_itemsize */
1369 /* methods */
1370 (destructor)takewhile_dealloc, /* tp_dealloc */
1371 0, /* tp_print */
1372 0, /* tp_getattr */
1373 0, /* tp_setattr */
1374 0, /* tp_reserved */
1375 0, /* tp_repr */
1376 0, /* tp_as_number */
1377 0, /* tp_as_sequence */
1378 0, /* tp_as_mapping */
1379 0, /* tp_hash */
1380 0, /* tp_call */
1381 0, /* tp_str */
1382 PyObject_GenericGetAttr, /* tp_getattro */
1383 0, /* tp_setattro */
1384 0, /* tp_as_buffer */
1385 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1386 Py_TPFLAGS_BASETYPE, /* tp_flags */
1387 takewhile_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001388 (traverseproc)takewhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001389 0, /* tp_clear */
1390 0, /* tp_richcompare */
1391 0, /* tp_weaklistoffset */
1392 PyObject_SelfIter, /* tp_iter */
1393 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001394 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001395 0, /* tp_members */
1396 0, /* tp_getset */
1397 0, /* tp_base */
1398 0, /* tp_dict */
1399 0, /* tp_descr_get */
1400 0, /* tp_descr_set */
1401 0, /* tp_dictoffset */
1402 0, /* tp_init */
1403 0, /* tp_alloc */
1404 takewhile_new, /* tp_new */
1405 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001406};
1407
1408
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001409/* islice object *************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001410
1411typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001412 PyObject_HEAD
1413 PyObject *it;
1414 Py_ssize_t next;
1415 Py_ssize_t stop;
1416 Py_ssize_t step;
1417 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001418} isliceobject;
1419
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001420static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001421
1422static PyObject *
1423islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1424{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001425 PyObject *seq;
1426 Py_ssize_t start=0, stop=-1, step=1;
1427 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1428 Py_ssize_t numargs;
1429 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001430
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001431 if (type == &islice_type && !_PyArg_NoKeywords("islice", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001432 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001433
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001434 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1435 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001436
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001437 numargs = PyTuple_Size(args);
1438 if (numargs == 2) {
1439 if (a1 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001440 stop = PyNumber_AsSsize_t(a1, 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 } else {
1451 if (a1 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001452 start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001453 if (start == -1 && PyErr_Occurred())
1454 PyErr_Clear();
1455 if (a2 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001456 stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001457 if (stop == -1) {
1458 if (PyErr_Occurred())
1459 PyErr_Clear();
1460 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001461 "Stop argument for islice() must be None or "
1462 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001463 return NULL;
1464 }
1465 }
1466 }
1467 if (start<0 || stop<-1) {
1468 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001469 "Indices for islice() must be None or "
1470 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001471 return NULL;
1472 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001473
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001474 if (a3 != NULL) {
1475 if (a3 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001476 step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001477 if (step == -1 && PyErr_Occurred())
1478 PyErr_Clear();
1479 }
1480 if (step<1) {
1481 PyErr_SetString(PyExc_ValueError,
1482 "Step for islice() must be a positive integer or None.");
1483 return NULL;
1484 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001485
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001486 /* Get iterator. */
1487 it = PyObject_GetIter(seq);
1488 if (it == NULL)
1489 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001490
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001491 /* create isliceobject structure */
1492 lz = (isliceobject *)type->tp_alloc(type, 0);
1493 if (lz == NULL) {
1494 Py_DECREF(it);
1495 return NULL;
1496 }
1497 lz->it = it;
1498 lz->next = start;
1499 lz->stop = stop;
1500 lz->step = step;
1501 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001502
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001503 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001504}
1505
1506static void
1507islice_dealloc(isliceobject *lz)
1508{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001509 PyObject_GC_UnTrack(lz);
1510 Py_XDECREF(lz->it);
1511 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001512}
1513
1514static int
1515islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1516{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001517 Py_VISIT(lz->it);
1518 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001519}
1520
1521static PyObject *
1522islice_next(isliceobject *lz)
1523{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001524 PyObject *item;
1525 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001526 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001527 Py_ssize_t oldnext;
1528 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001529
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001530 if (it == NULL)
1531 return NULL;
1532
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001533 iternext = *Py_TYPE(it)->tp_iternext;
1534 while (lz->cnt < lz->next) {
1535 item = iternext(it);
1536 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001537 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001538 Py_DECREF(item);
1539 lz->cnt++;
1540 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001541 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001542 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001543 item = iternext(it);
1544 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001545 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001546 lz->cnt++;
1547 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001548 /* The (size_t) cast below avoids the danger of undefined
1549 behaviour from signed integer overflow. */
1550 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001551 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1552 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001553 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001554
1555empty:
1556 Py_CLEAR(lz->it);
1557 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001558}
1559
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001560static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301561islice_reduce(isliceobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001562{
1563 /* When unpickled, generate a new object with the same bounds,
1564 * then 'setstate' with the next and count
1565 */
1566 PyObject *stop;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001567
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001568 if (lz->it == NULL) {
1569 PyObject *empty_list;
1570 PyObject *empty_it;
1571 empty_list = PyList_New(0);
1572 if (empty_list == NULL)
1573 return NULL;
1574 empty_it = PyObject_GetIter(empty_list);
1575 Py_DECREF(empty_list);
1576 if (empty_it == NULL)
1577 return NULL;
1578 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1579 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001580 if (lz->stop == -1) {
1581 stop = Py_None;
1582 Py_INCREF(stop);
1583 } else {
1584 stop = PyLong_FromSsize_t(lz->stop);
1585 if (stop == NULL)
1586 return NULL;
1587 }
1588 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1589 lz->it, lz->next, stop, lz->step,
1590 lz->cnt);
1591}
1592
1593static PyObject *
1594islice_setstate(isliceobject *lz, PyObject *state)
1595{
1596 Py_ssize_t cnt = PyLong_AsSsize_t(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001597
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001598 if (cnt == -1 && PyErr_Occurred())
1599 return NULL;
1600 lz->cnt = cnt;
1601 Py_RETURN_NONE;
1602}
1603
1604static PyMethodDef islice_methods[] = {
1605 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1606 reduce_doc},
1607 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1608 setstate_doc},
1609 {NULL, NULL} /* sentinel */
1610};
1611
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001612PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001613"islice(iterable, stop) --> islice object\n\
1614islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001615\n\
1616Return an iterator whose next() method returns selected values from an\n\
1617iterable. If start is specified, will skip all preceding elements;\n\
1618otherwise, start defaults to zero. Step defaults to one. If\n\
oldkaa0735f2018-02-02 16:52:55 +08001619specified as another value, step determines how many values are\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001620skipped between successive calls. Works like a slice() on a list\n\
1621but returns an iterator.");
1622
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001623static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001624 PyVarObject_HEAD_INIT(NULL, 0)
1625 "itertools.islice", /* tp_name */
1626 sizeof(isliceobject), /* tp_basicsize */
1627 0, /* tp_itemsize */
1628 /* methods */
1629 (destructor)islice_dealloc, /* tp_dealloc */
1630 0, /* tp_print */
1631 0, /* tp_getattr */
1632 0, /* tp_setattr */
1633 0, /* tp_reserved */
1634 0, /* tp_repr */
1635 0, /* tp_as_number */
1636 0, /* tp_as_sequence */
1637 0, /* tp_as_mapping */
1638 0, /* tp_hash */
1639 0, /* tp_call */
1640 0, /* tp_str */
1641 PyObject_GenericGetAttr, /* tp_getattro */
1642 0, /* tp_setattro */
1643 0, /* tp_as_buffer */
1644 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1645 Py_TPFLAGS_BASETYPE, /* tp_flags */
1646 islice_doc, /* tp_doc */
1647 (traverseproc)islice_traverse, /* tp_traverse */
1648 0, /* tp_clear */
1649 0, /* tp_richcompare */
1650 0, /* tp_weaklistoffset */
1651 PyObject_SelfIter, /* tp_iter */
1652 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001653 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001654 0, /* tp_members */
1655 0, /* tp_getset */
1656 0, /* tp_base */
1657 0, /* tp_dict */
1658 0, /* tp_descr_get */
1659 0, /* tp_descr_set */
1660 0, /* tp_dictoffset */
1661 0, /* tp_init */
1662 0, /* tp_alloc */
1663 islice_new, /* tp_new */
1664 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001665};
1666
1667
1668/* starmap object ************************************************************/
1669
1670typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001671 PyObject_HEAD
1672 PyObject *func;
1673 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001674} starmapobject;
1675
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001676static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001677
1678static PyObject *
1679starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1680{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001681 PyObject *func, *seq;
1682 PyObject *it;
1683 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001684
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001685 if (type == &starmap_type && !_PyArg_NoKeywords("starmap", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001686 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001687
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001688 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1689 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001690
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001691 /* Get iterator. */
1692 it = PyObject_GetIter(seq);
1693 if (it == NULL)
1694 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001695
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001696 /* create starmapobject structure */
1697 lz = (starmapobject *)type->tp_alloc(type, 0);
1698 if (lz == NULL) {
1699 Py_DECREF(it);
1700 return NULL;
1701 }
1702 Py_INCREF(func);
1703 lz->func = func;
1704 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001705
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001706 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001707}
1708
1709static void
1710starmap_dealloc(starmapobject *lz)
1711{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001712 PyObject_GC_UnTrack(lz);
1713 Py_XDECREF(lz->func);
1714 Py_XDECREF(lz->it);
1715 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001716}
1717
1718static int
1719starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1720{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001721 Py_VISIT(lz->it);
1722 Py_VISIT(lz->func);
1723 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001724}
1725
1726static PyObject *
1727starmap_next(starmapobject *lz)
1728{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001729 PyObject *args;
1730 PyObject *result;
1731 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001732
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001733 args = (*Py_TYPE(it)->tp_iternext)(it);
1734 if (args == NULL)
1735 return NULL;
1736 if (!PyTuple_CheckExact(args)) {
1737 PyObject *newargs = PySequence_Tuple(args);
1738 Py_DECREF(args);
1739 if (newargs == NULL)
1740 return NULL;
1741 args = newargs;
1742 }
1743 result = PyObject_Call(lz->func, args, NULL);
1744 Py_DECREF(args);
1745 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001746}
1747
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001748static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301749starmap_reduce(starmapobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001750{
1751 /* Just pickle the iterator */
1752 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1753}
1754
1755static PyMethodDef starmap_methods[] = {
1756 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1757 reduce_doc},
1758 {NULL, NULL} /* sentinel */
1759};
1760
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001761PyDoc_STRVAR(starmap_doc,
1762"starmap(function, sequence) --> starmap object\n\
1763\n\
1764Return an iterator whose values are returned from the function evaluated\n\
Serhiy Storchakad65c9492015-11-02 14:10:23 +02001765with an argument tuple taken from the given sequence.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001766
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001767static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001768 PyVarObject_HEAD_INIT(NULL, 0)
1769 "itertools.starmap", /* tp_name */
1770 sizeof(starmapobject), /* tp_basicsize */
1771 0, /* tp_itemsize */
1772 /* methods */
1773 (destructor)starmap_dealloc, /* tp_dealloc */
1774 0, /* tp_print */
1775 0, /* tp_getattr */
1776 0, /* tp_setattr */
1777 0, /* tp_reserved */
1778 0, /* tp_repr */
1779 0, /* tp_as_number */
1780 0, /* tp_as_sequence */
1781 0, /* tp_as_mapping */
1782 0, /* tp_hash */
1783 0, /* tp_call */
1784 0, /* tp_str */
1785 PyObject_GenericGetAttr, /* tp_getattro */
1786 0, /* tp_setattro */
1787 0, /* tp_as_buffer */
1788 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1789 Py_TPFLAGS_BASETYPE, /* tp_flags */
1790 starmap_doc, /* tp_doc */
1791 (traverseproc)starmap_traverse, /* tp_traverse */
1792 0, /* tp_clear */
1793 0, /* tp_richcompare */
1794 0, /* tp_weaklistoffset */
1795 PyObject_SelfIter, /* tp_iter */
1796 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001797 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001798 0, /* tp_members */
1799 0, /* tp_getset */
1800 0, /* tp_base */
1801 0, /* tp_dict */
1802 0, /* tp_descr_get */
1803 0, /* tp_descr_set */
1804 0, /* tp_dictoffset */
1805 0, /* tp_init */
1806 0, /* tp_alloc */
1807 starmap_new, /* tp_new */
1808 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001809};
1810
1811
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001812/* chain object **************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001813
1814typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001815 PyObject_HEAD
1816 PyObject *source; /* Iterator over input iterables */
1817 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001818} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001819
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001820static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001821
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001822static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001823chain_new_internal(PyTypeObject *type, PyObject *source)
1824{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001825 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001826
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001827 lz = (chainobject *)type->tp_alloc(type, 0);
1828 if (lz == NULL) {
1829 Py_DECREF(source);
1830 return NULL;
1831 }
1832
1833 lz->source = source;
1834 lz->active = NULL;
1835 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001836}
1837
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001838static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001839chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001840{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001841 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001842
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001843 if (type == &chain_type && !_PyArg_NoKeywords("chain", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001844 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001845
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001846 source = PyObject_GetIter(args);
1847 if (source == NULL)
1848 return NULL;
1849
1850 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001851}
1852
1853static PyObject *
1854chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1855{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001856 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001857
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001858 source = PyObject_GetIter(arg);
1859 if (source == NULL)
1860 return NULL;
1861
1862 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001863}
1864
1865static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001866chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001867{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001868 PyObject_GC_UnTrack(lz);
1869 Py_XDECREF(lz->active);
1870 Py_XDECREF(lz->source);
1871 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001872}
1873
Raymond Hettinger2012f172003-02-07 05:32:58 +00001874static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001875chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001876{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001877 Py_VISIT(lz->source);
1878 Py_VISIT(lz->active);
1879 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001880}
1881
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001882static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001883chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001884{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001885 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001886
T. Wouters5466d4a2017-03-30 09:58:35 -07001887 /* lz->source is the iterator of iterables. If it's NULL, we've already
1888 * consumed them all. lz->active is the current iterator. If it's NULL,
1889 * we should grab a new one from lz->source. */
1890 while (lz->source != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001891 if (lz->active == NULL) {
T. Wouters5466d4a2017-03-30 09:58:35 -07001892 PyObject *iterable = PyIter_Next(lz->source);
1893 if (iterable == NULL) {
1894 Py_CLEAR(lz->source);
1895 return NULL; /* no more input sources */
1896 }
1897 lz->active = PyObject_GetIter(iterable);
1898 Py_DECREF(iterable);
1899 if (lz->active == NULL) {
1900 Py_CLEAR(lz->source);
1901 return NULL; /* input not iterable */
1902 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001903 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001904 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1905 if (item != NULL)
1906 return item;
1907 if (PyErr_Occurred()) {
1908 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1909 PyErr_Clear();
1910 else
1911 return NULL; /* input raised an exception */
1912 }
1913 /* lz->active is consumed, try with the next iterable. */
1914 Py_CLEAR(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001915 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001916 /* Everything had been consumed already. */
1917 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001918}
1919
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001920static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301921chain_reduce(chainobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001922{
1923 if (lz->source) {
1924 /* we can't pickle function objects (itertools.from_iterable) so
1925 * we must use setstate to replace the iterable. One day we
1926 * will fix pickling of functions
1927 */
1928 if (lz->active) {
1929 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1930 } else {
1931 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1932 }
1933 } else {
1934 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1935 }
1936 return NULL;
1937}
1938
1939static PyObject *
1940chain_setstate(chainobject *lz, PyObject *state)
1941{
1942 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001943
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001944 if (!PyTuple_Check(state)) {
1945 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001946 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001947 }
1948 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
1949 return NULL;
1950 }
1951 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
1952 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
1953 return NULL;
1954 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001955
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02001956 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001957 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02001958 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001959 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001960 Py_RETURN_NONE;
1961}
1962
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001963PyDoc_STRVAR(chain_doc,
1964"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001965\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001966Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001967first iterable until it is exhausted, then elements from the next\n\
1968iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001969
Christian Heimesf16baeb2008-02-29 14:57:44 +00001970PyDoc_STRVAR(chain_from_iterable_doc,
1971"chain.from_iterable(iterable) --> chain object\n\
1972\n\
Martin Pantereb995702016-07-28 01:11:04 +00001973Alternate chain() constructor taking a single iterable argument\n\
Christian Heimesf16baeb2008-02-29 14:57:44 +00001974that evaluates lazily.");
1975
1976static PyMethodDef chain_methods[] = {
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001977 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1978 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001979 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1980 reduce_doc},
1981 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1982 setstate_doc},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001983 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001984};
1985
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001986static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001987 PyVarObject_HEAD_INIT(NULL, 0)
1988 "itertools.chain", /* tp_name */
1989 sizeof(chainobject), /* tp_basicsize */
1990 0, /* tp_itemsize */
1991 /* methods */
1992 (destructor)chain_dealloc, /* tp_dealloc */
1993 0, /* tp_print */
1994 0, /* tp_getattr */
1995 0, /* tp_setattr */
1996 0, /* tp_reserved */
1997 0, /* tp_repr */
1998 0, /* tp_as_number */
1999 0, /* tp_as_sequence */
2000 0, /* tp_as_mapping */
2001 0, /* tp_hash */
2002 0, /* tp_call */
2003 0, /* tp_str */
2004 PyObject_GenericGetAttr, /* tp_getattro */
2005 0, /* tp_setattro */
2006 0, /* tp_as_buffer */
2007 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2008 Py_TPFLAGS_BASETYPE, /* tp_flags */
2009 chain_doc, /* tp_doc */
2010 (traverseproc)chain_traverse, /* tp_traverse */
2011 0, /* tp_clear */
2012 0, /* tp_richcompare */
2013 0, /* tp_weaklistoffset */
2014 PyObject_SelfIter, /* tp_iter */
2015 (iternextfunc)chain_next, /* tp_iternext */
2016 chain_methods, /* tp_methods */
2017 0, /* tp_members */
2018 0, /* tp_getset */
2019 0, /* tp_base */
2020 0, /* tp_dict */
2021 0, /* tp_descr_get */
2022 0, /* tp_descr_set */
2023 0, /* tp_dictoffset */
2024 0, /* tp_init */
2025 0, /* tp_alloc */
2026 chain_new, /* tp_new */
2027 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002028};
2029
2030
Christian Heimesc3f30c42008-02-22 16:37:40 +00002031/* product object ************************************************************/
2032
2033typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002034 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002035 PyObject *pools; /* tuple of pool tuples */
2036 Py_ssize_t *indices; /* one index per pool */
2037 PyObject *result; /* most recently returned result tuple */
2038 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002039} productobject;
2040
2041static PyTypeObject product_type;
2042
2043static PyObject *
2044product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2045{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002046 productobject *lz;
2047 Py_ssize_t nargs, npools, repeat=1;
2048 PyObject *pools = NULL;
2049 Py_ssize_t *indices = NULL;
2050 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002051
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002052 if (kwds != NULL) {
2053 char *kwlist[] = {"repeat", 0};
2054 PyObject *tmpargs = PyTuple_New(0);
2055 if (tmpargs == NULL)
2056 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002057 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2058 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002059 Py_DECREF(tmpargs);
2060 return NULL;
2061 }
2062 Py_DECREF(tmpargs);
2063 if (repeat < 0) {
2064 PyErr_SetString(PyExc_ValueError,
2065 "repeat argument cannot be negative");
2066 return NULL;
2067 }
2068 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002069
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002070 assert(PyTuple_CheckExact(args));
2071 if (repeat == 0) {
2072 nargs = 0;
2073 } else {
2074 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002075 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002076 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2077 return NULL;
2078 }
2079 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002080 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002081
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002082 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002083 if (indices == NULL) {
2084 PyErr_NoMemory();
2085 goto error;
2086 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002087
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002088 pools = PyTuple_New(npools);
2089 if (pools == NULL)
2090 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002091
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002092 for (i=0; i < nargs ; ++i) {
2093 PyObject *item = PyTuple_GET_ITEM(args, i);
2094 PyObject *pool = PySequence_Tuple(item);
2095 if (pool == NULL)
2096 goto error;
2097 PyTuple_SET_ITEM(pools, i, pool);
2098 indices[i] = 0;
2099 }
2100 for ( ; i < npools; ++i) {
2101 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2102 Py_INCREF(pool);
2103 PyTuple_SET_ITEM(pools, i, pool);
2104 indices[i] = 0;
2105 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002106
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002107 /* create productobject structure */
2108 lz = (productobject *)type->tp_alloc(type, 0);
2109 if (lz == NULL)
2110 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002111
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002112 lz->pools = pools;
2113 lz->indices = indices;
2114 lz->result = NULL;
2115 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002116
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002117 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002118
2119error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002120 if (indices != NULL)
2121 PyMem_Free(indices);
2122 Py_XDECREF(pools);
2123 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002124}
2125
2126static void
2127product_dealloc(productobject *lz)
2128{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002129 PyObject_GC_UnTrack(lz);
2130 Py_XDECREF(lz->pools);
2131 Py_XDECREF(lz->result);
2132 if (lz->indices != NULL)
2133 PyMem_Free(lz->indices);
2134 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002135}
2136
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002137static PyObject *
2138product_sizeof(productobject *lz, void *unused)
2139{
2140 Py_ssize_t res;
2141
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002142 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002143 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2144 return PyLong_FromSsize_t(res);
2145}
2146
2147PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2148
Christian Heimesc3f30c42008-02-22 16:37:40 +00002149static int
2150product_traverse(productobject *lz, visitproc visit, void *arg)
2151{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002152 Py_VISIT(lz->pools);
2153 Py_VISIT(lz->result);
2154 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002155}
2156
2157static PyObject *
2158product_next(productobject *lz)
2159{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002160 PyObject *pool;
2161 PyObject *elem;
2162 PyObject *oldelem;
2163 PyObject *pools = lz->pools;
2164 PyObject *result = lz->result;
2165 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2166 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002167
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002168 if (lz->stopped)
2169 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002170
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002171 if (result == NULL) {
2172 /* On the first pass, return an initial tuple filled with the
2173 first element from each pool. */
2174 result = PyTuple_New(npools);
2175 if (result == NULL)
2176 goto empty;
2177 lz->result = result;
2178 for (i=0; i < npools; i++) {
2179 pool = PyTuple_GET_ITEM(pools, i);
2180 if (PyTuple_GET_SIZE(pool) == 0)
2181 goto empty;
2182 elem = PyTuple_GET_ITEM(pool, 0);
2183 Py_INCREF(elem);
2184 PyTuple_SET_ITEM(result, i, elem);
2185 }
2186 } else {
2187 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002188
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002189 /* Copy the previous result tuple or re-use it if available */
2190 if (Py_REFCNT(result) > 1) {
2191 PyObject *old_result = result;
2192 result = PyTuple_New(npools);
2193 if (result == NULL)
2194 goto empty;
2195 lz->result = result;
2196 for (i=0; i < npools; i++) {
2197 elem = PyTuple_GET_ITEM(old_result, i);
2198 Py_INCREF(elem);
2199 PyTuple_SET_ITEM(result, i, elem);
2200 }
2201 Py_DECREF(old_result);
2202 }
2203 /* Now, we've got the only copy so we can update it in-place */
2204 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002205
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002206 /* Update the pool indices right-to-left. Only advance to the
2207 next pool when the previous one rolls-over */
2208 for (i=npools-1 ; i >= 0 ; i--) {
2209 pool = PyTuple_GET_ITEM(pools, i);
2210 indices[i]++;
2211 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2212 /* Roll-over and advance to next pool */
2213 indices[i] = 0;
2214 elem = PyTuple_GET_ITEM(pool, 0);
2215 Py_INCREF(elem);
2216 oldelem = PyTuple_GET_ITEM(result, i);
2217 PyTuple_SET_ITEM(result, i, elem);
2218 Py_DECREF(oldelem);
2219 } else {
2220 /* No rollover. Just increment and stop here. */
2221 elem = PyTuple_GET_ITEM(pool, indices[i]);
2222 Py_INCREF(elem);
2223 oldelem = PyTuple_GET_ITEM(result, i);
2224 PyTuple_SET_ITEM(result, i, elem);
2225 Py_DECREF(oldelem);
2226 break;
2227 }
2228 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002229
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002230 /* If i is negative, then the indices have all rolled-over
2231 and we're done. */
2232 if (i < 0)
2233 goto empty;
2234 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002235
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002236 Py_INCREF(result);
2237 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002238
2239empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002240 lz->stopped = 1;
2241 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002242}
2243
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002244static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302245product_reduce(productobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002246{
2247 if (lz->stopped) {
2248 return Py_BuildValue("O(())", Py_TYPE(lz));
2249 } else if (lz->result == NULL) {
2250 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2251 } else {
2252 PyObject *indices;
2253 Py_ssize_t n, i;
2254
2255 /* we must pickle the indices use them for setstate, and
2256 * additionally indicate that the iterator has started
2257 */
2258 n = PyTuple_GET_SIZE(lz->pools);
2259 indices = PyTuple_New(n);
2260 if (indices == NULL)
2261 return NULL;
2262 for (i=0; i<n; i++){
2263 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2264 if (!index) {
2265 Py_DECREF(indices);
2266 return NULL;
2267 }
2268 PyTuple_SET_ITEM(indices, i, index);
2269 }
2270 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2271 }
2272}
2273
2274static PyObject *
2275product_setstate(productobject *lz, PyObject *state)
2276{
2277 PyObject *result;
2278 Py_ssize_t n, i;
2279
2280 n = PyTuple_GET_SIZE(lz->pools);
2281 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2282 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2283 return NULL;
2284 }
2285 for (i=0; i<n; i++)
2286 {
2287 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2288 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002289 PyObject* pool;
2290 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002291 if (index < 0 && PyErr_Occurred())
2292 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002293 pool = PyTuple_GET_ITEM(lz->pools, i);
2294 poolsize = PyTuple_GET_SIZE(pool);
2295 if (poolsize == 0) {
2296 lz->stopped = 1;
2297 Py_RETURN_NONE;
2298 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002299 /* clamp the index */
2300 if (index < 0)
2301 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002302 else if (index > poolsize-1)
2303 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002304 lz->indices[i] = index;
2305 }
2306
2307 result = PyTuple_New(n);
2308 if (!result)
2309 return NULL;
2310 for (i=0; i<n; i++) {
2311 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2312 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2313 Py_INCREF(element);
2314 PyTuple_SET_ITEM(result, i, element);
2315 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002316 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002317 Py_RETURN_NONE;
2318}
2319
2320static PyMethodDef product_methods[] = {
2321 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2322 reduce_doc},
2323 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2324 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002325 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2326 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002327 {NULL, NULL} /* sentinel */
2328};
2329
Christian Heimesc3f30c42008-02-22 16:37:40 +00002330PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002331"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002332\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002333Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002334For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2335The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2336cycle in a manner similar to an odometer (with the rightmost element changing\n\
2337on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002338To compute the product of an iterable with itself, specify the number\n\
2339of repetitions with the optional repeat keyword argument. For example,\n\
2340product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002341product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2342product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2343
2344static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002345 PyVarObject_HEAD_INIT(NULL, 0)
2346 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002347 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002348 0, /* tp_itemsize */
2349 /* methods */
2350 (destructor)product_dealloc, /* tp_dealloc */
2351 0, /* tp_print */
2352 0, /* tp_getattr */
2353 0, /* tp_setattr */
2354 0, /* tp_reserved */
2355 0, /* tp_repr */
2356 0, /* tp_as_number */
2357 0, /* tp_as_sequence */
2358 0, /* tp_as_mapping */
2359 0, /* tp_hash */
2360 0, /* tp_call */
2361 0, /* tp_str */
2362 PyObject_GenericGetAttr, /* tp_getattro */
2363 0, /* tp_setattro */
2364 0, /* tp_as_buffer */
2365 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2366 Py_TPFLAGS_BASETYPE, /* tp_flags */
2367 product_doc, /* tp_doc */
2368 (traverseproc)product_traverse, /* tp_traverse */
2369 0, /* tp_clear */
2370 0, /* tp_richcompare */
2371 0, /* tp_weaklistoffset */
2372 PyObject_SelfIter, /* tp_iter */
2373 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002374 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002375 0, /* tp_members */
2376 0, /* tp_getset */
2377 0, /* tp_base */
2378 0, /* tp_dict */
2379 0, /* tp_descr_get */
2380 0, /* tp_descr_set */
2381 0, /* tp_dictoffset */
2382 0, /* tp_init */
2383 0, /* tp_alloc */
2384 product_new, /* tp_new */
2385 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002386};
2387
2388
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002389/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002390
2391typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002392 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002393 PyObject *pool; /* input converted to a tuple */
2394 Py_ssize_t *indices; /* one index per result element */
2395 PyObject *result; /* most recently returned result tuple */
2396 Py_ssize_t r; /* size of result tuple */
2397 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002398} combinationsobject;
2399
2400static PyTypeObject combinations_type;
2401
2402static PyObject *
2403combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2404{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002405 combinationsobject *co;
2406 Py_ssize_t n;
2407 Py_ssize_t r;
2408 PyObject *pool = NULL;
2409 PyObject *iterable = NULL;
2410 Py_ssize_t *indices = NULL;
2411 Py_ssize_t i;
2412 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002413
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002414 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2415 &iterable, &r))
2416 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002417
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002418 pool = PySequence_Tuple(iterable);
2419 if (pool == NULL)
2420 goto error;
2421 n = PyTuple_GET_SIZE(pool);
2422 if (r < 0) {
2423 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2424 goto error;
2425 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002426
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002427 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002428 if (indices == NULL) {
2429 PyErr_NoMemory();
2430 goto error;
2431 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002432
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002433 for (i=0 ; i<r ; i++)
2434 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002435
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002436 /* create combinationsobject structure */
2437 co = (combinationsobject *)type->tp_alloc(type, 0);
2438 if (co == NULL)
2439 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002440
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002441 co->pool = pool;
2442 co->indices = indices;
2443 co->result = NULL;
2444 co->r = r;
2445 co->stopped = r > n ? 1 : 0;
2446
2447 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002448
2449error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002450 if (indices != NULL)
2451 PyMem_Free(indices);
2452 Py_XDECREF(pool);
2453 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002454}
2455
2456static void
2457combinations_dealloc(combinationsobject *co)
2458{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002459 PyObject_GC_UnTrack(co);
2460 Py_XDECREF(co->pool);
2461 Py_XDECREF(co->result);
2462 if (co->indices != NULL)
2463 PyMem_Free(co->indices);
2464 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002465}
2466
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002467static PyObject *
2468combinations_sizeof(combinationsobject *co, void *unused)
2469{
2470 Py_ssize_t res;
2471
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002472 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002473 res += co->r * sizeof(Py_ssize_t);
2474 return PyLong_FromSsize_t(res);
2475}
2476
Christian Heimes380f7f22008-02-28 11:19:05 +00002477static int
2478combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2479{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002480 Py_VISIT(co->pool);
2481 Py_VISIT(co->result);
2482 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002483}
2484
2485static PyObject *
2486combinations_next(combinationsobject *co)
2487{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002488 PyObject *elem;
2489 PyObject *oldelem;
2490 PyObject *pool = co->pool;
2491 Py_ssize_t *indices = co->indices;
2492 PyObject *result = co->result;
2493 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2494 Py_ssize_t r = co->r;
2495 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002496
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002497 if (co->stopped)
2498 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002499
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002500 if (result == NULL) {
2501 /* On the first pass, initialize result tuple using the indices */
2502 result = PyTuple_New(r);
2503 if (result == NULL)
2504 goto empty;
2505 co->result = result;
2506 for (i=0; i<r ; i++) {
2507 index = indices[i];
2508 elem = PyTuple_GET_ITEM(pool, index);
2509 Py_INCREF(elem);
2510 PyTuple_SET_ITEM(result, i, elem);
2511 }
2512 } else {
2513 /* Copy the previous result tuple or re-use it if available */
2514 if (Py_REFCNT(result) > 1) {
2515 PyObject *old_result = result;
2516 result = PyTuple_New(r);
2517 if (result == NULL)
2518 goto empty;
2519 co->result = result;
2520 for (i=0; i<r ; i++) {
2521 elem = PyTuple_GET_ITEM(old_result, i);
2522 Py_INCREF(elem);
2523 PyTuple_SET_ITEM(result, i, elem);
2524 }
2525 Py_DECREF(old_result);
2526 }
2527 /* Now, we've got the only copy so we can update it in-place
2528 * CPython's empty tuple is a singleton and cached in
2529 * PyTuple's freelist.
2530 */
2531 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002532
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002533 /* Scan indices right-to-left until finding one that is not
2534 at its maximum (i + n - r). */
2535 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2536 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002537
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002538 /* If i is negative, then the indices are all at
2539 their maximum value and we're done. */
2540 if (i < 0)
2541 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002542
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002543 /* Increment the current index which we know is not at its
2544 maximum. Then move back to the right setting each index
2545 to its lowest possible value (one higher than the index
2546 to its left -- this maintains the sort order invariant). */
2547 indices[i]++;
2548 for (j=i+1 ; j<r ; j++)
2549 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002550
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002551 /* Update the result tuple for the new indices
2552 starting with i, the leftmost index that changed */
2553 for ( ; i<r ; i++) {
2554 index = indices[i];
2555 elem = PyTuple_GET_ITEM(pool, index);
2556 Py_INCREF(elem);
2557 oldelem = PyTuple_GET_ITEM(result, i);
2558 PyTuple_SET_ITEM(result, i, elem);
2559 Py_DECREF(oldelem);
2560 }
2561 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002562
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002563 Py_INCREF(result);
2564 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002565
2566empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002567 co->stopped = 1;
2568 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002569}
2570
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002571static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302572combinations_reduce(combinationsobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002573{
2574 if (lz->result == NULL) {
2575 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2576 } else if (lz->stopped) {
2577 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2578 } else {
2579 PyObject *indices;
2580 Py_ssize_t i;
2581
2582 /* we must pickle the indices and use them for setstate */
2583 indices = PyTuple_New(lz->r);
2584 if (!indices)
2585 return NULL;
2586 for (i=0; i<lz->r; i++)
2587 {
2588 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2589 if (!index) {
2590 Py_DECREF(indices);
2591 return NULL;
2592 }
2593 PyTuple_SET_ITEM(indices, i, index);
2594 }
2595
2596 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2597 }
2598}
2599
2600static PyObject *
2601combinations_setstate(combinationsobject *lz, PyObject *state)
2602{
2603 PyObject *result;
2604 Py_ssize_t i;
2605 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2606
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002607 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002608 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2609 return NULL;
2610 }
2611
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002612 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002613 Py_ssize_t max;
2614 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2615 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002616
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002617 if (index == -1 && PyErr_Occurred())
2618 return NULL; /* not an integer */
2619 max = i + n - lz->r;
2620 /* clamp the index (beware of negative max) */
2621 if (index > max)
2622 index = max;
2623 if (index < 0)
2624 index = 0;
2625 lz->indices[i] = index;
2626 }
2627
2628 result = PyTuple_New(lz->r);
2629 if (result == NULL)
2630 return NULL;
2631 for (i=0; i<lz->r; i++) {
2632 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2633 Py_INCREF(element);
2634 PyTuple_SET_ITEM(result, i, element);
2635 }
2636
Serhiy Storchaka48842712016-04-06 09:45:48 +03002637 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002638 Py_RETURN_NONE;
2639}
2640
2641static PyMethodDef combinations_methods[] = {
2642 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2643 reduce_doc},
2644 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2645 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002646 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2647 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002648 {NULL, NULL} /* sentinel */
2649};
2650
Christian Heimes380f7f22008-02-28 11:19:05 +00002651PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002652"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002653\n\
2654Return successive r-length combinations of elements in the iterable.\n\n\
2655combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2656
2657static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002658 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002659 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002660 sizeof(combinationsobject), /* tp_basicsize */
2661 0, /* tp_itemsize */
2662 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002663 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002664 0, /* tp_print */
2665 0, /* tp_getattr */
2666 0, /* tp_setattr */
2667 0, /* tp_reserved */
2668 0, /* tp_repr */
2669 0, /* tp_as_number */
2670 0, /* tp_as_sequence */
2671 0, /* tp_as_mapping */
2672 0, /* tp_hash */
2673 0, /* tp_call */
2674 0, /* tp_str */
2675 PyObject_GenericGetAttr, /* tp_getattro */
2676 0, /* tp_setattro */
2677 0, /* tp_as_buffer */
2678 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2679 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002680 combinations_doc, /* tp_doc */
2681 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002682 0, /* tp_clear */
2683 0, /* tp_richcompare */
2684 0, /* tp_weaklistoffset */
2685 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002686 (iternextfunc)combinations_next, /* tp_iternext */
2687 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002688 0, /* tp_members */
2689 0, /* tp_getset */
2690 0, /* tp_base */
2691 0, /* tp_dict */
2692 0, /* tp_descr_get */
2693 0, /* tp_descr_set */
2694 0, /* tp_dictoffset */
2695 0, /* tp_init */
2696 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002697 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002698 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002699};
2700
2701
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002702/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002703
2704/* Equivalent to:
2705
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002706 def combinations_with_replacement(iterable, r):
2707 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2708 # number items returned: (n+r-1)! / r! / (n-1)!
2709 pool = tuple(iterable)
2710 n = len(pool)
2711 indices = [0] * r
2712 yield tuple(pool[i] for i in indices)
2713 while 1:
2714 for i in reversed(range(r)):
2715 if indices[i] != n - 1:
2716 break
2717 else:
2718 return
2719 indices[i:] = [indices[i] + 1] * (r - i)
2720 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002721
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002722 def combinations_with_replacement2(iterable, r):
2723 'Alternate version that filters from product()'
2724 pool = tuple(iterable)
2725 n = len(pool)
2726 for indices in product(range(n), repeat=r):
2727 if sorted(indices) == list(indices):
2728 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002729*/
2730typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002731 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002732 PyObject *pool; /* input converted to a tuple */
2733 Py_ssize_t *indices; /* one index per result element */
2734 PyObject *result; /* most recently returned result tuple */
2735 Py_ssize_t r; /* size of result tuple */
2736 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002737} cwrobject;
2738
2739static PyTypeObject cwr_type;
2740
2741static PyObject *
2742cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2743{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002744 cwrobject *co;
2745 Py_ssize_t n;
2746 Py_ssize_t r;
2747 PyObject *pool = NULL;
2748 PyObject *iterable = NULL;
2749 Py_ssize_t *indices = NULL;
2750 Py_ssize_t i;
2751 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002752
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002753 if (!PyArg_ParseTupleAndKeywords(args, kwds,
2754 "On:combinations_with_replacement",
2755 kwargs, &iterable, &r))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002756 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002757
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002758 pool = PySequence_Tuple(iterable);
2759 if (pool == NULL)
2760 goto error;
2761 n = PyTuple_GET_SIZE(pool);
2762 if (r < 0) {
2763 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2764 goto error;
2765 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002766
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002767 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002768 if (indices == NULL) {
2769 PyErr_NoMemory();
2770 goto error;
2771 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002772
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002773 for (i=0 ; i<r ; i++)
2774 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002775
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002776 /* create cwrobject structure */
2777 co = (cwrobject *)type->tp_alloc(type, 0);
2778 if (co == NULL)
2779 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002780
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002781 co->pool = pool;
2782 co->indices = indices;
2783 co->result = NULL;
2784 co->r = r;
2785 co->stopped = !n && r;
2786
2787 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002788
2789error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002790 if (indices != NULL)
2791 PyMem_Free(indices);
2792 Py_XDECREF(pool);
2793 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002794}
2795
2796static void
2797cwr_dealloc(cwrobject *co)
2798{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002799 PyObject_GC_UnTrack(co);
2800 Py_XDECREF(co->pool);
2801 Py_XDECREF(co->result);
2802 if (co->indices != NULL)
2803 PyMem_Free(co->indices);
2804 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002805}
2806
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002807static PyObject *
2808cwr_sizeof(cwrobject *co, void *unused)
2809{
2810 Py_ssize_t res;
2811
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002812 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002813 res += co->r * sizeof(Py_ssize_t);
2814 return PyLong_FromSsize_t(res);
2815}
2816
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002817static int
2818cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2819{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002820 Py_VISIT(co->pool);
2821 Py_VISIT(co->result);
2822 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002823}
2824
2825static PyObject *
2826cwr_next(cwrobject *co)
2827{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002828 PyObject *elem;
2829 PyObject *oldelem;
2830 PyObject *pool = co->pool;
2831 Py_ssize_t *indices = co->indices;
2832 PyObject *result = co->result;
2833 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2834 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002835 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002836
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002837 if (co->stopped)
2838 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002839
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002840 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002841 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002842 result = PyTuple_New(r);
2843 if (result == NULL)
2844 goto empty;
2845 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002846 if (n > 0) {
2847 elem = PyTuple_GET_ITEM(pool, 0);
2848 for (i=0; i<r ; i++) {
2849 assert(indices[i] == 0);
2850 Py_INCREF(elem);
2851 PyTuple_SET_ITEM(result, i, elem);
2852 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002853 }
2854 } else {
2855 /* Copy the previous result tuple or re-use it if available */
2856 if (Py_REFCNT(result) > 1) {
2857 PyObject *old_result = result;
2858 result = PyTuple_New(r);
2859 if (result == NULL)
2860 goto empty;
2861 co->result = result;
2862 for (i=0; i<r ; i++) {
2863 elem = PyTuple_GET_ITEM(old_result, i);
2864 Py_INCREF(elem);
2865 PyTuple_SET_ITEM(result, i, elem);
2866 }
2867 Py_DECREF(old_result);
2868 }
2869 /* Now, we've got the only copy so we can update it in-place CPython's
2870 empty tuple is a singleton and cached in PyTuple's freelist. */
2871 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002872
Tim Peters9edb1682013-09-03 11:49:31 -05002873 /* Scan indices right-to-left until finding one that is not
2874 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002875 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2876 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002877
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002878 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002879 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002880 if (i < 0)
2881 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002882
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002883 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002884 maximum. Then set all to the right to the same value. */
2885 index = indices[i] + 1;
2886 assert(index < n);
2887 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002888 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002889 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002890 Py_INCREF(elem);
2891 oldelem = PyTuple_GET_ITEM(result, i);
2892 PyTuple_SET_ITEM(result, i, elem);
2893 Py_DECREF(oldelem);
2894 }
2895 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002896
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002897 Py_INCREF(result);
2898 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002899
2900empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002901 co->stopped = 1;
2902 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002903}
2904
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002905static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302906cwr_reduce(cwrobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002907{
2908 if (lz->result == NULL) {
2909 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2910 } else if (lz->stopped) {
2911 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2912 } else {
2913 PyObject *indices;
2914 Py_ssize_t i;
2915
2916 /* we must pickle the indices and use them for setstate */
2917 indices = PyTuple_New(lz->r);
2918 if (!indices)
2919 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002920 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002921 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2922 if (!index) {
2923 Py_DECREF(indices);
2924 return NULL;
2925 }
2926 PyTuple_SET_ITEM(indices, i, index);
2927 }
2928
2929 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2930 }
2931}
2932
2933static PyObject *
2934cwr_setstate(cwrobject *lz, PyObject *state)
2935{
2936 PyObject *result;
2937 Py_ssize_t n, i;
2938
2939 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2940 {
2941 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2942 return NULL;
2943 }
2944
2945 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002946 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002947 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2948 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002949
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002950 if (index < 0 && PyErr_Occurred())
2951 return NULL; /* not an integer */
2952 /* clamp the index */
2953 if (index < 0)
2954 index = 0;
2955 else if (index > n-1)
2956 index = n-1;
2957 lz->indices[i] = index;
2958 }
2959 result = PyTuple_New(lz->r);
2960 if (result == NULL)
2961 return NULL;
2962 for (i=0; i<lz->r; i++) {
2963 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2964 Py_INCREF(element);
2965 PyTuple_SET_ITEM(result, i, element);
2966 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002967 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002968 Py_RETURN_NONE;
2969}
2970
2971static PyMethodDef cwr_methods[] = {
2972 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2973 reduce_doc},
2974 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2975 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002976 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
2977 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002978 {NULL, NULL} /* sentinel */
2979};
2980
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002981PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002982"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002983\n\
2984Return successive r-length combinations of elements in the iterable\n\
2985allowing individual elements to have successive repeats.\n\
2986combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2987
2988static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002989 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002990 "itertools.combinations_with_replacement", /* tp_name */
2991 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002992 0, /* tp_itemsize */
2993 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002994 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002995 0, /* tp_print */
2996 0, /* tp_getattr */
2997 0, /* tp_setattr */
2998 0, /* tp_reserved */
2999 0, /* tp_repr */
3000 0, /* tp_as_number */
3001 0, /* tp_as_sequence */
3002 0, /* tp_as_mapping */
3003 0, /* tp_hash */
3004 0, /* tp_call */
3005 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003006 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003007 0, /* tp_setattro */
3008 0, /* tp_as_buffer */
3009 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003010 Py_TPFLAGS_BASETYPE, /* tp_flags */
3011 cwr_doc, /* tp_doc */
3012 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003013 0, /* tp_clear */
3014 0, /* tp_richcompare */
3015 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003016 PyObject_SelfIter, /* tp_iter */
3017 (iternextfunc)cwr_next, /* tp_iternext */
3018 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003019 0, /* tp_members */
3020 0, /* tp_getset */
3021 0, /* tp_base */
3022 0, /* tp_dict */
3023 0, /* tp_descr_get */
3024 0, /* tp_descr_set */
3025 0, /* tp_dictoffset */
3026 0, /* tp_init */
3027 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003028 cwr_new, /* tp_new */
3029 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003030};
3031
3032
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003033/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003034
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003035def permutations(iterable, r=None):
3036 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
3037 pool = tuple(iterable)
3038 n = len(pool)
3039 r = n if r is None else r
3040 indices = range(n)
3041 cycles = range(n-r+1, n+1)[::-1]
3042 yield tuple(pool[i] for i in indices[:r])
3043 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003044 for i in reversed(range(r)):
3045 cycles[i] -= 1
3046 if cycles[i] == 0:
3047 indices[i:] = indices[i+1:] + indices[i:i+1]
3048 cycles[i] = n - i
3049 else:
3050 j = cycles[i]
3051 indices[i], indices[-j] = indices[-j], indices[i]
3052 yield tuple(pool[i] for i in indices[:r])
3053 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003054 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003055 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003056*/
3057
3058typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003059 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003060 PyObject *pool; /* input converted to a tuple */
3061 Py_ssize_t *indices; /* one index per element in the pool */
3062 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3063 PyObject *result; /* most recently returned result tuple */
3064 Py_ssize_t r; /* size of result tuple */
3065 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003066} permutationsobject;
3067
3068static PyTypeObject permutations_type;
3069
3070static PyObject *
3071permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3072{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003073 permutationsobject *po;
3074 Py_ssize_t n;
3075 Py_ssize_t r;
3076 PyObject *robj = Py_None;
3077 PyObject *pool = NULL;
3078 PyObject *iterable = NULL;
3079 Py_ssize_t *indices = NULL;
3080 Py_ssize_t *cycles = NULL;
3081 Py_ssize_t i;
3082 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003083
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003084 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
3085 &iterable, &robj))
3086 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003087
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003088 pool = PySequence_Tuple(iterable);
3089 if (pool == NULL)
3090 goto error;
3091 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003092
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003093 r = n;
3094 if (robj != Py_None) {
3095 if (!PyLong_Check(robj)) {
3096 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3097 goto error;
3098 }
3099 r = PyLong_AsSsize_t(robj);
3100 if (r == -1 && PyErr_Occurred())
3101 goto error;
3102 }
3103 if (r < 0) {
3104 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3105 goto error;
3106 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003107
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003108 indices = PyMem_New(Py_ssize_t, n);
3109 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003110 if (indices == NULL || cycles == NULL) {
3111 PyErr_NoMemory();
3112 goto error;
3113 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003114
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003115 for (i=0 ; i<n ; i++)
3116 indices[i] = i;
3117 for (i=0 ; i<r ; i++)
3118 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003119
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003120 /* create permutationsobject structure */
3121 po = (permutationsobject *)type->tp_alloc(type, 0);
3122 if (po == NULL)
3123 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003124
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003125 po->pool = pool;
3126 po->indices = indices;
3127 po->cycles = cycles;
3128 po->result = NULL;
3129 po->r = r;
3130 po->stopped = r > n ? 1 : 0;
3131
3132 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003133
3134error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003135 if (indices != NULL)
3136 PyMem_Free(indices);
3137 if (cycles != NULL)
3138 PyMem_Free(cycles);
3139 Py_XDECREF(pool);
3140 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003141}
3142
3143static void
3144permutations_dealloc(permutationsobject *po)
3145{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003146 PyObject_GC_UnTrack(po);
3147 Py_XDECREF(po->pool);
3148 Py_XDECREF(po->result);
3149 PyMem_Free(po->indices);
3150 PyMem_Free(po->cycles);
3151 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003152}
3153
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003154static PyObject *
3155permutations_sizeof(permutationsobject *po, void *unused)
3156{
3157 Py_ssize_t res;
3158
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003159 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003160 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3161 res += po->r * sizeof(Py_ssize_t);
3162 return PyLong_FromSsize_t(res);
3163}
3164
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003165static int
3166permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3167{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003168 Py_VISIT(po->pool);
3169 Py_VISIT(po->result);
3170 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003171}
3172
3173static PyObject *
3174permutations_next(permutationsobject *po)
3175{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003176 PyObject *elem;
3177 PyObject *oldelem;
3178 PyObject *pool = po->pool;
3179 Py_ssize_t *indices = po->indices;
3180 Py_ssize_t *cycles = po->cycles;
3181 PyObject *result = po->result;
3182 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3183 Py_ssize_t r = po->r;
3184 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003185
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003186 if (po->stopped)
3187 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003188
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003189 if (result == NULL) {
3190 /* On the first pass, initialize result tuple using the indices */
3191 result = PyTuple_New(r);
3192 if (result == NULL)
3193 goto empty;
3194 po->result = result;
3195 for (i=0; i<r ; i++) {
3196 index = indices[i];
3197 elem = PyTuple_GET_ITEM(pool, index);
3198 Py_INCREF(elem);
3199 PyTuple_SET_ITEM(result, i, elem);
3200 }
3201 } else {
3202 if (n == 0)
3203 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003204
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003205 /* Copy the previous result tuple or re-use it if available */
3206 if (Py_REFCNT(result) > 1) {
3207 PyObject *old_result = result;
3208 result = PyTuple_New(r);
3209 if (result == NULL)
3210 goto empty;
3211 po->result = result;
3212 for (i=0; i<r ; i++) {
3213 elem = PyTuple_GET_ITEM(old_result, i);
3214 Py_INCREF(elem);
3215 PyTuple_SET_ITEM(result, i, elem);
3216 }
3217 Py_DECREF(old_result);
3218 }
3219 /* Now, we've got the only copy so we can update it in-place */
3220 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003221
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003222 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3223 for (i=r-1 ; i>=0 ; i--) {
3224 cycles[i] -= 1;
3225 if (cycles[i] == 0) {
3226 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3227 index = indices[i];
3228 for (j=i ; j<n-1 ; j++)
3229 indices[j] = indices[j+1];
3230 indices[n-1] = index;
3231 cycles[i] = n - i;
3232 } else {
3233 j = cycles[i];
3234 index = indices[i];
3235 indices[i] = indices[n-j];
3236 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003237
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003238 for (k=i; k<r ; k++) {
3239 /* start with i, the leftmost element that changed */
3240 /* yield tuple(pool[k] for k in indices[:r]) */
3241 index = indices[k];
3242 elem = PyTuple_GET_ITEM(pool, index);
3243 Py_INCREF(elem);
3244 oldelem = PyTuple_GET_ITEM(result, k);
3245 PyTuple_SET_ITEM(result, k, elem);
3246 Py_DECREF(oldelem);
3247 }
3248 break;
3249 }
3250 }
3251 /* If i is negative, then the cycles have all
3252 rolled-over and we're done. */
3253 if (i < 0)
3254 goto empty;
3255 }
3256 Py_INCREF(result);
3257 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003258
3259empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003260 po->stopped = 1;
3261 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003262}
3263
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003264static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303265permutations_reduce(permutationsobject *po, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003266{
3267 if (po->result == NULL) {
3268 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3269 } else if (po->stopped) {
3270 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3271 } else {
3272 PyObject *indices=NULL, *cycles=NULL;
3273 Py_ssize_t n, i;
3274
3275 /* we must pickle the indices and cycles and use them for setstate */
3276 n = PyTuple_GET_SIZE(po->pool);
3277 indices = PyTuple_New(n);
3278 if (indices == NULL)
3279 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003280 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003281 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3282 if (!index)
3283 goto err;
3284 PyTuple_SET_ITEM(indices, i, index);
3285 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003286
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003287 cycles = PyTuple_New(po->r);
3288 if (cycles == NULL)
3289 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003290 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003291 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3292 if (!index)
3293 goto err;
3294 PyTuple_SET_ITEM(cycles, i, index);
3295 }
3296 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3297 po->pool, po->r,
3298 indices, cycles);
3299 err:
3300 Py_XDECREF(indices);
3301 Py_XDECREF(cycles);
3302 return NULL;
3303 }
3304}
3305
3306static PyObject *
3307permutations_setstate(permutationsobject *po, PyObject *state)
3308{
3309 PyObject *indices, *cycles, *result;
3310 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003311
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003312 if (!PyTuple_Check(state)) {
3313 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3314 return NULL;
3315 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003316 if (!PyArg_ParseTuple(state, "O!O!",
3317 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003318 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003319 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003320 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003321
3322 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003323 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003324 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3325 return NULL;
3326 }
3327
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003328 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003329 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3330 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3331 if (index < 0 && PyErr_Occurred())
3332 return NULL; /* not an integer */
3333 /* clamp the index */
3334 if (index < 0)
3335 index = 0;
3336 else if (index > n-1)
3337 index = n-1;
3338 po->indices[i] = index;
3339 }
3340
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003341 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003342 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3343 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3344 if (index < 0 && PyErr_Occurred())
3345 return NULL; /* not an integer */
3346 if (index < 1)
3347 index = 1;
3348 else if (index > n-i)
3349 index = n-i;
3350 po->cycles[i] = index;
3351 }
3352 result = PyTuple_New(po->r);
3353 if (result == NULL)
3354 return NULL;
3355 for (i=0; i<po->r; i++) {
3356 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3357 Py_INCREF(element);
3358 PyTuple_SET_ITEM(result, i, element);
3359 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003360 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003361 Py_RETURN_NONE;
3362}
3363
3364static PyMethodDef permuations_methods[] = {
3365 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3366 reduce_doc},
3367 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3368 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003369 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3370 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003371 {NULL, NULL} /* sentinel */
3372};
3373
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003374PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003375"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003376\n\
3377Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003378permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003379
3380static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003381 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003382 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003383 sizeof(permutationsobject), /* tp_basicsize */
3384 0, /* tp_itemsize */
3385 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003386 (destructor)permutations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003387 0, /* tp_print */
3388 0, /* tp_getattr */
3389 0, /* tp_setattr */
3390 0, /* tp_reserved */
3391 0, /* tp_repr */
3392 0, /* tp_as_number */
3393 0, /* tp_as_sequence */
3394 0, /* tp_as_mapping */
3395 0, /* tp_hash */
3396 0, /* tp_call */
3397 0, /* tp_str */
3398 PyObject_GenericGetAttr, /* tp_getattro */
3399 0, /* tp_setattro */
3400 0, /* tp_as_buffer */
3401 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3402 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003403 permutations_doc, /* tp_doc */
3404 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003405 0, /* tp_clear */
3406 0, /* tp_richcompare */
3407 0, /* tp_weaklistoffset */
3408 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003409 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003410 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003411 0, /* tp_members */
3412 0, /* tp_getset */
3413 0, /* tp_base */
3414 0, /* tp_dict */
3415 0, /* tp_descr_get */
3416 0, /* tp_descr_set */
3417 0, /* tp_dictoffset */
3418 0, /* tp_init */
3419 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003420 permutations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003421 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003422};
3423
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003424/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003425
3426typedef struct {
3427 PyObject_HEAD
3428 PyObject *total;
3429 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003430 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003431} accumulateobject;
3432
3433static PyTypeObject accumulate_type;
3434
3435static PyObject *
3436accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3437{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003438 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003439 PyObject *iterable;
3440 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003441 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003442 accumulateobject *lz;
3443
Raymond Hettinger5d446132011-03-27 18:52:10 -07003444 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3445 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003446 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003447
3448 /* Get iterator. */
3449 it = PyObject_GetIter(iterable);
3450 if (it == NULL)
3451 return NULL;
3452
Raymond Hettinger482ba772010-12-01 22:48:00 +00003453 /* create accumulateobject structure */
3454 lz = (accumulateobject *)type->tp_alloc(type, 0);
3455 if (lz == NULL) {
3456 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003457 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003458 }
3459
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003460 if (binop != Py_None) {
3461 Py_XINCREF(binop);
3462 lz->binop = binop;
3463 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003464 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003465 lz->it = it;
3466 return (PyObject *)lz;
3467}
3468
3469static void
3470accumulate_dealloc(accumulateobject *lz)
3471{
3472 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003473 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003474 Py_XDECREF(lz->total);
3475 Py_XDECREF(lz->it);
3476 Py_TYPE(lz)->tp_free(lz);
3477}
3478
3479static int
3480accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3481{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003482 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003483 Py_VISIT(lz->it);
3484 Py_VISIT(lz->total);
3485 return 0;
3486}
3487
3488static PyObject *
3489accumulate_next(accumulateobject *lz)
3490{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003491 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003492
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003493 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003494 if (val == NULL)
3495 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003496
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003497 if (lz->total == NULL) {
3498 Py_INCREF(val);
3499 lz->total = val;
3500 return lz->total;
3501 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003502
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003503 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003504 newtotal = PyNumber_Add(lz->total, val);
3505 else
3506 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003507 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003508 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003509 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003510
Raymond Hettinger482ba772010-12-01 22:48:00 +00003511 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003512 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003513 return newtotal;
3514}
3515
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003516static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303517accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003518{
Serhiy Storchakad5516252016-03-06 14:00:45 +02003519 if (lz->total == Py_None) {
3520 PyObject *it;
3521
3522 if (PyType_Ready(&chain_type) < 0)
3523 return NULL;
3524 if (PyType_Ready(&islice_type) < 0)
3525 return NULL;
3526 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3527 lz->total, lz->it);
3528 if (it == NULL)
3529 return NULL;
3530 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3531 it, lz->binop ? lz->binop : Py_None);
3532 if (it == NULL)
3533 return NULL;
3534 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3535 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003536 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3537 lz->it, lz->binop?lz->binop:Py_None,
3538 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003539}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003540
3541static PyObject *
3542accumulate_setstate(accumulateobject *lz, PyObject *state)
3543{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003544 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003545 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003546 Py_RETURN_NONE;
3547}
3548
3549static PyMethodDef accumulate_methods[] = {
3550 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3551 reduce_doc},
3552 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3553 setstate_doc},
3554 {NULL, NULL} /* sentinel */
3555};
3556
Raymond Hettinger482ba772010-12-01 22:48:00 +00003557PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003558"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003559\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003560Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003561
3562static PyTypeObject accumulate_type = {
3563 PyVarObject_HEAD_INIT(NULL, 0)
3564 "itertools.accumulate", /* tp_name */
3565 sizeof(accumulateobject), /* tp_basicsize */
3566 0, /* tp_itemsize */
3567 /* methods */
3568 (destructor)accumulate_dealloc, /* tp_dealloc */
3569 0, /* tp_print */
3570 0, /* tp_getattr */
3571 0, /* tp_setattr */
3572 0, /* tp_reserved */
3573 0, /* tp_repr */
3574 0, /* tp_as_number */
3575 0, /* tp_as_sequence */
3576 0, /* tp_as_mapping */
3577 0, /* tp_hash */
3578 0, /* tp_call */
3579 0, /* tp_str */
3580 PyObject_GenericGetAttr, /* tp_getattro */
3581 0, /* tp_setattro */
3582 0, /* tp_as_buffer */
3583 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3584 Py_TPFLAGS_BASETYPE, /* tp_flags */
3585 accumulate_doc, /* tp_doc */
3586 (traverseproc)accumulate_traverse, /* tp_traverse */
3587 0, /* tp_clear */
3588 0, /* tp_richcompare */
3589 0, /* tp_weaklistoffset */
3590 PyObject_SelfIter, /* tp_iter */
3591 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003592 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003593 0, /* tp_members */
3594 0, /* tp_getset */
3595 0, /* tp_base */
3596 0, /* tp_dict */
3597 0, /* tp_descr_get */
3598 0, /* tp_descr_set */
3599 0, /* tp_dictoffset */
3600 0, /* tp_init */
3601 0, /* tp_alloc */
3602 accumulate_new, /* tp_new */
3603 PyObject_GC_Del, /* tp_free */
3604};
3605
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003606
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003607/* compress object ************************************************************/
3608
3609/* Equivalent to:
3610
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003611 def compress(data, selectors):
3612 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3613 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003614*/
3615
3616typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003617 PyObject_HEAD
3618 PyObject *data;
3619 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003620} compressobject;
3621
3622static PyTypeObject compress_type;
3623
3624static PyObject *
3625compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3626{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003627 PyObject *seq1, *seq2;
3628 PyObject *data=NULL, *selectors=NULL;
3629 compressobject *lz;
3630 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003631
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003632 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3633 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003634
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003635 data = PyObject_GetIter(seq1);
3636 if (data == NULL)
3637 goto fail;
3638 selectors = PyObject_GetIter(seq2);
3639 if (selectors == NULL)
3640 goto fail;
3641
3642 /* create compressobject structure */
3643 lz = (compressobject *)type->tp_alloc(type, 0);
3644 if (lz == NULL)
3645 goto fail;
3646 lz->data = data;
3647 lz->selectors = selectors;
3648 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003649
3650fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003651 Py_XDECREF(data);
3652 Py_XDECREF(selectors);
3653 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003654}
3655
3656static void
3657compress_dealloc(compressobject *lz)
3658{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003659 PyObject_GC_UnTrack(lz);
3660 Py_XDECREF(lz->data);
3661 Py_XDECREF(lz->selectors);
3662 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003663}
3664
3665static int
3666compress_traverse(compressobject *lz, visitproc visit, void *arg)
3667{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003668 Py_VISIT(lz->data);
3669 Py_VISIT(lz->selectors);
3670 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003671}
3672
3673static PyObject *
3674compress_next(compressobject *lz)
3675{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003676 PyObject *data = lz->data, *selectors = lz->selectors;
3677 PyObject *datum, *selector;
3678 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3679 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3680 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003681
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003682 while (1) {
3683 /* Steps: get datum, get selector, evaluate selector.
3684 Order is important (to match the pure python version
3685 in terms of which input gets a chance to raise an
3686 exception first).
3687 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003688
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003689 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003690 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003691 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003692
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003693 selector = selectornext(selectors);
3694 if (selector == NULL) {
3695 Py_DECREF(datum);
3696 return NULL;
3697 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003698
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003699 ok = PyObject_IsTrue(selector);
3700 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003701 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003702 return datum;
3703 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003704 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003705 return NULL;
3706 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003707}
3708
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003709static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303710compress_reduce(compressobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003711{
3712 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3713 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003714}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003715
3716static PyMethodDef compress_methods[] = {
3717 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3718 reduce_doc},
3719 {NULL, NULL} /* sentinel */
3720};
3721
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003722PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003723"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003724\n\
3725Return data elements corresponding to true selector elements.\n\
3726Forms a shorter iterator from selected data elements using the\n\
3727selectors to choose the data elements.");
3728
3729static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003730 PyVarObject_HEAD_INIT(NULL, 0)
3731 "itertools.compress", /* tp_name */
3732 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003733 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003734 /* methods */
3735 (destructor)compress_dealloc, /* tp_dealloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003736 0, /* tp_print */
3737 0, /* tp_getattr */
3738 0, /* tp_setattr */
3739 0, /* tp_reserved */
3740 0, /* tp_repr */
3741 0, /* tp_as_number */
3742 0, /* tp_as_sequence */
3743 0, /* tp_as_mapping */
3744 0, /* tp_hash */
3745 0, /* tp_call */
3746 0, /* tp_str */
3747 PyObject_GenericGetAttr, /* tp_getattro */
3748 0, /* tp_setattro */
3749 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003750 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003751 Py_TPFLAGS_BASETYPE, /* tp_flags */
3752 compress_doc, /* tp_doc */
3753 (traverseproc)compress_traverse, /* tp_traverse */
3754 0, /* tp_clear */
3755 0, /* tp_richcompare */
3756 0, /* tp_weaklistoffset */
3757 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003758 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003759 compress_methods, /* tp_methods */
3760 0, /* tp_members */
3761 0, /* tp_getset */
3762 0, /* tp_base */
3763 0, /* tp_dict */
3764 0, /* tp_descr_get */
3765 0, /* tp_descr_set */
3766 0, /* tp_dictoffset */
3767 0, /* tp_init */
3768 0, /* tp_alloc */
3769 compress_new, /* tp_new */
3770 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003771};
3772
3773
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003774/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003775
3776typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003777 PyObject_HEAD
3778 PyObject *func;
3779 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003780} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003781
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003782static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003783
3784static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003785filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003786{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003787 PyObject *func, *seq;
3788 PyObject *it;
3789 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003790
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003791 if (type == &filterfalse_type &&
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03003792 !_PyArg_NoKeywords("filterfalse", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003793 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003794
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003795 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3796 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003797
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003798 /* Get iterator. */
3799 it = PyObject_GetIter(seq);
3800 if (it == NULL)
3801 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003802
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003803 /* create filterfalseobject structure */
3804 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3805 if (lz == NULL) {
3806 Py_DECREF(it);
3807 return NULL;
3808 }
3809 Py_INCREF(func);
3810 lz->func = func;
3811 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003812
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003813 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003814}
3815
3816static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003817filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003818{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003819 PyObject_GC_UnTrack(lz);
3820 Py_XDECREF(lz->func);
3821 Py_XDECREF(lz->it);
3822 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003823}
3824
3825static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003826filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003827{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003828 Py_VISIT(lz->it);
3829 Py_VISIT(lz->func);
3830 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003831}
3832
3833static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003834filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003835{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003836 PyObject *item;
3837 PyObject *it = lz->it;
3838 long ok;
3839 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003840
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003841 iternext = *Py_TYPE(it)->tp_iternext;
3842 for (;;) {
3843 item = iternext(it);
3844 if (item == NULL)
3845 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003846
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003847 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3848 ok = PyObject_IsTrue(item);
3849 } else {
3850 PyObject *good;
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01003851 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003852 if (good == NULL) {
3853 Py_DECREF(item);
3854 return NULL;
3855 }
3856 ok = PyObject_IsTrue(good);
3857 Py_DECREF(good);
3858 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003859 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003860 return item;
3861 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003862 if (ok < 0)
3863 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003864 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003865}
3866
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003867static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303868filterfalse_reduce(filterfalseobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003869{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003870 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3871}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003872
3873static PyMethodDef filterfalse_methods[] = {
3874 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3875 reduce_doc},
3876 {NULL, NULL} /* sentinel */
3877};
3878
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003879PyDoc_STRVAR(filterfalse_doc,
3880"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003881\n\
3882Return those items of sequence for which function(item) is false.\n\
3883If function is None, return the items that are false.");
3884
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003885static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003886 PyVarObject_HEAD_INIT(NULL, 0)
3887 "itertools.filterfalse", /* tp_name */
3888 sizeof(filterfalseobject), /* tp_basicsize */
3889 0, /* tp_itemsize */
3890 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003891 (destructor)filterfalse_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003892 0, /* tp_print */
3893 0, /* tp_getattr */
3894 0, /* tp_setattr */
3895 0, /* tp_reserved */
3896 0, /* tp_repr */
3897 0, /* tp_as_number */
3898 0, /* tp_as_sequence */
3899 0, /* tp_as_mapping */
3900 0, /* tp_hash */
3901 0, /* tp_call */
3902 0, /* tp_str */
3903 PyObject_GenericGetAttr, /* tp_getattro */
3904 0, /* tp_setattro */
3905 0, /* tp_as_buffer */
3906 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3907 Py_TPFLAGS_BASETYPE, /* tp_flags */
3908 filterfalse_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003909 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003910 0, /* tp_clear */
3911 0, /* tp_richcompare */
3912 0, /* tp_weaklistoffset */
3913 PyObject_SelfIter, /* tp_iter */
3914 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003915 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003916 0, /* tp_members */
3917 0, /* tp_getset */
3918 0, /* tp_base */
3919 0, /* tp_dict */
3920 0, /* tp_descr_get */
3921 0, /* tp_descr_set */
3922 0, /* tp_dictoffset */
3923 0, /* tp_init */
3924 0, /* tp_alloc */
3925 filterfalse_new, /* tp_new */
3926 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003927};
3928
3929
3930/* count object ************************************************************/
3931
3932typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003933 PyObject_HEAD
3934 Py_ssize_t cnt;
3935 PyObject *long_cnt;
3936 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003937} countobject;
3938
Raymond Hettinger30729212009-02-12 06:28:27 +00003939/* Counting logic and invariants:
3940
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003941fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003942
Serhiy Storchaka483405b2015-02-17 10:14:30 +02003943 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003944 Advances with: cnt += 1
3945 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003946
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003947slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003948
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003949 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3950 All counting is done with python objects (no overflows or underflows).
3951 Advances with: long_cnt += long_step
3952 Step may be zero -- effectively a slow version of repeat(cnt).
3953 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003954*/
3955
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003956static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003957
3958static PyObject *
3959count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3960{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003961 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003962 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003963 Py_ssize_t cnt = 0;
3964 PyObject *long_cnt = NULL;
3965 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003966 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003967 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003968
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003969 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3970 kwlist, &long_cnt, &long_step))
3971 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003972
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003973 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3974 (long_step != NULL && !PyNumber_Check(long_step))) {
3975 PyErr_SetString(PyExc_TypeError, "a number is required");
3976 return NULL;
3977 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003978
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003979 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
3980 (long_step == NULL || PyLong_Check(long_step));
3981
3982 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003983 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003984 if (fast_mode) {
3985 assert(PyLong_Check(long_cnt));
3986 cnt = PyLong_AsSsize_t(long_cnt);
3987 if (cnt == -1 && PyErr_Occurred()) {
3988 PyErr_Clear();
3989 fast_mode = 0;
3990 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003991 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003992 } else {
3993 cnt = 0;
Serhiy Storchakaba85d692017-03-30 09:09:41 +03003994 long_cnt = _PyLong_Zero;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003995 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03003996 Py_INCREF(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003997
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003998 /* If not specified, step defaults to 1 */
Serhiy Storchakaba85d692017-03-30 09:09:41 +03003999 if (long_step == NULL)
4000 long_step = _PyLong_One;
4001 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004002
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004003 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004005 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004006 if (fast_mode) {
4007 assert(PyLong_Check(long_step));
4008 step = PyLong_AsLong(long_step);
4009 if (step != 1) {
4010 fast_mode = 0;
4011 if (step == -1 && PyErr_Occurred())
4012 PyErr_Clear();
4013 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004014 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004015
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004016 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004017 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004018 else
4019 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004020
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004021 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4022 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4023 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004024 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004025
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004026 /* create countobject structure */
4027 lz = (countobject *)type->tp_alloc(type, 0);
4028 if (lz == NULL) {
4029 Py_XDECREF(long_cnt);
4030 return NULL;
4031 }
4032 lz->cnt = cnt;
4033 lz->long_cnt = long_cnt;
4034 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004035
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004036 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004037}
4038
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004039static void
4040count_dealloc(countobject *lz)
4041{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004042 PyObject_GC_UnTrack(lz);
4043 Py_XDECREF(lz->long_cnt);
4044 Py_XDECREF(lz->long_step);
4045 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004046}
4047
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004048static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004049count_traverse(countobject *lz, visitproc visit, void *arg)
4050{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004051 Py_VISIT(lz->long_cnt);
4052 Py_VISIT(lz->long_step);
4053 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004054}
4055
4056static PyObject *
4057count_nextlong(countobject *lz)
4058{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004059 PyObject *long_cnt;
4060 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004061
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004062 long_cnt = lz->long_cnt;
4063 if (long_cnt == NULL) {
4064 /* Switch to slow_mode */
4065 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4066 if (long_cnt == NULL)
4067 return NULL;
4068 }
4069 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004070
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004071 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4072 if (stepped_up == NULL)
4073 return NULL;
4074 lz->long_cnt = stepped_up;
4075 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004076}
4077
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004078static PyObject *
4079count_next(countobject *lz)
4080{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004081 if (lz->cnt == PY_SSIZE_T_MAX)
4082 return count_nextlong(lz);
4083 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004084}
4085
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004086static PyObject *
4087count_repr(countobject *lz)
4088{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004089 if (lz->cnt != PY_SSIZE_T_MAX)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004090 return PyUnicode_FromFormat("%s(%zd)",
4091 _PyType_Name(Py_TYPE(lz)), lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004092
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004093 if (PyLong_Check(lz->long_step)) {
4094 long step = PyLong_AsLong(lz->long_step);
4095 if (step == -1 && PyErr_Occurred()) {
4096 PyErr_Clear();
4097 }
4098 if (step == 1) {
4099 /* Don't display step when it is an integer equal to 1 */
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004100 return PyUnicode_FromFormat("%s(%R)",
4101 _PyType_Name(Py_TYPE(lz)),
4102 lz->long_cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004103 }
4104 }
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004105 return PyUnicode_FromFormat("%s(%R, %R)",
4106 _PyType_Name(Py_TYPE(lz)),
4107 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004108}
4109
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004110static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304111count_reduce(countobject *lz, PyObject *Py_UNUSED(ignored))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004112{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004113 if (lz->cnt == PY_SSIZE_T_MAX)
4114 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4115 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004116}
4117
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004118static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004119 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004120 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004121 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004122};
4123
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004124PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00004125 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004126\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004127Return a count object whose .__next__() method returns consecutive values.\n\
4128Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00004129 def count(firstval=0, step=1):\n\
Eli Benderskyc554f722013-09-03 06:37:19 -07004130 x = firstval\n\
4131 while 1:\n\
4132 yield x\n\
4133 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004134
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004135static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004136 PyVarObject_HEAD_INIT(NULL, 0)
4137 "itertools.count", /* tp_name */
4138 sizeof(countobject), /* tp_basicsize */
4139 0, /* tp_itemsize */
4140 /* methods */
4141 (destructor)count_dealloc, /* tp_dealloc */
4142 0, /* tp_print */
4143 0, /* tp_getattr */
4144 0, /* tp_setattr */
4145 0, /* tp_reserved */
4146 (reprfunc)count_repr, /* tp_repr */
4147 0, /* tp_as_number */
4148 0, /* tp_as_sequence */
4149 0, /* tp_as_mapping */
4150 0, /* tp_hash */
4151 0, /* tp_call */
4152 0, /* tp_str */
4153 PyObject_GenericGetAttr, /* tp_getattro */
4154 0, /* tp_setattro */
4155 0, /* tp_as_buffer */
4156 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004157 Py_TPFLAGS_BASETYPE, /* tp_flags */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004158 count_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004159 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004160 0, /* tp_clear */
4161 0, /* tp_richcompare */
4162 0, /* tp_weaklistoffset */
4163 PyObject_SelfIter, /* tp_iter */
4164 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004165 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004166 0, /* tp_members */
4167 0, /* tp_getset */
4168 0, /* tp_base */
4169 0, /* tp_dict */
4170 0, /* tp_descr_get */
4171 0, /* tp_descr_set */
4172 0, /* tp_dictoffset */
4173 0, /* tp_init */
4174 0, /* tp_alloc */
4175 count_new, /* tp_new */
4176 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004177};
4178
4179
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004180/* repeat object ************************************************************/
4181
4182typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004183 PyObject_HEAD
4184 PyObject *element;
4185 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004186} repeatobject;
4187
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004188static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004189
4190static PyObject *
4191repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4192{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004193 repeatobject *ro;
4194 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004195 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004196 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004197
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004198 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4199 &element, &cnt))
4200 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004201
Raymond Hettinger97d35552014-06-24 21:36:58 -07004202 if (kwds != NULL)
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004203 n_kwds = PyDict_GET_SIZE(kwds);
Raymond Hettinger97d35552014-06-24 21:36:58 -07004204 /* Does user supply times argument? */
4205 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004206 cnt = 0;
4207
4208 ro = (repeatobject *)type->tp_alloc(type, 0);
4209 if (ro == NULL)
4210 return NULL;
4211 Py_INCREF(element);
4212 ro->element = element;
4213 ro->cnt = cnt;
4214 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004215}
4216
4217static void
4218repeat_dealloc(repeatobject *ro)
4219{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004220 PyObject_GC_UnTrack(ro);
4221 Py_XDECREF(ro->element);
4222 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004223}
4224
4225static int
4226repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4227{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004228 Py_VISIT(ro->element);
4229 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004230}
4231
4232static PyObject *
4233repeat_next(repeatobject *ro)
4234{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004235 if (ro->cnt == 0)
4236 return NULL;
4237 if (ro->cnt > 0)
4238 ro->cnt--;
4239 Py_INCREF(ro->element);
4240 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004241}
4242
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004243static PyObject *
4244repeat_repr(repeatobject *ro)
4245{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004246 if (ro->cnt == -1)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004247 return PyUnicode_FromFormat("%s(%R)",
4248 _PyType_Name(Py_TYPE(ro)), ro->element);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004249 else
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004250 return PyUnicode_FromFormat("%s(%R, %zd)",
4251 _PyType_Name(Py_TYPE(ro)), ro->element,
4252 ro->cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004253}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004254
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004255static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304256repeat_len(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004257{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004258 if (ro->cnt == -1) {
4259 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4260 return NULL;
4261 }
4262 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004263}
4264
Armin Rigof5b3e362006-02-11 21:32:43 +00004265PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004266
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004267static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304268repeat_reduce(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004269{
4270 /* unpickle this so that a new repeat iterator is constructed with an
4271 * object, then call __setstate__ on it to set cnt
4272 */
4273 if (ro->cnt >= 0)
4274 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4275 else
4276 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4277}
4278
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004279static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004280 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004281 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004282 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004283};
4284
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004285PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004286"repeat(object [,times]) -> create an iterator which returns the object\n\
4287for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004288endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004289
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004290static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004291 PyVarObject_HEAD_INIT(NULL, 0)
4292 "itertools.repeat", /* tp_name */
4293 sizeof(repeatobject), /* tp_basicsize */
4294 0, /* tp_itemsize */
4295 /* methods */
4296 (destructor)repeat_dealloc, /* tp_dealloc */
4297 0, /* tp_print */
4298 0, /* tp_getattr */
4299 0, /* tp_setattr */
4300 0, /* tp_reserved */
4301 (reprfunc)repeat_repr, /* tp_repr */
4302 0, /* tp_as_number */
4303 0, /* tp_as_sequence */
4304 0, /* tp_as_mapping */
4305 0, /* tp_hash */
4306 0, /* tp_call */
4307 0, /* tp_str */
4308 PyObject_GenericGetAttr, /* tp_getattro */
4309 0, /* tp_setattro */
4310 0, /* tp_as_buffer */
4311 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4312 Py_TPFLAGS_BASETYPE, /* tp_flags */
4313 repeat_doc, /* tp_doc */
4314 (traverseproc)repeat_traverse, /* tp_traverse */
4315 0, /* tp_clear */
4316 0, /* tp_richcompare */
4317 0, /* tp_weaklistoffset */
4318 PyObject_SelfIter, /* tp_iter */
4319 (iternextfunc)repeat_next, /* tp_iternext */
4320 repeat_methods, /* tp_methods */
4321 0, /* tp_members */
4322 0, /* tp_getset */
4323 0, /* tp_base */
4324 0, /* tp_dict */
4325 0, /* tp_descr_get */
4326 0, /* tp_descr_set */
4327 0, /* tp_dictoffset */
4328 0, /* tp_init */
4329 0, /* tp_alloc */
4330 repeat_new, /* tp_new */
4331 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004332};
4333
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004334/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004335
4336typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004337 PyObject_HEAD
4338 Py_ssize_t tuplesize;
4339 Py_ssize_t numactive;
4340 PyObject *ittuple; /* tuple of iterators */
4341 PyObject *result;
4342 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004343} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004344
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004345static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004346
4347static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004348zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004349{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004350 ziplongestobject *lz;
4351 Py_ssize_t i;
4352 PyObject *ittuple; /* tuple of iterators */
4353 PyObject *result;
4354 PyObject *fillvalue = Py_None;
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004355 Py_ssize_t tuplesize;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004356
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004357 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004358 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004359 if (fillvalue == NULL || PyDict_GET_SIZE(kwds) > 1) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004360 PyErr_SetString(PyExc_TypeError,
4361 "zip_longest() got an unexpected keyword argument");
4362 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004363 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004364 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004365
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004366 /* args must be a tuple */
4367 assert(PyTuple_Check(args));
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004368 tuplesize = PyTuple_GET_SIZE(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004369
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004370 /* obtain iterators */
4371 ittuple = PyTuple_New(tuplesize);
4372 if (ittuple == NULL)
4373 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004374 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004375 PyObject *item = PyTuple_GET_ITEM(args, i);
4376 PyObject *it = PyObject_GetIter(item);
4377 if (it == NULL) {
4378 if (PyErr_ExceptionMatches(PyExc_TypeError))
4379 PyErr_Format(PyExc_TypeError,
4380 "zip_longest argument #%zd must support iteration",
4381 i+1);
4382 Py_DECREF(ittuple);
4383 return NULL;
4384 }
4385 PyTuple_SET_ITEM(ittuple, i, it);
4386 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004387
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004388 /* create a result holder */
4389 result = PyTuple_New(tuplesize);
4390 if (result == NULL) {
4391 Py_DECREF(ittuple);
4392 return NULL;
4393 }
4394 for (i=0 ; i < tuplesize ; i++) {
4395 Py_INCREF(Py_None);
4396 PyTuple_SET_ITEM(result, i, Py_None);
4397 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004398
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004399 /* create ziplongestobject structure */
4400 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4401 if (lz == NULL) {
4402 Py_DECREF(ittuple);
4403 Py_DECREF(result);
4404 return NULL;
4405 }
4406 lz->ittuple = ittuple;
4407 lz->tuplesize = tuplesize;
4408 lz->numactive = tuplesize;
4409 lz->result = result;
4410 Py_INCREF(fillvalue);
4411 lz->fillvalue = fillvalue;
4412 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004413}
4414
4415static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004416zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004417{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004418 PyObject_GC_UnTrack(lz);
4419 Py_XDECREF(lz->ittuple);
4420 Py_XDECREF(lz->result);
4421 Py_XDECREF(lz->fillvalue);
4422 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004423}
4424
4425static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004426zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004427{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004428 Py_VISIT(lz->ittuple);
4429 Py_VISIT(lz->result);
4430 Py_VISIT(lz->fillvalue);
4431 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004432}
4433
4434static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004435zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004436{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004437 Py_ssize_t i;
4438 Py_ssize_t tuplesize = lz->tuplesize;
4439 PyObject *result = lz->result;
4440 PyObject *it;
4441 PyObject *item;
4442 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004443
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004444 if (tuplesize == 0)
4445 return NULL;
4446 if (lz->numactive == 0)
4447 return NULL;
4448 if (Py_REFCNT(result) == 1) {
4449 Py_INCREF(result);
4450 for (i=0 ; i < tuplesize ; i++) {
4451 it = PyTuple_GET_ITEM(lz->ittuple, i);
4452 if (it == NULL) {
4453 Py_INCREF(lz->fillvalue);
4454 item = lz->fillvalue;
4455 } else {
4456 item = PyIter_Next(it);
4457 if (item == NULL) {
4458 lz->numactive -= 1;
4459 if (lz->numactive == 0 || PyErr_Occurred()) {
4460 lz->numactive = 0;
4461 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004462 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004463 } else {
4464 Py_INCREF(lz->fillvalue);
4465 item = lz->fillvalue;
4466 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4467 Py_DECREF(it);
4468 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004469 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004470 }
4471 olditem = PyTuple_GET_ITEM(result, i);
4472 PyTuple_SET_ITEM(result, i, item);
4473 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004474 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004475 } else {
4476 result = PyTuple_New(tuplesize);
4477 if (result == NULL)
4478 return NULL;
4479 for (i=0 ; i < tuplesize ; i++) {
4480 it = PyTuple_GET_ITEM(lz->ittuple, i);
4481 if (it == NULL) {
4482 Py_INCREF(lz->fillvalue);
4483 item = lz->fillvalue;
4484 } else {
4485 item = PyIter_Next(it);
4486 if (item == NULL) {
4487 lz->numactive -= 1;
4488 if (lz->numactive == 0 || PyErr_Occurred()) {
4489 lz->numactive = 0;
4490 Py_DECREF(result);
4491 return NULL;
4492 } else {
4493 Py_INCREF(lz->fillvalue);
4494 item = lz->fillvalue;
4495 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4496 Py_DECREF(it);
4497 }
4498 }
4499 }
4500 PyTuple_SET_ITEM(result, i, item);
4501 }
4502 }
4503 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004504}
4505
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004506static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304507zip_longest_reduce(ziplongestobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004508{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004509
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004510 /* Create a new tuple with empty sequences where appropriate to pickle.
4511 * Then use setstate to set the fillvalue
4512 */
4513 int i;
4514 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004515
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004516 if (args == NULL)
4517 return NULL;
4518 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4519 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4520 if (elem == NULL) {
4521 elem = PyTuple_New(0);
4522 if (elem == NULL) {
4523 Py_DECREF(args);
4524 return NULL;
4525 }
4526 } else
4527 Py_INCREF(elem);
4528 PyTuple_SET_ITEM(args, i, elem);
4529 }
4530 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4531}
4532
4533static PyObject *
4534zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4535{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004536 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004537 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004538 Py_RETURN_NONE;
4539}
4540
4541static PyMethodDef zip_longest_methods[] = {
4542 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4543 reduce_doc},
4544 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4545 setstate_doc},
4546 {NULL, NULL} /* sentinel */
4547};
4548
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004549PyDoc_STRVAR(zip_longest_doc,
4550"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004551\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004552Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004553the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004554method continues until the longest iterable in the argument sequence\n\
4555is exhausted and then it raises StopIteration. When the shorter iterables\n\
4556are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4557defaults to None or can be specified by a keyword argument.\n\
4558");
4559
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004560static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004561 PyVarObject_HEAD_INIT(NULL, 0)
4562 "itertools.zip_longest", /* tp_name */
4563 sizeof(ziplongestobject), /* tp_basicsize */
4564 0, /* tp_itemsize */
4565 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004566 (destructor)zip_longest_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004567 0, /* tp_print */
4568 0, /* tp_getattr */
4569 0, /* tp_setattr */
4570 0, /* tp_reserved */
4571 0, /* tp_repr */
4572 0, /* tp_as_number */
4573 0, /* tp_as_sequence */
4574 0, /* tp_as_mapping */
4575 0, /* tp_hash */
4576 0, /* tp_call */
4577 0, /* tp_str */
4578 PyObject_GenericGetAttr, /* tp_getattro */
4579 0, /* tp_setattro */
4580 0, /* tp_as_buffer */
4581 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4582 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004583 zip_longest_doc, /* tp_doc */
4584 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004585 0, /* tp_clear */
4586 0, /* tp_richcompare */
4587 0, /* tp_weaklistoffset */
4588 PyObject_SelfIter, /* tp_iter */
4589 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004590 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004591 0, /* tp_members */
4592 0, /* tp_getset */
4593 0, /* tp_base */
4594 0, /* tp_dict */
4595 0, /* tp_descr_get */
4596 0, /* tp_descr_set */
4597 0, /* tp_dictoffset */
4598 0, /* tp_init */
4599 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004600 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004601 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004602};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004603
4604/* module level code ********************************************************/
4605
4606PyDoc_STRVAR(module_doc,
4607"Functional tools for creating and using iterators.\n\
4608\n\
4609Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004610count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004611cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004612repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004613\n\
4614Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004615accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
oldkaa0735f2018-02-02 16:52:55 +08004616chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
4617chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004618compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4619dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4620groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004621filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004622islice(seq, [start,] stop [, step]) --> elements from\n\
4623 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004624starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004625tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004626takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
oldkaa0735f2018-02-02 16:52:55 +08004627zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004628\n\
4629Combinatoric generators:\n\
4630product(p, q, ... [repeat=1]) --> cartesian product\n\
4631permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004632combinations(p, r)\n\
4633combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004634");
4635
4636
Raymond Hettingerad983e72003-11-12 14:32:26 +00004637static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004638 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4639 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004640};
4641
Martin v. Löwis1a214512008-06-11 05:26:20 +00004642
4643static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004644 PyModuleDef_HEAD_INIT,
4645 "itertools",
4646 module_doc,
4647 -1,
4648 module_methods,
4649 NULL,
4650 NULL,
4651 NULL,
4652 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004653};
4654
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004655PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004656PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004657{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004658 int i;
4659 PyObject *m;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004660 const char *name;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004661 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004662 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004663 &combinations_type,
4664 &cwr_type,
4665 &cycle_type,
4666 &dropwhile_type,
4667 &takewhile_type,
4668 &islice_type,
4669 &starmap_type,
4670 &chain_type,
4671 &compress_type,
4672 &filterfalse_type,
4673 &count_type,
4674 &ziplongest_type,
4675 &permutations_type,
4676 &product_type,
4677 &repeat_type,
4678 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004679 &_grouper_type,
4680 &tee_type,
4681 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004682 NULL
4683 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004684
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004685 Py_TYPE(&teedataobject_type) = &PyType_Type;
4686 m = PyModule_Create(&itertoolsmodule);
4687 if (m == NULL)
4688 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004689
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004690 for (i=0 ; typelist[i] != NULL ; i++) {
4691 if (PyType_Ready(typelist[i]) < 0)
4692 return NULL;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004693 name = _PyType_Name(typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004694 Py_INCREF(typelist[i]);
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004695 PyModule_AddObject(m, name, (PyObject *)typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004696 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004697
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004698 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004699}