blob: c966376e472cad7429b2e3b43c67d433ea9b295b [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
Raymond Hettingerca3788c2015-08-16 14:49:24 -07002#define PY_SSIZE_T_CLEAN
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003#include "Python.h"
Fred Drake08ebfec2004-10-17 19:36:57 +00004#include "structmember.h"
Raymond Hettinger96ef8112003-02-01 00:10:11 +00005
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00006/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00007 by Raymond D. Hettinger <python@rcn.com>
Raymond Hettinger96ef8112003-02-01 00:10:11 +00008*/
9
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000010
Raymond Hettingera6ea44a2015-08-17 23:55:28 -070011/* groupby object ************************************************************/
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000012
13typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000014 PyObject_HEAD
15 PyObject *it;
16 PyObject *keyfunc;
17 PyObject *tgtkey;
18 PyObject *currkey;
19 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000020} groupbyobject;
21
22static PyTypeObject groupby_type;
23static PyObject *_grouper_create(groupbyobject *, PyObject *);
24
25static PyObject *
26groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
27{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000028 static char *kwargs[] = {"iterable", "key", NULL};
29 groupbyobject *gbo;
30 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000031
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000032 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
33 &it, &keyfunc))
34 return NULL;
35
36 gbo = (groupbyobject *)type->tp_alloc(type, 0);
37 if (gbo == NULL)
38 return NULL;
39 gbo->tgtkey = NULL;
40 gbo->currkey = NULL;
41 gbo->currvalue = NULL;
42 gbo->keyfunc = keyfunc;
43 Py_INCREF(keyfunc);
44 gbo->it = PyObject_GetIter(it);
45 if (gbo->it == NULL) {
46 Py_DECREF(gbo);
47 return NULL;
48 }
49 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000050}
51
52static void
53groupby_dealloc(groupbyobject *gbo)
54{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000055 PyObject_GC_UnTrack(gbo);
56 Py_XDECREF(gbo->it);
57 Py_XDECREF(gbo->keyfunc);
58 Py_XDECREF(gbo->tgtkey);
59 Py_XDECREF(gbo->currkey);
60 Py_XDECREF(gbo->currvalue);
61 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000062}
63
64static int
65groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
66{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000067 Py_VISIT(gbo->it);
68 Py_VISIT(gbo->keyfunc);
69 Py_VISIT(gbo->tgtkey);
70 Py_VISIT(gbo->currkey);
71 Py_VISIT(gbo->currvalue);
72 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000073}
74
75static PyObject *
76groupby_next(groupbyobject *gbo)
77{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +020078 PyObject *newvalue, *newkey, *r, *grouper;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000079
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000080 /* skip to next iteration group */
81 for (;;) {
82 if (gbo->currkey == NULL)
83 /* pass */;
84 else if (gbo->tgtkey == NULL)
85 break;
86 else {
87 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000088
Raymond Hettingera6ea44a2015-08-17 23:55:28 -070089 rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000090 if (rcmp == -1)
91 return NULL;
92 else if (rcmp == 0)
93 break;
94 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000095
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000096 newvalue = PyIter_Next(gbo->it);
97 if (newvalue == NULL)
98 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000099
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000100 if (gbo->keyfunc == Py_None) {
101 newkey = newvalue;
102 Py_INCREF(newvalue);
103 } else {
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700104 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000105 if (newkey == NULL) {
106 Py_DECREF(newvalue);
107 return NULL;
108 }
109 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000110
Serhiy Storchakaec397562016-04-06 09:50:03 +0300111 Py_XSETREF(gbo->currkey, newkey);
112 Py_XSETREF(gbo->currvalue, newvalue);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000113 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000114
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000115 Py_INCREF(gbo->currkey);
Serhiy Storchakaec397562016-04-06 09:50:03 +0300116 Py_XSETREF(gbo->tgtkey, gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000117
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000118 grouper = _grouper_create(gbo, gbo->tgtkey);
119 if (grouper == NULL)
120 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000121
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000122 r = PyTuple_Pack(2, gbo->currkey, grouper);
123 Py_DECREF(grouper);
124 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000125}
126
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000127static PyObject *
128groupby_reduce(groupbyobject *lz)
129{
130 /* reduce as a 'new' call with an optional 'setstate' if groupby
131 * has started
132 */
133 PyObject *value;
134 if (lz->tgtkey && lz->currkey && lz->currvalue)
135 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
136 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
137 else
138 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
139 lz->it, lz->keyfunc);
140
141 return value;
142}
143
144PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
145
146static PyObject *
147groupby_setstate(groupbyobject *lz, PyObject *state)
148{
149 PyObject *currkey, *currvalue, *tgtkey;
150 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey))
151 return NULL;
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200152 Py_INCREF(currkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300153 Py_XSETREF(lz->currkey, currkey);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200154 Py_INCREF(currvalue);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300155 Py_XSETREF(lz->currvalue, currvalue);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200156 Py_INCREF(tgtkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300157 Py_XSETREF(lz->tgtkey, tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000158 Py_RETURN_NONE;
159}
160
161PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
162
163static PyMethodDef groupby_methods[] = {
164 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
165 reduce_doc},
166 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
167 setstate_doc},
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700168 {NULL, NULL} /* sentinel */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000169};
170
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000171PyDoc_STRVAR(groupby_doc,
172"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
173(key, sub-iterator) grouped by each value of key(value).\n");
174
175static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000176 PyVarObject_HEAD_INIT(NULL, 0)
177 "itertools.groupby", /* tp_name */
178 sizeof(groupbyobject), /* tp_basicsize */
179 0, /* tp_itemsize */
180 /* methods */
181 (destructor)groupby_dealloc, /* tp_dealloc */
182 0, /* tp_print */
183 0, /* tp_getattr */
184 0, /* tp_setattr */
185 0, /* tp_reserved */
186 0, /* tp_repr */
187 0, /* tp_as_number */
188 0, /* tp_as_sequence */
189 0, /* tp_as_mapping */
190 0, /* tp_hash */
191 0, /* tp_call */
192 0, /* tp_str */
193 PyObject_GenericGetAttr, /* tp_getattro */
194 0, /* tp_setattro */
195 0, /* tp_as_buffer */
196 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
197 Py_TPFLAGS_BASETYPE, /* tp_flags */
198 groupby_doc, /* tp_doc */
199 (traverseproc)groupby_traverse, /* tp_traverse */
200 0, /* tp_clear */
201 0, /* tp_richcompare */
202 0, /* tp_weaklistoffset */
203 PyObject_SelfIter, /* tp_iter */
204 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000205 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000206 0, /* tp_members */
207 0, /* tp_getset */
208 0, /* tp_base */
209 0, /* tp_dict */
210 0, /* tp_descr_get */
211 0, /* tp_descr_set */
212 0, /* tp_dictoffset */
213 0, /* tp_init */
214 0, /* tp_alloc */
215 groupby_new, /* tp_new */
216 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000217};
218
219
220/* _grouper object (internal) ************************************************/
221
222typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000223 PyObject_HEAD
224 PyObject *parent;
225 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000226} _grouperobject;
227
228static PyTypeObject _grouper_type;
229
230static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000231_grouper_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
232{
233 PyObject *parent, *tgtkey;
234
235 if (!PyArg_ParseTuple(args, "O!O", &groupby_type, &parent, &tgtkey))
236 return NULL;
237
238 return _grouper_create((groupbyobject*) parent, tgtkey);
239}
240
241static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000242_grouper_create(groupbyobject *parent, PyObject *tgtkey)
243{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000244 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000245
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000246 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
247 if (igo == NULL)
248 return NULL;
249 igo->parent = (PyObject *)parent;
250 Py_INCREF(parent);
251 igo->tgtkey = tgtkey;
252 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000253
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000254 PyObject_GC_Track(igo);
255 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000256}
257
258static void
259_grouper_dealloc(_grouperobject *igo)
260{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000261 PyObject_GC_UnTrack(igo);
262 Py_DECREF(igo->parent);
263 Py_DECREF(igo->tgtkey);
264 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000265}
266
267static int
268_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
269{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000270 Py_VISIT(igo->parent);
271 Py_VISIT(igo->tgtkey);
272 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000273}
274
275static PyObject *
276_grouper_next(_grouperobject *igo)
277{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000278 groupbyobject *gbo = (groupbyobject *)igo->parent;
279 PyObject *newvalue, *newkey, *r;
280 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000281
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000282 if (gbo->currvalue == NULL) {
283 newvalue = PyIter_Next(gbo->it);
284 if (newvalue == NULL)
285 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000286
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000287 if (gbo->keyfunc == Py_None) {
288 newkey = newvalue;
289 Py_INCREF(newvalue);
290 } else {
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700291 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000292 if (newkey == NULL) {
293 Py_DECREF(newvalue);
294 return NULL;
295 }
296 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000297
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000298 assert(gbo->currkey == NULL);
299 gbo->currkey = newkey;
300 gbo->currvalue = newvalue;
301 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000302
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000303 assert(gbo->currkey != NULL);
304 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
305 if (rcmp <= 0)
306 /* got any error or current group is end */
307 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000308
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000309 r = gbo->currvalue;
310 gbo->currvalue = NULL;
311 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000312
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000313 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000314}
315
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000316static PyObject *
317_grouper_reduce(_grouperobject *lz)
318{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700319 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000320}
321
322static PyMethodDef _grouper_methods[] = {
323 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
324 reduce_doc},
325 {NULL, NULL} /* sentinel */
326};
327
328
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000329static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000330 PyVarObject_HEAD_INIT(NULL, 0)
331 "itertools._grouper", /* tp_name */
332 sizeof(_grouperobject), /* tp_basicsize */
333 0, /* tp_itemsize */
334 /* methods */
335 (destructor)_grouper_dealloc, /* tp_dealloc */
336 0, /* tp_print */
337 0, /* tp_getattr */
338 0, /* tp_setattr */
339 0, /* tp_reserved */
340 0, /* tp_repr */
341 0, /* tp_as_number */
342 0, /* tp_as_sequence */
343 0, /* tp_as_mapping */
344 0, /* tp_hash */
345 0, /* tp_call */
346 0, /* tp_str */
347 PyObject_GenericGetAttr, /* tp_getattro */
348 0, /* tp_setattro */
349 0, /* tp_as_buffer */
350 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
351 0, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700352 (traverseproc)_grouper_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000353 0, /* tp_clear */
354 0, /* tp_richcompare */
355 0, /* tp_weaklistoffset */
356 PyObject_SelfIter, /* tp_iter */
357 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000358 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000359 0, /* tp_members */
360 0, /* tp_getset */
361 0, /* tp_base */
362 0, /* tp_dict */
363 0, /* tp_descr_get */
364 0, /* tp_descr_set */
365 0, /* tp_dictoffset */
366 0, /* tp_init */
367 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000368 _grouper_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000369 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000370};
371
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000372
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700373/* tee object and with supporting function and objects ***********************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000374
Raymond Hettingerad983e72003-11-12 14:32:26 +0000375/* The teedataobject pre-allocates space for LINKCELLS number of objects.
376 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000377 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000378 the other structure members including PyHEAD overhead. The larger the
379 value, the less memory overhead per object and the less time spent
380 allocating/deallocating new links. The smaller the number, the less
381 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000382*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000383#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000384
385typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000386 PyObject_HEAD
387 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700388 int numread; /* 0 <= numread <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000389 PyObject *nextlink;
390 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000391} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000392
393typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000394 PyObject_HEAD
395 teedataobject *dataobj;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700396 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000397 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000398} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000399
Raymond Hettingerad983e72003-11-12 14:32:26 +0000400static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000401
402static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000403teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000404{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000405 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000406
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000407 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
408 if (tdo == NULL)
409 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000410
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000411 tdo->numread = 0;
412 tdo->nextlink = NULL;
413 Py_INCREF(it);
414 tdo->it = it;
415 PyObject_GC_Track(tdo);
416 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000417}
418
419static PyObject *
420teedataobject_jumplink(teedataobject *tdo)
421{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000422 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000423 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000424 Py_XINCREF(tdo->nextlink);
425 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000426}
427
428static PyObject *
429teedataobject_getitem(teedataobject *tdo, int i)
430{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000431 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000432
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000433 assert(i < LINKCELLS);
434 if (i < tdo->numread)
435 value = tdo->values[i];
436 else {
437 /* this is the lead iterator, so fetch more data */
438 assert(i == tdo->numread);
439 value = PyIter_Next(tdo->it);
440 if (value == NULL)
441 return NULL;
442 tdo->numread++;
443 tdo->values[i] = value;
444 }
445 Py_INCREF(value);
446 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000447}
448
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000449static int
450teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
451{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000452 int i;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700453
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000454 Py_VISIT(tdo->it);
455 for (i = 0; i < tdo->numread; i++)
456 Py_VISIT(tdo->values[i]);
457 Py_VISIT(tdo->nextlink);
458 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000459}
460
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200461static void
462teedataobject_safe_decref(PyObject *obj)
463{
464 while (obj && Py_TYPE(obj) == &teedataobject_type &&
465 Py_REFCNT(obj) == 1) {
466 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
467 ((teedataobject *)obj)->nextlink = NULL;
468 Py_DECREF(obj);
469 obj = nextlink;
470 }
471 Py_XDECREF(obj);
472}
473
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000474static int
475teedataobject_clear(teedataobject *tdo)
476{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000477 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200478 PyObject *tmp;
479
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000480 Py_CLEAR(tdo->it);
481 for (i=0 ; i<tdo->numread ; i++)
482 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200483 tmp = tdo->nextlink;
484 tdo->nextlink = NULL;
485 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000486 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000487}
488
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000489static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000490teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000491{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000492 PyObject_GC_UnTrack(tdo);
493 teedataobject_clear(tdo);
494 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000495}
496
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000497static PyObject *
498teedataobject_reduce(teedataobject *tdo)
499{
500 int i;
501 /* create a temporary list of already iterated values */
502 PyObject *values = PyList_New(tdo->numread);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700503
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000504 if (!values)
505 return NULL;
506 for (i=0 ; i<tdo->numread ; i++) {
507 Py_INCREF(tdo->values[i]);
508 PyList_SET_ITEM(values, i, tdo->values[i]);
509 }
510 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
511 values,
512 tdo->nextlink ? tdo->nextlink : Py_None);
513}
514
515static PyTypeObject teedataobject_type;
516
517static PyObject *
518teedataobject_new(PyTypeObject *type, PyObject *args, PyObject *kw)
519{
520 teedataobject *tdo;
521 PyObject *it, *values, *next;
522 Py_ssize_t i, len;
523
524 assert(type == &teedataobject_type);
525 if (!PyArg_ParseTuple(args, "OO!O", &it, &PyList_Type, &values, &next))
526 return NULL;
527
528 tdo = (teedataobject *)teedataobject_newinternal(it);
529 if (!tdo)
530 return NULL;
531
532 len = PyList_GET_SIZE(values);
533 if (len > LINKCELLS)
534 goto err;
535 for (i=0; i<len; i++) {
536 tdo->values[i] = PyList_GET_ITEM(values, i);
537 Py_INCREF(tdo->values[i]);
538 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200539 /* len <= LINKCELLS < INT_MAX */
540 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000541
542 if (len == LINKCELLS) {
543 if (next != Py_None) {
544 if (Py_TYPE(next) != &teedataobject_type)
545 goto err;
546 assert(tdo->nextlink == NULL);
547 Py_INCREF(next);
548 tdo->nextlink = next;
549 }
550 } else {
551 if (next != Py_None)
552 goto err; /* shouldn't have a next if we are not full */
553 }
554 return (PyObject*)tdo;
555
556err:
557 Py_XDECREF(tdo);
558 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
559 return NULL;
560}
561
562static PyMethodDef teedataobject_methods[] = {
563 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
564 reduce_doc},
565 {NULL, NULL} /* sentinel */
566};
567
Raymond Hettingerad983e72003-11-12 14:32:26 +0000568PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000569
Raymond Hettingerad983e72003-11-12 14:32:26 +0000570static PyTypeObject teedataobject_type = {
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700571 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000572 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000573 sizeof(teedataobject), /* tp_basicsize */
574 0, /* tp_itemsize */
575 /* methods */
576 (destructor)teedataobject_dealloc, /* tp_dealloc */
577 0, /* tp_print */
578 0, /* tp_getattr */
579 0, /* tp_setattr */
580 0, /* tp_reserved */
581 0, /* tp_repr */
582 0, /* tp_as_number */
583 0, /* tp_as_sequence */
584 0, /* tp_as_mapping */
585 0, /* tp_hash */
586 0, /* tp_call */
587 0, /* tp_str */
588 PyObject_GenericGetAttr, /* tp_getattro */
589 0, /* tp_setattro */
590 0, /* tp_as_buffer */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700591 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000592 teedataobject_doc, /* tp_doc */
593 (traverseproc)teedataobject_traverse, /* tp_traverse */
594 (inquiry)teedataobject_clear, /* tp_clear */
595 0, /* tp_richcompare */
596 0, /* tp_weaklistoffset */
597 0, /* tp_iter */
598 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000599 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000600 0, /* tp_members */
601 0, /* tp_getset */
602 0, /* tp_base */
603 0, /* tp_dict */
604 0, /* tp_descr_get */
605 0, /* tp_descr_set */
606 0, /* tp_dictoffset */
607 0, /* tp_init */
608 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000609 teedataobject_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000610 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000611};
612
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000613
614static PyTypeObject tee_type;
615
616static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000617tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000618{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000619 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000620
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000621 if (to->index >= LINKCELLS) {
622 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200623 if (link == NULL)
624 return NULL;
Serhiy Storchaka57a01d32016-04-10 18:05:40 +0300625 Py_SETREF(to->dataobj, (teedataobject *)link);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000626 to->index = 0;
627 }
628 value = teedataobject_getitem(to->dataobj, to->index);
629 if (value == NULL)
630 return NULL;
631 to->index++;
632 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000633}
634
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000635static int
636tee_traverse(teeobject *to, visitproc visit, void *arg)
637{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000638 Py_VISIT((PyObject *)to->dataobj);
639 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000640}
641
Raymond Hettingerad983e72003-11-12 14:32:26 +0000642static PyObject *
643tee_copy(teeobject *to)
644{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000645 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000646
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000647 newto = PyObject_GC_New(teeobject, &tee_type);
648 if (newto == NULL)
649 return NULL;
650 Py_INCREF(to->dataobj);
651 newto->dataobj = to->dataobj;
652 newto->index = to->index;
653 newto->weakreflist = NULL;
654 PyObject_GC_Track(newto);
655 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000656}
657
658PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
659
660static PyObject *
661tee_fromiterable(PyObject *iterable)
662{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000663 teeobject *to;
664 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000665
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000666 it = PyObject_GetIter(iterable);
667 if (it == NULL)
668 return NULL;
669 if (PyObject_TypeCheck(it, &tee_type)) {
670 to = (teeobject *)tee_copy((teeobject *)it);
671 goto done;
672 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000673
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000674 to = PyObject_GC_New(teeobject, &tee_type);
675 if (to == NULL)
676 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000677 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000678 if (!to->dataobj) {
679 PyObject_GC_Del(to);
680 to = NULL;
681 goto done;
682 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000683
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000684 to->index = 0;
685 to->weakreflist = NULL;
686 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000687done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000688 Py_XDECREF(it);
689 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000690}
691
692static PyObject *
693tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
694{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000695 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000696
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000697 if (!PyArg_UnpackTuple(args, "_tee", 1, 1, &iterable))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000698 return NULL;
699 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000700}
701
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000702static int
703tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000704{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000705 if (to->weakreflist != NULL)
706 PyObject_ClearWeakRefs((PyObject *) to);
707 Py_CLEAR(to->dataobj);
708 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000709}
710
711static void
712tee_dealloc(teeobject *to)
713{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000714 PyObject_GC_UnTrack(to);
715 tee_clear(to);
716 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000717}
718
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000719static PyObject *
720tee_reduce(teeobject *to)
721{
722 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
723}
724
725static PyObject *
726tee_setstate(teeobject *to, PyObject *state)
727{
728 teedataobject *tdo;
729 int index;
730 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index))
731 return NULL;
732 if (index < 0 || index > LINKCELLS) {
733 PyErr_SetString(PyExc_ValueError, "Index out of range");
734 return NULL;
735 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200736 Py_INCREF(tdo);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300737 Py_XSETREF(to->dataobj, tdo);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000738 to->index = index;
739 Py_RETURN_NONE;
740}
741
Raymond Hettingerad983e72003-11-12 14:32:26 +0000742PyDoc_STRVAR(teeobject_doc,
743"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000744
Raymond Hettingerad983e72003-11-12 14:32:26 +0000745static PyMethodDef tee_methods[] = {
Raymond Hettingera6a2d442015-08-16 14:38:07 -0700746 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
747 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
748 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000749 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000750};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000751
752static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000753 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000754 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000755 sizeof(teeobject), /* tp_basicsize */
756 0, /* tp_itemsize */
757 /* methods */
758 (destructor)tee_dealloc, /* tp_dealloc */
759 0, /* tp_print */
760 0, /* tp_getattr */
761 0, /* tp_setattr */
762 0, /* tp_reserved */
763 0, /* tp_repr */
764 0, /* tp_as_number */
765 0, /* tp_as_sequence */
766 0, /* tp_as_mapping */
767 0, /* tp_hash */
768 0, /* tp_call */
769 0, /* tp_str */
770 0, /* tp_getattro */
771 0, /* tp_setattro */
772 0, /* tp_as_buffer */
773 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
774 teeobject_doc, /* tp_doc */
775 (traverseproc)tee_traverse, /* tp_traverse */
776 (inquiry)tee_clear, /* tp_clear */
777 0, /* tp_richcompare */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700778 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000779 PyObject_SelfIter, /* tp_iter */
780 (iternextfunc)tee_next, /* tp_iternext */
781 tee_methods, /* tp_methods */
782 0, /* tp_members */
783 0, /* tp_getset */
784 0, /* tp_base */
785 0, /* tp_dict */
786 0, /* tp_descr_get */
787 0, /* tp_descr_set */
788 0, /* tp_dictoffset */
789 0, /* tp_init */
790 0, /* tp_alloc */
791 tee_new, /* tp_new */
792 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000793};
794
Raymond Hettingerad983e72003-11-12 14:32:26 +0000795static PyObject *
796tee(PyObject *self, PyObject *args)
797{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000798 Py_ssize_t i, n=2;
799 PyObject *it, *iterable, *copyable, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200800 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000801
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000802 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
803 return NULL;
804 if (n < 0) {
805 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
806 return NULL;
807 }
808 result = PyTuple_New(n);
809 if (result == NULL)
810 return NULL;
811 if (n == 0)
812 return result;
813 it = PyObject_GetIter(iterable);
814 if (it == NULL) {
815 Py_DECREF(result);
816 return NULL;
817 }
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200818 if (!_PyObject_HasAttrId(it, &PyId___copy__)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000819 copyable = tee_fromiterable(it);
820 Py_DECREF(it);
821 if (copyable == NULL) {
822 Py_DECREF(result);
823 return NULL;
824 }
825 } else
826 copyable = it;
827 PyTuple_SET_ITEM(result, 0, copyable);
828 for (i=1 ; i<n ; i++) {
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200829
830 copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000831 if (copyable == NULL) {
832 Py_DECREF(result);
833 return NULL;
834 }
835 PyTuple_SET_ITEM(result, i, copyable);
836 }
837 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000838}
839
840PyDoc_STRVAR(tee_doc,
841"tee(iterable, n=2) --> tuple of n independent iterators.");
842
843
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700844/* cycle object **************************************************************/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000845
846typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000847 PyObject_HEAD
848 PyObject *it;
849 PyObject *saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700850 Py_ssize_t index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000851 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000852} cycleobject;
853
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000854static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000855
856static PyObject *
857cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
858{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000859 PyObject *it;
860 PyObject *iterable;
861 PyObject *saved;
862 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000863
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000864 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
865 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000866
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000867 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
868 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000869
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000870 /* Get iterator. */
871 it = PyObject_GetIter(iterable);
872 if (it == NULL)
873 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000874
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000875 saved = PyList_New(0);
876 if (saved == NULL) {
877 Py_DECREF(it);
878 return NULL;
879 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000880
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000881 /* create cycleobject structure */
882 lz = (cycleobject *)type->tp_alloc(type, 0);
883 if (lz == NULL) {
884 Py_DECREF(it);
885 Py_DECREF(saved);
886 return NULL;
887 }
888 lz->it = it;
889 lz->saved = saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700890 lz->index = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000891 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000892
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000893 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000894}
895
896static void
897cycle_dealloc(cycleobject *lz)
898{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000899 PyObject_GC_UnTrack(lz);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000900 Py_XDECREF(lz->it);
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700901 Py_XDECREF(lz->saved);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000902 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000903}
904
905static int
906cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
907{
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700908 if (lz->it)
909 Py_VISIT(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000910 Py_VISIT(lz->saved);
911 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000912}
913
914static PyObject *
915cycle_next(cycleobject *lz)
916{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000917 PyObject *item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000918
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700919 if (lz->it != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000920 item = PyIter_Next(lz->it);
921 if (item != NULL) {
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700922 if (lz->firstpass)
923 return item;
924 if (PyList_Append(lz->saved, item)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000925 Py_DECREF(item);
926 return NULL;
927 }
928 return item;
929 }
Raymond Hettinger98958fe2015-08-15 15:09:30 -0700930 /* Note: StopIteration is already cleared by PyIter_Next() */
931 if (PyErr_Occurred())
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000932 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700933 Py_CLEAR(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000934 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700935 if (Py_SIZE(lz->saved) == 0)
936 return NULL;
937 item = PyList_GET_ITEM(lz->saved, lz->index);
938 lz->index++;
939 if (lz->index >= Py_SIZE(lz->saved))
940 lz->index = 0;
941 Py_INCREF(item);
942 return item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000943}
944
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000945static PyObject *
946cycle_reduce(cycleobject *lz)
947{
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700948 /* Create a new cycle with the iterator tuple, then set the saved state */
949 if (lz->it == NULL) {
950 PyObject *it = PyObject_GetIter(lz->saved);
951 if (it == NULL)
952 return NULL;
953 if (lz->index != 0) {
954 _Py_IDENTIFIER(__setstate__);
955 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
956 "n", lz->index);
957 if (res == NULL) {
958 Py_DECREF(it);
959 return NULL;
960 }
961 Py_DECREF(res);
962 }
963 return Py_BuildValue("O(N)(Oi)", Py_TYPE(lz), it, lz->saved, 1);
964 }
965 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), lz->it, lz->saved,
966 lz->firstpass);
Raymond Hettingerc39786d2015-08-15 15:16:12 -0700967}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000968
969static PyObject *
970cycle_setstate(cycleobject *lz, PyObject *state)
971{
972 PyObject *saved=NULL;
973 int firstpass;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700974
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700975 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000976 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700977 Py_INCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300978 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000979 lz->firstpass = firstpass != 0;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700980 lz->index = 0;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000981 Py_RETURN_NONE;
982}
983
984static PyMethodDef cycle_methods[] = {
985 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
986 reduce_doc},
987 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
988 setstate_doc},
989 {NULL, NULL} /* sentinel */
990};
991
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000992PyDoc_STRVAR(cycle_doc,
993"cycle(iterable) --> cycle object\n\
994\n\
995Return elements from the iterable until it is exhausted.\n\
996Then repeat the sequence indefinitely.");
997
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000998static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000999 PyVarObject_HEAD_INIT(NULL, 0)
1000 "itertools.cycle", /* tp_name */
1001 sizeof(cycleobject), /* tp_basicsize */
1002 0, /* tp_itemsize */
1003 /* methods */
1004 (destructor)cycle_dealloc, /* tp_dealloc */
1005 0, /* tp_print */
1006 0, /* tp_getattr */
1007 0, /* tp_setattr */
1008 0, /* tp_reserved */
1009 0, /* tp_repr */
1010 0, /* tp_as_number */
1011 0, /* tp_as_sequence */
1012 0, /* tp_as_mapping */
1013 0, /* tp_hash */
1014 0, /* tp_call */
1015 0, /* tp_str */
1016 PyObject_GenericGetAttr, /* tp_getattro */
1017 0, /* tp_setattro */
1018 0, /* tp_as_buffer */
1019 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1020 Py_TPFLAGS_BASETYPE, /* tp_flags */
1021 cycle_doc, /* tp_doc */
1022 (traverseproc)cycle_traverse, /* tp_traverse */
1023 0, /* tp_clear */
1024 0, /* tp_richcompare */
1025 0, /* tp_weaklistoffset */
1026 PyObject_SelfIter, /* tp_iter */
1027 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001028 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001029 0, /* tp_members */
1030 0, /* tp_getset */
1031 0, /* tp_base */
1032 0, /* tp_dict */
1033 0, /* tp_descr_get */
1034 0, /* tp_descr_set */
1035 0, /* tp_dictoffset */
1036 0, /* tp_init */
1037 0, /* tp_alloc */
1038 cycle_new, /* tp_new */
1039 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001040};
1041
1042
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001043/* dropwhile object **********************************************************/
1044
1045typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001046 PyObject_HEAD
1047 PyObject *func;
1048 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001049 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001050} dropwhileobject;
1051
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001052static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053
1054static PyObject *
1055dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1056{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001057 PyObject *func, *seq;
1058 PyObject *it;
1059 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001060
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001061 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
1062 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001063
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001064 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
1065 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001066
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001067 /* Get iterator. */
1068 it = PyObject_GetIter(seq);
1069 if (it == NULL)
1070 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001071
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001072 /* create dropwhileobject structure */
1073 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1074 if (lz == NULL) {
1075 Py_DECREF(it);
1076 return NULL;
1077 }
1078 Py_INCREF(func);
1079 lz->func = func;
1080 lz->it = it;
1081 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001082
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001083 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001084}
1085
1086static void
1087dropwhile_dealloc(dropwhileobject *lz)
1088{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001089 PyObject_GC_UnTrack(lz);
1090 Py_XDECREF(lz->func);
1091 Py_XDECREF(lz->it);
1092 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001093}
1094
1095static int
1096dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1097{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001098 Py_VISIT(lz->it);
1099 Py_VISIT(lz->func);
1100 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001101}
1102
1103static PyObject *
1104dropwhile_next(dropwhileobject *lz)
1105{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001106 PyObject *item, *good;
1107 PyObject *it = lz->it;
1108 long ok;
1109 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001110
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001111 iternext = *Py_TYPE(it)->tp_iternext;
1112 for (;;) {
1113 item = iternext(it);
1114 if (item == NULL)
1115 return NULL;
1116 if (lz->start == 1)
1117 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001118
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001119 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1120 if (good == NULL) {
1121 Py_DECREF(item);
1122 return NULL;
1123 }
1124 ok = PyObject_IsTrue(good);
1125 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001126 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001127 lz->start = 1;
1128 return item;
1129 }
1130 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001131 if (ok < 0)
1132 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001133 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001134}
1135
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001136static PyObject *
1137dropwhile_reduce(dropwhileobject *lz)
1138{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001139 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 +00001140}
1141
1142static PyObject *
1143dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1144{
1145 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001146 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001147 return NULL;
1148 lz->start = start;
1149 Py_RETURN_NONE;
1150}
1151
1152static PyMethodDef dropwhile_methods[] = {
1153 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1154 reduce_doc},
1155 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1156 setstate_doc},
1157 {NULL, NULL} /* sentinel */
1158};
1159
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001160PyDoc_STRVAR(dropwhile_doc,
1161"dropwhile(predicate, iterable) --> dropwhile object\n\
1162\n\
1163Drop items from the iterable while predicate(item) is true.\n\
1164Afterwards, return every element until the iterable is exhausted.");
1165
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001166static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001167 PyVarObject_HEAD_INIT(NULL, 0)
1168 "itertools.dropwhile", /* tp_name */
1169 sizeof(dropwhileobject), /* tp_basicsize */
1170 0, /* tp_itemsize */
1171 /* methods */
1172 (destructor)dropwhile_dealloc, /* tp_dealloc */
1173 0, /* tp_print */
1174 0, /* tp_getattr */
1175 0, /* tp_setattr */
1176 0, /* tp_reserved */
1177 0, /* tp_repr */
1178 0, /* tp_as_number */
1179 0, /* tp_as_sequence */
1180 0, /* tp_as_mapping */
1181 0, /* tp_hash */
1182 0, /* tp_call */
1183 0, /* tp_str */
1184 PyObject_GenericGetAttr, /* tp_getattro */
1185 0, /* tp_setattro */
1186 0, /* tp_as_buffer */
1187 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1188 Py_TPFLAGS_BASETYPE, /* tp_flags */
1189 dropwhile_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001190 (traverseproc)dropwhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001191 0, /* tp_clear */
1192 0, /* tp_richcompare */
1193 0, /* tp_weaklistoffset */
1194 PyObject_SelfIter, /* tp_iter */
1195 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001196 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001197 0, /* tp_members */
1198 0, /* tp_getset */
1199 0, /* tp_base */
1200 0, /* tp_dict */
1201 0, /* tp_descr_get */
1202 0, /* tp_descr_set */
1203 0, /* tp_dictoffset */
1204 0, /* tp_init */
1205 0, /* tp_alloc */
1206 dropwhile_new, /* tp_new */
1207 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001208};
1209
1210
1211/* takewhile object **********************************************************/
1212
1213typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001214 PyObject_HEAD
1215 PyObject *func;
1216 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001217 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001218} takewhileobject;
1219
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001220static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001221
1222static PyObject *
1223takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1224{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001225 PyObject *func, *seq;
1226 PyObject *it;
1227 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001228
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001229 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1230 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001231
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001232 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1233 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001234
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001235 /* Get iterator. */
1236 it = PyObject_GetIter(seq);
1237 if (it == NULL)
1238 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001239
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001240 /* create takewhileobject structure */
1241 lz = (takewhileobject *)type->tp_alloc(type, 0);
1242 if (lz == NULL) {
1243 Py_DECREF(it);
1244 return NULL;
1245 }
1246 Py_INCREF(func);
1247 lz->func = func;
1248 lz->it = it;
1249 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001250
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001251 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001252}
1253
1254static void
1255takewhile_dealloc(takewhileobject *lz)
1256{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001257 PyObject_GC_UnTrack(lz);
1258 Py_XDECREF(lz->func);
1259 Py_XDECREF(lz->it);
1260 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001261}
1262
1263static int
1264takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1265{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001266 Py_VISIT(lz->it);
1267 Py_VISIT(lz->func);
1268 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001269}
1270
1271static PyObject *
1272takewhile_next(takewhileobject *lz)
1273{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001274 PyObject *item, *good;
1275 PyObject *it = lz->it;
1276 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001277
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001278 if (lz->stop == 1)
1279 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001280
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001281 item = (*Py_TYPE(it)->tp_iternext)(it);
1282 if (item == NULL)
1283 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001284
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001285 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1286 if (good == NULL) {
1287 Py_DECREF(item);
1288 return NULL;
1289 }
1290 ok = PyObject_IsTrue(good);
1291 Py_DECREF(good);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001292 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001293 return item;
1294 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001295 if (ok == 0)
1296 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001297 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001298}
1299
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001300static PyObject *
1301takewhile_reduce(takewhileobject *lz)
1302{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001303 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 +00001304}
1305
1306static PyObject *
1307takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1308{
1309 int stop = PyObject_IsTrue(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001310
Antoine Pitrou721738f2012-08-15 23:20:39 +02001311 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001312 return NULL;
1313 lz->stop = stop;
1314 Py_RETURN_NONE;
1315}
1316
1317static PyMethodDef takewhile_reduce_methods[] = {
1318 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1319 reduce_doc},
1320 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1321 setstate_doc},
1322 {NULL, NULL} /* sentinel */
1323};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001324PyDoc_STRVAR(takewhile_doc,
1325"takewhile(predicate, iterable) --> takewhile object\n\
1326\n\
1327Return successive entries from an iterable as long as the \n\
1328predicate evaluates to true for each entry.");
1329
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001330static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001331 PyVarObject_HEAD_INIT(NULL, 0)
1332 "itertools.takewhile", /* tp_name */
1333 sizeof(takewhileobject), /* tp_basicsize */
1334 0, /* tp_itemsize */
1335 /* methods */
1336 (destructor)takewhile_dealloc, /* tp_dealloc */
1337 0, /* tp_print */
1338 0, /* tp_getattr */
1339 0, /* tp_setattr */
1340 0, /* tp_reserved */
1341 0, /* tp_repr */
1342 0, /* tp_as_number */
1343 0, /* tp_as_sequence */
1344 0, /* tp_as_mapping */
1345 0, /* tp_hash */
1346 0, /* tp_call */
1347 0, /* tp_str */
1348 PyObject_GenericGetAttr, /* tp_getattro */
1349 0, /* tp_setattro */
1350 0, /* tp_as_buffer */
1351 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1352 Py_TPFLAGS_BASETYPE, /* tp_flags */
1353 takewhile_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001354 (traverseproc)takewhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001355 0, /* tp_clear */
1356 0, /* tp_richcompare */
1357 0, /* tp_weaklistoffset */
1358 PyObject_SelfIter, /* tp_iter */
1359 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001360 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001361 0, /* tp_members */
1362 0, /* tp_getset */
1363 0, /* tp_base */
1364 0, /* tp_dict */
1365 0, /* tp_descr_get */
1366 0, /* tp_descr_set */
1367 0, /* tp_dictoffset */
1368 0, /* tp_init */
1369 0, /* tp_alloc */
1370 takewhile_new, /* tp_new */
1371 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001372};
1373
1374
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001375/* islice object *************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001376
1377typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001378 PyObject_HEAD
1379 PyObject *it;
1380 Py_ssize_t next;
1381 Py_ssize_t stop;
1382 Py_ssize_t step;
1383 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001384} isliceobject;
1385
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001386static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001387
1388static PyObject *
1389islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1390{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001391 PyObject *seq;
1392 Py_ssize_t start=0, stop=-1, step=1;
1393 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1394 Py_ssize_t numargs;
1395 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001396
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001397 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1398 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001399
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001400 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1401 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001402
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001403 numargs = PyTuple_Size(args);
1404 if (numargs == 2) {
1405 if (a1 != Py_None) {
1406 stop = PyLong_AsSsize_t(a1);
1407 if (stop == -1) {
1408 if (PyErr_Occurred())
1409 PyErr_Clear();
1410 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001411 "Stop argument for islice() must be None or "
1412 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001413 return NULL;
1414 }
1415 }
1416 } else {
1417 if (a1 != Py_None)
1418 start = PyLong_AsSsize_t(a1);
1419 if (start == -1 && PyErr_Occurred())
1420 PyErr_Clear();
1421 if (a2 != Py_None) {
1422 stop = PyLong_AsSsize_t(a2);
1423 if (stop == -1) {
1424 if (PyErr_Occurred())
1425 PyErr_Clear();
1426 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001427 "Stop argument for islice() must be None or "
1428 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001429 return NULL;
1430 }
1431 }
1432 }
1433 if (start<0 || stop<-1) {
1434 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001435 "Indices for islice() must be None or "
1436 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001437 return NULL;
1438 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001439
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001440 if (a3 != NULL) {
1441 if (a3 != Py_None)
1442 step = PyLong_AsSsize_t(a3);
1443 if (step == -1 && PyErr_Occurred())
1444 PyErr_Clear();
1445 }
1446 if (step<1) {
1447 PyErr_SetString(PyExc_ValueError,
1448 "Step for islice() must be a positive integer or None.");
1449 return NULL;
1450 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001451
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001452 /* Get iterator. */
1453 it = PyObject_GetIter(seq);
1454 if (it == NULL)
1455 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001456
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001457 /* create isliceobject structure */
1458 lz = (isliceobject *)type->tp_alloc(type, 0);
1459 if (lz == NULL) {
1460 Py_DECREF(it);
1461 return NULL;
1462 }
1463 lz->it = it;
1464 lz->next = start;
1465 lz->stop = stop;
1466 lz->step = step;
1467 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001468
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001469 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001470}
1471
1472static void
1473islice_dealloc(isliceobject *lz)
1474{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001475 PyObject_GC_UnTrack(lz);
1476 Py_XDECREF(lz->it);
1477 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001478}
1479
1480static int
1481islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1482{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001483 Py_VISIT(lz->it);
1484 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001485}
1486
1487static PyObject *
1488islice_next(isliceobject *lz)
1489{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001490 PyObject *item;
1491 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001492 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001493 Py_ssize_t oldnext;
1494 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001495
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001496 if (it == NULL)
1497 return NULL;
1498
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001499 iternext = *Py_TYPE(it)->tp_iternext;
1500 while (lz->cnt < lz->next) {
1501 item = iternext(it);
1502 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001503 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001504 Py_DECREF(item);
1505 lz->cnt++;
1506 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001507 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001508 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001509 item = iternext(it);
1510 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001511 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001512 lz->cnt++;
1513 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001514 /* The (size_t) cast below avoids the danger of undefined
1515 behaviour from signed integer overflow. */
1516 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001517 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1518 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001519 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001520
1521empty:
1522 Py_CLEAR(lz->it);
1523 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001524}
1525
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001526static PyObject *
1527islice_reduce(isliceobject *lz)
1528{
1529 /* When unpickled, generate a new object with the same bounds,
1530 * then 'setstate' with the next and count
1531 */
1532 PyObject *stop;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001533
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001534 if (lz->it == NULL) {
1535 PyObject *empty_list;
1536 PyObject *empty_it;
1537 empty_list = PyList_New(0);
1538 if (empty_list == NULL)
1539 return NULL;
1540 empty_it = PyObject_GetIter(empty_list);
1541 Py_DECREF(empty_list);
1542 if (empty_it == NULL)
1543 return NULL;
1544 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1545 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001546 if (lz->stop == -1) {
1547 stop = Py_None;
1548 Py_INCREF(stop);
1549 } else {
1550 stop = PyLong_FromSsize_t(lz->stop);
1551 if (stop == NULL)
1552 return NULL;
1553 }
1554 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1555 lz->it, lz->next, stop, lz->step,
1556 lz->cnt);
1557}
1558
1559static PyObject *
1560islice_setstate(isliceobject *lz, PyObject *state)
1561{
1562 Py_ssize_t cnt = PyLong_AsSsize_t(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001563
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001564 if (cnt == -1 && PyErr_Occurred())
1565 return NULL;
1566 lz->cnt = cnt;
1567 Py_RETURN_NONE;
1568}
1569
1570static PyMethodDef islice_methods[] = {
1571 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1572 reduce_doc},
1573 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1574 setstate_doc},
1575 {NULL, NULL} /* sentinel */
1576};
1577
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001578PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001579"islice(iterable, stop) --> islice object\n\
1580islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001581\n\
1582Return an iterator whose next() method returns selected values from an\n\
1583iterable. If start is specified, will skip all preceding elements;\n\
1584otherwise, start defaults to zero. Step defaults to one. If\n\
1585specified as another value, step determines how many values are \n\
1586skipped between successive calls. Works like a slice() on a list\n\
1587but returns an iterator.");
1588
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001589static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001590 PyVarObject_HEAD_INIT(NULL, 0)
1591 "itertools.islice", /* tp_name */
1592 sizeof(isliceobject), /* tp_basicsize */
1593 0, /* tp_itemsize */
1594 /* methods */
1595 (destructor)islice_dealloc, /* tp_dealloc */
1596 0, /* tp_print */
1597 0, /* tp_getattr */
1598 0, /* tp_setattr */
1599 0, /* tp_reserved */
1600 0, /* tp_repr */
1601 0, /* tp_as_number */
1602 0, /* tp_as_sequence */
1603 0, /* tp_as_mapping */
1604 0, /* tp_hash */
1605 0, /* tp_call */
1606 0, /* tp_str */
1607 PyObject_GenericGetAttr, /* tp_getattro */
1608 0, /* tp_setattro */
1609 0, /* tp_as_buffer */
1610 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1611 Py_TPFLAGS_BASETYPE, /* tp_flags */
1612 islice_doc, /* tp_doc */
1613 (traverseproc)islice_traverse, /* tp_traverse */
1614 0, /* tp_clear */
1615 0, /* tp_richcompare */
1616 0, /* tp_weaklistoffset */
1617 PyObject_SelfIter, /* tp_iter */
1618 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001619 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001620 0, /* tp_members */
1621 0, /* tp_getset */
1622 0, /* tp_base */
1623 0, /* tp_dict */
1624 0, /* tp_descr_get */
1625 0, /* tp_descr_set */
1626 0, /* tp_dictoffset */
1627 0, /* tp_init */
1628 0, /* tp_alloc */
1629 islice_new, /* tp_new */
1630 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001631};
1632
1633
1634/* starmap object ************************************************************/
1635
1636typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001637 PyObject_HEAD
1638 PyObject *func;
1639 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001640} starmapobject;
1641
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001642static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001643
1644static PyObject *
1645starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1646{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001647 PyObject *func, *seq;
1648 PyObject *it;
1649 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001650
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001651 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1652 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001653
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001654 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1655 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001656
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001657 /* Get iterator. */
1658 it = PyObject_GetIter(seq);
1659 if (it == NULL)
1660 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001661
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001662 /* create starmapobject structure */
1663 lz = (starmapobject *)type->tp_alloc(type, 0);
1664 if (lz == NULL) {
1665 Py_DECREF(it);
1666 return NULL;
1667 }
1668 Py_INCREF(func);
1669 lz->func = func;
1670 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001671
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001672 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001673}
1674
1675static void
1676starmap_dealloc(starmapobject *lz)
1677{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001678 PyObject_GC_UnTrack(lz);
1679 Py_XDECREF(lz->func);
1680 Py_XDECREF(lz->it);
1681 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001682}
1683
1684static int
1685starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1686{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001687 Py_VISIT(lz->it);
1688 Py_VISIT(lz->func);
1689 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001690}
1691
1692static PyObject *
1693starmap_next(starmapobject *lz)
1694{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001695 PyObject *args;
1696 PyObject *result;
1697 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001698
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001699 args = (*Py_TYPE(it)->tp_iternext)(it);
1700 if (args == NULL)
1701 return NULL;
1702 if (!PyTuple_CheckExact(args)) {
1703 PyObject *newargs = PySequence_Tuple(args);
1704 Py_DECREF(args);
1705 if (newargs == NULL)
1706 return NULL;
1707 args = newargs;
1708 }
1709 result = PyObject_Call(lz->func, args, NULL);
1710 Py_DECREF(args);
1711 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001712}
1713
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001714static PyObject *
1715starmap_reduce(starmapobject *lz)
1716{
1717 /* Just pickle the iterator */
1718 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1719}
1720
1721static PyMethodDef starmap_methods[] = {
1722 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1723 reduce_doc},
1724 {NULL, NULL} /* sentinel */
1725};
1726
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001727PyDoc_STRVAR(starmap_doc,
1728"starmap(function, sequence) --> starmap object\n\
1729\n\
1730Return an iterator whose values are returned from the function evaluated\n\
Serhiy Storchakad65c9492015-11-02 14:10:23 +02001731with an argument tuple taken from the given sequence.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001732
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001733static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001734 PyVarObject_HEAD_INIT(NULL, 0)
1735 "itertools.starmap", /* tp_name */
1736 sizeof(starmapobject), /* tp_basicsize */
1737 0, /* tp_itemsize */
1738 /* methods */
1739 (destructor)starmap_dealloc, /* tp_dealloc */
1740 0, /* tp_print */
1741 0, /* tp_getattr */
1742 0, /* tp_setattr */
1743 0, /* tp_reserved */
1744 0, /* tp_repr */
1745 0, /* tp_as_number */
1746 0, /* tp_as_sequence */
1747 0, /* tp_as_mapping */
1748 0, /* tp_hash */
1749 0, /* tp_call */
1750 0, /* tp_str */
1751 PyObject_GenericGetAttr, /* tp_getattro */
1752 0, /* tp_setattro */
1753 0, /* tp_as_buffer */
1754 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1755 Py_TPFLAGS_BASETYPE, /* tp_flags */
1756 starmap_doc, /* tp_doc */
1757 (traverseproc)starmap_traverse, /* tp_traverse */
1758 0, /* tp_clear */
1759 0, /* tp_richcompare */
1760 0, /* tp_weaklistoffset */
1761 PyObject_SelfIter, /* tp_iter */
1762 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001763 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001764 0, /* tp_members */
1765 0, /* tp_getset */
1766 0, /* tp_base */
1767 0, /* tp_dict */
1768 0, /* tp_descr_get */
1769 0, /* tp_descr_set */
1770 0, /* tp_dictoffset */
1771 0, /* tp_init */
1772 0, /* tp_alloc */
1773 starmap_new, /* tp_new */
1774 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001775};
1776
1777
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001778/* chain object **************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001779
1780typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001781 PyObject_HEAD
1782 PyObject *source; /* Iterator over input iterables */
1783 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001784} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001785
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001786static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001787
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001788static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001789chain_new_internal(PyTypeObject *type, PyObject *source)
1790{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001791 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001792
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001793 lz = (chainobject *)type->tp_alloc(type, 0);
1794 if (lz == NULL) {
1795 Py_DECREF(source);
1796 return NULL;
1797 }
1798
1799 lz->source = source;
1800 lz->active = NULL;
1801 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001802}
1803
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001804static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001805chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001806{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001807 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001808
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001809 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1810 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001811
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001812 source = PyObject_GetIter(args);
1813 if (source == NULL)
1814 return NULL;
1815
1816 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001817}
1818
1819static PyObject *
1820chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1821{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001822 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001823
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001824 source = PyObject_GetIter(arg);
1825 if (source == NULL)
1826 return NULL;
1827
1828 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001829}
1830
1831static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001832chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001833{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001834 PyObject_GC_UnTrack(lz);
1835 Py_XDECREF(lz->active);
1836 Py_XDECREF(lz->source);
1837 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001838}
1839
Raymond Hettinger2012f172003-02-07 05:32:58 +00001840static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001841chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001842{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001843 Py_VISIT(lz->source);
1844 Py_VISIT(lz->active);
1845 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001846}
1847
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001848static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001849chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001850{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001851 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001852
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001853 if (lz->source == NULL)
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001854 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001855
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001856 if (lz->active == NULL) {
1857 PyObject *iterable = PyIter_Next(lz->source);
1858 if (iterable == NULL) {
1859 Py_CLEAR(lz->source);
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001860 return NULL; /* no more input sources */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001861 }
1862 lz->active = PyObject_GetIter(iterable);
1863 Py_DECREF(iterable);
1864 if (lz->active == NULL) {
1865 Py_CLEAR(lz->source);
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001866 return NULL; /* input not iterable */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001867 }
1868 }
Raymond Hettingerb5244a32015-08-16 14:24:20 -07001869 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001870 if (item != NULL)
1871 return item;
1872 if (PyErr_Occurred()) {
1873 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1874 PyErr_Clear();
1875 else
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001876 return NULL; /* input raised an exception */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001877 }
1878 Py_CLEAR(lz->active);
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001879 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001880}
1881
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001882static PyObject *
1883chain_reduce(chainobject *lz)
1884{
1885 if (lz->source) {
1886 /* we can't pickle function objects (itertools.from_iterable) so
1887 * we must use setstate to replace the iterable. One day we
1888 * will fix pickling of functions
1889 */
1890 if (lz->active) {
1891 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1892 } else {
1893 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1894 }
1895 } else {
1896 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1897 }
1898 return NULL;
1899}
1900
1901static PyObject *
1902chain_setstate(chainobject *lz, PyObject *state)
1903{
1904 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001905
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001906 if (! PyArg_ParseTuple(state, "O|O", &source, &active))
1907 return NULL;
1908
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02001909 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001910 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02001911 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001912 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001913 Py_RETURN_NONE;
1914}
1915
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001916PyDoc_STRVAR(chain_doc,
1917"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001918\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001919Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001920first iterable until it is exhausted, then elements from the next\n\
1921iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001922
Christian Heimesf16baeb2008-02-29 14:57:44 +00001923PyDoc_STRVAR(chain_from_iterable_doc,
1924"chain.from_iterable(iterable) --> chain object\n\
1925\n\
Martin Pantereb995702016-07-28 01:11:04 +00001926Alternate chain() constructor taking a single iterable argument\n\
Christian Heimesf16baeb2008-02-29 14:57:44 +00001927that evaluates lazily.");
1928
1929static PyMethodDef chain_methods[] = {
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001930 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1931 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001932 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1933 reduce_doc},
1934 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1935 setstate_doc},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001936 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001937};
1938
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001939static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001940 PyVarObject_HEAD_INIT(NULL, 0)
1941 "itertools.chain", /* tp_name */
1942 sizeof(chainobject), /* tp_basicsize */
1943 0, /* tp_itemsize */
1944 /* methods */
1945 (destructor)chain_dealloc, /* tp_dealloc */
1946 0, /* tp_print */
1947 0, /* tp_getattr */
1948 0, /* tp_setattr */
1949 0, /* tp_reserved */
1950 0, /* tp_repr */
1951 0, /* tp_as_number */
1952 0, /* tp_as_sequence */
1953 0, /* tp_as_mapping */
1954 0, /* tp_hash */
1955 0, /* tp_call */
1956 0, /* tp_str */
1957 PyObject_GenericGetAttr, /* tp_getattro */
1958 0, /* tp_setattro */
1959 0, /* tp_as_buffer */
1960 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1961 Py_TPFLAGS_BASETYPE, /* tp_flags */
1962 chain_doc, /* tp_doc */
1963 (traverseproc)chain_traverse, /* tp_traverse */
1964 0, /* tp_clear */
1965 0, /* tp_richcompare */
1966 0, /* tp_weaklistoffset */
1967 PyObject_SelfIter, /* tp_iter */
1968 (iternextfunc)chain_next, /* tp_iternext */
1969 chain_methods, /* tp_methods */
1970 0, /* tp_members */
1971 0, /* tp_getset */
1972 0, /* tp_base */
1973 0, /* tp_dict */
1974 0, /* tp_descr_get */
1975 0, /* tp_descr_set */
1976 0, /* tp_dictoffset */
1977 0, /* tp_init */
1978 0, /* tp_alloc */
1979 chain_new, /* tp_new */
1980 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001981};
1982
1983
Christian Heimesc3f30c42008-02-22 16:37:40 +00001984/* product object ************************************************************/
1985
1986typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001987 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001988 PyObject *pools; /* tuple of pool tuples */
1989 Py_ssize_t *indices; /* one index per pool */
1990 PyObject *result; /* most recently returned result tuple */
1991 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001992} productobject;
1993
1994static PyTypeObject product_type;
1995
1996static PyObject *
1997product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1998{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001999 productobject *lz;
2000 Py_ssize_t nargs, npools, repeat=1;
2001 PyObject *pools = NULL;
2002 Py_ssize_t *indices = NULL;
2003 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002005 if (kwds != NULL) {
2006 char *kwlist[] = {"repeat", 0};
2007 PyObject *tmpargs = PyTuple_New(0);
2008 if (tmpargs == NULL)
2009 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002010 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2011 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002012 Py_DECREF(tmpargs);
2013 return NULL;
2014 }
2015 Py_DECREF(tmpargs);
2016 if (repeat < 0) {
2017 PyErr_SetString(PyExc_ValueError,
2018 "repeat argument cannot be negative");
2019 return NULL;
2020 }
2021 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002022
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002023 assert(PyTuple_CheckExact(args));
2024 if (repeat == 0) {
2025 nargs = 0;
2026 } else {
2027 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002028 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002029 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2030 return NULL;
2031 }
2032 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002033 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002034
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002035 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002036 if (indices == NULL) {
2037 PyErr_NoMemory();
2038 goto error;
2039 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002040
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002041 pools = PyTuple_New(npools);
2042 if (pools == NULL)
2043 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002044
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002045 for (i=0; i < nargs ; ++i) {
2046 PyObject *item = PyTuple_GET_ITEM(args, i);
2047 PyObject *pool = PySequence_Tuple(item);
2048 if (pool == NULL)
2049 goto error;
2050 PyTuple_SET_ITEM(pools, i, pool);
2051 indices[i] = 0;
2052 }
2053 for ( ; i < npools; ++i) {
2054 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2055 Py_INCREF(pool);
2056 PyTuple_SET_ITEM(pools, i, pool);
2057 indices[i] = 0;
2058 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002059
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002060 /* create productobject structure */
2061 lz = (productobject *)type->tp_alloc(type, 0);
2062 if (lz == NULL)
2063 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002064
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002065 lz->pools = pools;
2066 lz->indices = indices;
2067 lz->result = NULL;
2068 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002069
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002070 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002071
2072error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002073 if (indices != NULL)
2074 PyMem_Free(indices);
2075 Py_XDECREF(pools);
2076 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002077}
2078
2079static void
2080product_dealloc(productobject *lz)
2081{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002082 PyObject_GC_UnTrack(lz);
2083 Py_XDECREF(lz->pools);
2084 Py_XDECREF(lz->result);
2085 if (lz->indices != NULL)
2086 PyMem_Free(lz->indices);
2087 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002088}
2089
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002090static PyObject *
2091product_sizeof(productobject *lz, void *unused)
2092{
2093 Py_ssize_t res;
2094
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002095 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002096 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2097 return PyLong_FromSsize_t(res);
2098}
2099
2100PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2101
Christian Heimesc3f30c42008-02-22 16:37:40 +00002102static int
2103product_traverse(productobject *lz, visitproc visit, void *arg)
2104{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002105 Py_VISIT(lz->pools);
2106 Py_VISIT(lz->result);
2107 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002108}
2109
2110static PyObject *
2111product_next(productobject *lz)
2112{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002113 PyObject *pool;
2114 PyObject *elem;
2115 PyObject *oldelem;
2116 PyObject *pools = lz->pools;
2117 PyObject *result = lz->result;
2118 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2119 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002120
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002121 if (lz->stopped)
2122 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002123
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002124 if (result == NULL) {
2125 /* On the first pass, return an initial tuple filled with the
2126 first element from each pool. */
2127 result = PyTuple_New(npools);
2128 if (result == NULL)
2129 goto empty;
2130 lz->result = result;
2131 for (i=0; i < npools; i++) {
2132 pool = PyTuple_GET_ITEM(pools, i);
2133 if (PyTuple_GET_SIZE(pool) == 0)
2134 goto empty;
2135 elem = PyTuple_GET_ITEM(pool, 0);
2136 Py_INCREF(elem);
2137 PyTuple_SET_ITEM(result, i, elem);
2138 }
2139 } else {
2140 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002141
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002142 /* Copy the previous result tuple or re-use it if available */
2143 if (Py_REFCNT(result) > 1) {
2144 PyObject *old_result = result;
2145 result = PyTuple_New(npools);
2146 if (result == NULL)
2147 goto empty;
2148 lz->result = result;
2149 for (i=0; i < npools; i++) {
2150 elem = PyTuple_GET_ITEM(old_result, i);
2151 Py_INCREF(elem);
2152 PyTuple_SET_ITEM(result, i, elem);
2153 }
2154 Py_DECREF(old_result);
2155 }
2156 /* Now, we've got the only copy so we can update it in-place */
2157 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002158
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002159 /* Update the pool indices right-to-left. Only advance to the
2160 next pool when the previous one rolls-over */
2161 for (i=npools-1 ; i >= 0 ; i--) {
2162 pool = PyTuple_GET_ITEM(pools, i);
2163 indices[i]++;
2164 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2165 /* Roll-over and advance to next pool */
2166 indices[i] = 0;
2167 elem = PyTuple_GET_ITEM(pool, 0);
2168 Py_INCREF(elem);
2169 oldelem = PyTuple_GET_ITEM(result, i);
2170 PyTuple_SET_ITEM(result, i, elem);
2171 Py_DECREF(oldelem);
2172 } else {
2173 /* No rollover. Just increment and stop here. */
2174 elem = PyTuple_GET_ITEM(pool, indices[i]);
2175 Py_INCREF(elem);
2176 oldelem = PyTuple_GET_ITEM(result, i);
2177 PyTuple_SET_ITEM(result, i, elem);
2178 Py_DECREF(oldelem);
2179 break;
2180 }
2181 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002182
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002183 /* If i is negative, then the indices have all rolled-over
2184 and we're done. */
2185 if (i < 0)
2186 goto empty;
2187 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002188
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002189 Py_INCREF(result);
2190 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002191
2192empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002193 lz->stopped = 1;
2194 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002195}
2196
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002197static PyObject *
2198product_reduce(productobject *lz)
2199{
2200 if (lz->stopped) {
2201 return Py_BuildValue("O(())", Py_TYPE(lz));
2202 } else if (lz->result == NULL) {
2203 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2204 } else {
2205 PyObject *indices;
2206 Py_ssize_t n, i;
2207
2208 /* we must pickle the indices use them for setstate, and
2209 * additionally indicate that the iterator has started
2210 */
2211 n = PyTuple_GET_SIZE(lz->pools);
2212 indices = PyTuple_New(n);
2213 if (indices == NULL)
2214 return NULL;
2215 for (i=0; i<n; i++){
2216 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2217 if (!index) {
2218 Py_DECREF(indices);
2219 return NULL;
2220 }
2221 PyTuple_SET_ITEM(indices, i, index);
2222 }
2223 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2224 }
2225}
2226
2227static PyObject *
2228product_setstate(productobject *lz, PyObject *state)
2229{
2230 PyObject *result;
2231 Py_ssize_t n, i;
2232
2233 n = PyTuple_GET_SIZE(lz->pools);
2234 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2235 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2236 return NULL;
2237 }
2238 for (i=0; i<n; i++)
2239 {
2240 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2241 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002242 PyObject* pool;
2243 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002244 if (index < 0 && PyErr_Occurred())
2245 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002246 pool = PyTuple_GET_ITEM(lz->pools, i);
2247 poolsize = PyTuple_GET_SIZE(pool);
2248 if (poolsize == 0) {
2249 lz->stopped = 1;
2250 Py_RETURN_NONE;
2251 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002252 /* clamp the index */
2253 if (index < 0)
2254 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002255 else if (index > poolsize-1)
2256 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002257 lz->indices[i] = index;
2258 }
2259
2260 result = PyTuple_New(n);
2261 if (!result)
2262 return NULL;
2263 for (i=0; i<n; i++) {
2264 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2265 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2266 Py_INCREF(element);
2267 PyTuple_SET_ITEM(result, i, element);
2268 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002269 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002270 Py_RETURN_NONE;
2271}
2272
2273static PyMethodDef product_methods[] = {
2274 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2275 reduce_doc},
2276 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2277 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002278 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2279 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002280 {NULL, NULL} /* sentinel */
2281};
2282
Christian Heimesc3f30c42008-02-22 16:37:40 +00002283PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002284"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002285\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002286Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002287For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2288The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2289cycle in a manner similar to an odometer (with the rightmost element changing\n\
2290on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002291To compute the product of an iterable with itself, specify the number\n\
2292of repetitions with the optional repeat keyword argument. For example,\n\
2293product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002294product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2295product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2296
2297static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002298 PyVarObject_HEAD_INIT(NULL, 0)
2299 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002300 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002301 0, /* tp_itemsize */
2302 /* methods */
2303 (destructor)product_dealloc, /* tp_dealloc */
2304 0, /* tp_print */
2305 0, /* tp_getattr */
2306 0, /* tp_setattr */
2307 0, /* tp_reserved */
2308 0, /* tp_repr */
2309 0, /* tp_as_number */
2310 0, /* tp_as_sequence */
2311 0, /* tp_as_mapping */
2312 0, /* tp_hash */
2313 0, /* tp_call */
2314 0, /* tp_str */
2315 PyObject_GenericGetAttr, /* tp_getattro */
2316 0, /* tp_setattro */
2317 0, /* tp_as_buffer */
2318 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2319 Py_TPFLAGS_BASETYPE, /* tp_flags */
2320 product_doc, /* tp_doc */
2321 (traverseproc)product_traverse, /* tp_traverse */
2322 0, /* tp_clear */
2323 0, /* tp_richcompare */
2324 0, /* tp_weaklistoffset */
2325 PyObject_SelfIter, /* tp_iter */
2326 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002327 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002328 0, /* tp_members */
2329 0, /* tp_getset */
2330 0, /* tp_base */
2331 0, /* tp_dict */
2332 0, /* tp_descr_get */
2333 0, /* tp_descr_set */
2334 0, /* tp_dictoffset */
2335 0, /* tp_init */
2336 0, /* tp_alloc */
2337 product_new, /* tp_new */
2338 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002339};
2340
2341
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002342/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002343
2344typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002345 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002346 PyObject *pool; /* input converted to a tuple */
2347 Py_ssize_t *indices; /* one index per result element */
2348 PyObject *result; /* most recently returned result tuple */
2349 Py_ssize_t r; /* size of result tuple */
2350 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002351} combinationsobject;
2352
2353static PyTypeObject combinations_type;
2354
2355static PyObject *
2356combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2357{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002358 combinationsobject *co;
2359 Py_ssize_t n;
2360 Py_ssize_t r;
2361 PyObject *pool = NULL;
2362 PyObject *iterable = NULL;
2363 Py_ssize_t *indices = NULL;
2364 Py_ssize_t i;
2365 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002366
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002367 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2368 &iterable, &r))
2369 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002370
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002371 pool = PySequence_Tuple(iterable);
2372 if (pool == NULL)
2373 goto error;
2374 n = PyTuple_GET_SIZE(pool);
2375 if (r < 0) {
2376 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2377 goto error;
2378 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002379
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002380 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002381 if (indices == NULL) {
2382 PyErr_NoMemory();
2383 goto error;
2384 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002385
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002386 for (i=0 ; i<r ; i++)
2387 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002388
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002389 /* create combinationsobject structure */
2390 co = (combinationsobject *)type->tp_alloc(type, 0);
2391 if (co == NULL)
2392 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002393
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002394 co->pool = pool;
2395 co->indices = indices;
2396 co->result = NULL;
2397 co->r = r;
2398 co->stopped = r > n ? 1 : 0;
2399
2400 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002401
2402error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002403 if (indices != NULL)
2404 PyMem_Free(indices);
2405 Py_XDECREF(pool);
2406 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002407}
2408
2409static void
2410combinations_dealloc(combinationsobject *co)
2411{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002412 PyObject_GC_UnTrack(co);
2413 Py_XDECREF(co->pool);
2414 Py_XDECREF(co->result);
2415 if (co->indices != NULL)
2416 PyMem_Free(co->indices);
2417 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002418}
2419
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002420static PyObject *
2421combinations_sizeof(combinationsobject *co, void *unused)
2422{
2423 Py_ssize_t res;
2424
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002425 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002426 res += co->r * sizeof(Py_ssize_t);
2427 return PyLong_FromSsize_t(res);
2428}
2429
Christian Heimes380f7f22008-02-28 11:19:05 +00002430static int
2431combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2432{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002433 Py_VISIT(co->pool);
2434 Py_VISIT(co->result);
2435 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002436}
2437
2438static PyObject *
2439combinations_next(combinationsobject *co)
2440{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002441 PyObject *elem;
2442 PyObject *oldelem;
2443 PyObject *pool = co->pool;
2444 Py_ssize_t *indices = co->indices;
2445 PyObject *result = co->result;
2446 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2447 Py_ssize_t r = co->r;
2448 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002449
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002450 if (co->stopped)
2451 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002452
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002453 if (result == NULL) {
2454 /* On the first pass, initialize result tuple using the indices */
2455 result = PyTuple_New(r);
2456 if (result == NULL)
2457 goto empty;
2458 co->result = result;
2459 for (i=0; i<r ; i++) {
2460 index = indices[i];
2461 elem = PyTuple_GET_ITEM(pool, index);
2462 Py_INCREF(elem);
2463 PyTuple_SET_ITEM(result, i, elem);
2464 }
2465 } else {
2466 /* Copy the previous result tuple or re-use it if available */
2467 if (Py_REFCNT(result) > 1) {
2468 PyObject *old_result = result;
2469 result = PyTuple_New(r);
2470 if (result == NULL)
2471 goto empty;
2472 co->result = result;
2473 for (i=0; i<r ; i++) {
2474 elem = PyTuple_GET_ITEM(old_result, i);
2475 Py_INCREF(elem);
2476 PyTuple_SET_ITEM(result, i, elem);
2477 }
2478 Py_DECREF(old_result);
2479 }
2480 /* Now, we've got the only copy so we can update it in-place
2481 * CPython's empty tuple is a singleton and cached in
2482 * PyTuple's freelist.
2483 */
2484 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002485
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002486 /* Scan indices right-to-left until finding one that is not
2487 at its maximum (i + n - r). */
2488 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2489 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002490
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002491 /* If i is negative, then the indices are all at
2492 their maximum value and we're done. */
2493 if (i < 0)
2494 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002495
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002496 /* Increment the current index which we know is not at its
2497 maximum. Then move back to the right setting each index
2498 to its lowest possible value (one higher than the index
2499 to its left -- this maintains the sort order invariant). */
2500 indices[i]++;
2501 for (j=i+1 ; j<r ; j++)
2502 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002503
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002504 /* Update the result tuple for the new indices
2505 starting with i, the leftmost index that changed */
2506 for ( ; i<r ; i++) {
2507 index = indices[i];
2508 elem = PyTuple_GET_ITEM(pool, index);
2509 Py_INCREF(elem);
2510 oldelem = PyTuple_GET_ITEM(result, i);
2511 PyTuple_SET_ITEM(result, i, elem);
2512 Py_DECREF(oldelem);
2513 }
2514 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002515
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002516 Py_INCREF(result);
2517 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002518
2519empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002520 co->stopped = 1;
2521 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002522}
2523
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002524static PyObject *
2525combinations_reduce(combinationsobject *lz)
2526{
2527 if (lz->result == NULL) {
2528 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2529 } else if (lz->stopped) {
2530 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2531 } else {
2532 PyObject *indices;
2533 Py_ssize_t i;
2534
2535 /* we must pickle the indices and use them for setstate */
2536 indices = PyTuple_New(lz->r);
2537 if (!indices)
2538 return NULL;
2539 for (i=0; i<lz->r; i++)
2540 {
2541 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2542 if (!index) {
2543 Py_DECREF(indices);
2544 return NULL;
2545 }
2546 PyTuple_SET_ITEM(indices, i, index);
2547 }
2548
2549 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2550 }
2551}
2552
2553static PyObject *
2554combinations_setstate(combinationsobject *lz, PyObject *state)
2555{
2556 PyObject *result;
2557 Py_ssize_t i;
2558 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2559
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002560 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002561 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2562 return NULL;
2563 }
2564
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002565 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002566 Py_ssize_t max;
2567 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2568 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002569
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002570 if (index == -1 && PyErr_Occurred())
2571 return NULL; /* not an integer */
2572 max = i + n - lz->r;
2573 /* clamp the index (beware of negative max) */
2574 if (index > max)
2575 index = max;
2576 if (index < 0)
2577 index = 0;
2578 lz->indices[i] = index;
2579 }
2580
2581 result = PyTuple_New(lz->r);
2582 if (result == NULL)
2583 return NULL;
2584 for (i=0; i<lz->r; i++) {
2585 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2586 Py_INCREF(element);
2587 PyTuple_SET_ITEM(result, i, element);
2588 }
2589
Serhiy Storchaka48842712016-04-06 09:45:48 +03002590 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002591 Py_RETURN_NONE;
2592}
2593
2594static PyMethodDef combinations_methods[] = {
2595 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2596 reduce_doc},
2597 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2598 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002599 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2600 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002601 {NULL, NULL} /* sentinel */
2602};
2603
Christian Heimes380f7f22008-02-28 11:19:05 +00002604PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002605"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002606\n\
2607Return successive r-length combinations of elements in the iterable.\n\n\
2608combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2609
2610static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002611 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002612 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002613 sizeof(combinationsobject), /* tp_basicsize */
2614 0, /* tp_itemsize */
2615 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002616 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002617 0, /* tp_print */
2618 0, /* tp_getattr */
2619 0, /* tp_setattr */
2620 0, /* tp_reserved */
2621 0, /* tp_repr */
2622 0, /* tp_as_number */
2623 0, /* tp_as_sequence */
2624 0, /* tp_as_mapping */
2625 0, /* tp_hash */
2626 0, /* tp_call */
2627 0, /* tp_str */
2628 PyObject_GenericGetAttr, /* tp_getattro */
2629 0, /* tp_setattro */
2630 0, /* tp_as_buffer */
2631 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2632 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002633 combinations_doc, /* tp_doc */
2634 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002635 0, /* tp_clear */
2636 0, /* tp_richcompare */
2637 0, /* tp_weaklistoffset */
2638 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002639 (iternextfunc)combinations_next, /* tp_iternext */
2640 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002641 0, /* tp_members */
2642 0, /* tp_getset */
2643 0, /* tp_base */
2644 0, /* tp_dict */
2645 0, /* tp_descr_get */
2646 0, /* tp_descr_set */
2647 0, /* tp_dictoffset */
2648 0, /* tp_init */
2649 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002650 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002651 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002652};
2653
2654
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002655/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002656
2657/* Equivalent to:
2658
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002659 def combinations_with_replacement(iterable, r):
2660 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2661 # number items returned: (n+r-1)! / r! / (n-1)!
2662 pool = tuple(iterable)
2663 n = len(pool)
2664 indices = [0] * r
2665 yield tuple(pool[i] for i in indices)
2666 while 1:
2667 for i in reversed(range(r)):
2668 if indices[i] != n - 1:
2669 break
2670 else:
2671 return
2672 indices[i:] = [indices[i] + 1] * (r - i)
2673 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002674
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002675 def combinations_with_replacement2(iterable, r):
2676 'Alternate version that filters from product()'
2677 pool = tuple(iterable)
2678 n = len(pool)
2679 for indices in product(range(n), repeat=r):
2680 if sorted(indices) == list(indices):
2681 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002682*/
2683typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002684 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002685 PyObject *pool; /* input converted to a tuple */
2686 Py_ssize_t *indices; /* one index per result element */
2687 PyObject *result; /* most recently returned result tuple */
2688 Py_ssize_t r; /* size of result tuple */
2689 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002690} cwrobject;
2691
2692static PyTypeObject cwr_type;
2693
2694static PyObject *
2695cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2696{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002697 cwrobject *co;
2698 Py_ssize_t n;
2699 Py_ssize_t r;
2700 PyObject *pool = NULL;
2701 PyObject *iterable = NULL;
2702 Py_ssize_t *indices = NULL;
2703 Py_ssize_t i;
2704 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002705
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002706 if (!PyArg_ParseTupleAndKeywords(args, kwds,
2707 "On:combinations_with_replacement",
2708 kwargs, &iterable, &r))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002709 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002710
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002711 pool = PySequence_Tuple(iterable);
2712 if (pool == NULL)
2713 goto error;
2714 n = PyTuple_GET_SIZE(pool);
2715 if (r < 0) {
2716 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2717 goto error;
2718 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002719
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002720 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002721 if (indices == NULL) {
2722 PyErr_NoMemory();
2723 goto error;
2724 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002725
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002726 for (i=0 ; i<r ; i++)
2727 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002728
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002729 /* create cwrobject structure */
2730 co = (cwrobject *)type->tp_alloc(type, 0);
2731 if (co == NULL)
2732 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002733
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002734 co->pool = pool;
2735 co->indices = indices;
2736 co->result = NULL;
2737 co->r = r;
2738 co->stopped = !n && r;
2739
2740 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002741
2742error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002743 if (indices != NULL)
2744 PyMem_Free(indices);
2745 Py_XDECREF(pool);
2746 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002747}
2748
2749static void
2750cwr_dealloc(cwrobject *co)
2751{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002752 PyObject_GC_UnTrack(co);
2753 Py_XDECREF(co->pool);
2754 Py_XDECREF(co->result);
2755 if (co->indices != NULL)
2756 PyMem_Free(co->indices);
2757 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002758}
2759
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002760static PyObject *
2761cwr_sizeof(cwrobject *co, void *unused)
2762{
2763 Py_ssize_t res;
2764
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002765 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002766 res += co->r * sizeof(Py_ssize_t);
2767 return PyLong_FromSsize_t(res);
2768}
2769
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002770static int
2771cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2772{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002773 Py_VISIT(co->pool);
2774 Py_VISIT(co->result);
2775 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002776}
2777
2778static PyObject *
2779cwr_next(cwrobject *co)
2780{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002781 PyObject *elem;
2782 PyObject *oldelem;
2783 PyObject *pool = co->pool;
2784 Py_ssize_t *indices = co->indices;
2785 PyObject *result = co->result;
2786 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2787 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002788 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002789
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002790 if (co->stopped)
2791 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002792
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002793 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002794 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002795 result = PyTuple_New(r);
2796 if (result == NULL)
2797 goto empty;
2798 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002799 if (n > 0) {
2800 elem = PyTuple_GET_ITEM(pool, 0);
2801 for (i=0; i<r ; i++) {
2802 assert(indices[i] == 0);
2803 Py_INCREF(elem);
2804 PyTuple_SET_ITEM(result, i, elem);
2805 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002806 }
2807 } else {
2808 /* Copy the previous result tuple or re-use it if available */
2809 if (Py_REFCNT(result) > 1) {
2810 PyObject *old_result = result;
2811 result = PyTuple_New(r);
2812 if (result == NULL)
2813 goto empty;
2814 co->result = result;
2815 for (i=0; i<r ; i++) {
2816 elem = PyTuple_GET_ITEM(old_result, i);
2817 Py_INCREF(elem);
2818 PyTuple_SET_ITEM(result, i, elem);
2819 }
2820 Py_DECREF(old_result);
2821 }
2822 /* Now, we've got the only copy so we can update it in-place CPython's
2823 empty tuple is a singleton and cached in PyTuple's freelist. */
2824 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002825
Tim Peters9edb1682013-09-03 11:49:31 -05002826 /* Scan indices right-to-left until finding one that is not
2827 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002828 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2829 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002830
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002831 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002832 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002833 if (i < 0)
2834 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002835
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002836 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002837 maximum. Then set all to the right to the same value. */
2838 index = indices[i] + 1;
2839 assert(index < n);
2840 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002841 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002842 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002843 Py_INCREF(elem);
2844 oldelem = PyTuple_GET_ITEM(result, i);
2845 PyTuple_SET_ITEM(result, i, elem);
2846 Py_DECREF(oldelem);
2847 }
2848 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002849
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002850 Py_INCREF(result);
2851 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002852
2853empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002854 co->stopped = 1;
2855 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002856}
2857
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002858static PyObject *
2859cwr_reduce(cwrobject *lz)
2860{
2861 if (lz->result == NULL) {
2862 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2863 } else if (lz->stopped) {
2864 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2865 } else {
2866 PyObject *indices;
2867 Py_ssize_t i;
2868
2869 /* we must pickle the indices and use them for setstate */
2870 indices = PyTuple_New(lz->r);
2871 if (!indices)
2872 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002873 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002874 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2875 if (!index) {
2876 Py_DECREF(indices);
2877 return NULL;
2878 }
2879 PyTuple_SET_ITEM(indices, i, index);
2880 }
2881
2882 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2883 }
2884}
2885
2886static PyObject *
2887cwr_setstate(cwrobject *lz, PyObject *state)
2888{
2889 PyObject *result;
2890 Py_ssize_t n, i;
2891
2892 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2893 {
2894 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2895 return NULL;
2896 }
2897
2898 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002899 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002900 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2901 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002902
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002903 if (index < 0 && PyErr_Occurred())
2904 return NULL; /* not an integer */
2905 /* clamp the index */
2906 if (index < 0)
2907 index = 0;
2908 else if (index > n-1)
2909 index = n-1;
2910 lz->indices[i] = index;
2911 }
2912 result = PyTuple_New(lz->r);
2913 if (result == NULL)
2914 return NULL;
2915 for (i=0; i<lz->r; i++) {
2916 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2917 Py_INCREF(element);
2918 PyTuple_SET_ITEM(result, i, element);
2919 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002920 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002921 Py_RETURN_NONE;
2922}
2923
2924static PyMethodDef cwr_methods[] = {
2925 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2926 reduce_doc},
2927 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2928 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002929 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
2930 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002931 {NULL, NULL} /* sentinel */
2932};
2933
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002934PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002935"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002936\n\
2937Return successive r-length combinations of elements in the iterable\n\
2938allowing individual elements to have successive repeats.\n\
2939combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2940
2941static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002942 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002943 "itertools.combinations_with_replacement", /* tp_name */
2944 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002945 0, /* tp_itemsize */
2946 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002947 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002948 0, /* tp_print */
2949 0, /* tp_getattr */
2950 0, /* tp_setattr */
2951 0, /* tp_reserved */
2952 0, /* tp_repr */
2953 0, /* tp_as_number */
2954 0, /* tp_as_sequence */
2955 0, /* tp_as_mapping */
2956 0, /* tp_hash */
2957 0, /* tp_call */
2958 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002959 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002960 0, /* tp_setattro */
2961 0, /* tp_as_buffer */
2962 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002963 Py_TPFLAGS_BASETYPE, /* tp_flags */
2964 cwr_doc, /* tp_doc */
2965 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002966 0, /* tp_clear */
2967 0, /* tp_richcompare */
2968 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002969 PyObject_SelfIter, /* tp_iter */
2970 (iternextfunc)cwr_next, /* tp_iternext */
2971 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002972 0, /* tp_members */
2973 0, /* tp_getset */
2974 0, /* tp_base */
2975 0, /* tp_dict */
2976 0, /* tp_descr_get */
2977 0, /* tp_descr_set */
2978 0, /* tp_dictoffset */
2979 0, /* tp_init */
2980 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002981 cwr_new, /* tp_new */
2982 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002983};
2984
2985
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002986/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002987
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002988def permutations(iterable, r=None):
2989 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2990 pool = tuple(iterable)
2991 n = len(pool)
2992 r = n if r is None else r
2993 indices = range(n)
2994 cycles = range(n-r+1, n+1)[::-1]
2995 yield tuple(pool[i] for i in indices[:r])
2996 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03002997 for i in reversed(range(r)):
2998 cycles[i] -= 1
2999 if cycles[i] == 0:
3000 indices[i:] = indices[i+1:] + indices[i:i+1]
3001 cycles[i] = n - i
3002 else:
3003 j = cycles[i]
3004 indices[i], indices[-j] = indices[-j], indices[i]
3005 yield tuple(pool[i] for i in indices[:r])
3006 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003007 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003008 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003009*/
3010
3011typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003012 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003013 PyObject *pool; /* input converted to a tuple */
3014 Py_ssize_t *indices; /* one index per element in the pool */
3015 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3016 PyObject *result; /* most recently returned result tuple */
3017 Py_ssize_t r; /* size of result tuple */
3018 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003019} permutationsobject;
3020
3021static PyTypeObject permutations_type;
3022
3023static PyObject *
3024permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3025{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003026 permutationsobject *po;
3027 Py_ssize_t n;
3028 Py_ssize_t r;
3029 PyObject *robj = Py_None;
3030 PyObject *pool = NULL;
3031 PyObject *iterable = NULL;
3032 Py_ssize_t *indices = NULL;
3033 Py_ssize_t *cycles = NULL;
3034 Py_ssize_t i;
3035 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003036
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003037 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
3038 &iterable, &robj))
3039 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003040
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003041 pool = PySequence_Tuple(iterable);
3042 if (pool == NULL)
3043 goto error;
3044 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003045
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003046 r = n;
3047 if (robj != Py_None) {
3048 if (!PyLong_Check(robj)) {
3049 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3050 goto error;
3051 }
3052 r = PyLong_AsSsize_t(robj);
3053 if (r == -1 && PyErr_Occurred())
3054 goto error;
3055 }
3056 if (r < 0) {
3057 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3058 goto error;
3059 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003060
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003061 indices = PyMem_New(Py_ssize_t, n);
3062 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003063 if (indices == NULL || cycles == NULL) {
3064 PyErr_NoMemory();
3065 goto error;
3066 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003067
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003068 for (i=0 ; i<n ; i++)
3069 indices[i] = i;
3070 for (i=0 ; i<r ; i++)
3071 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003072
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003073 /* create permutationsobject structure */
3074 po = (permutationsobject *)type->tp_alloc(type, 0);
3075 if (po == NULL)
3076 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003077
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003078 po->pool = pool;
3079 po->indices = indices;
3080 po->cycles = cycles;
3081 po->result = NULL;
3082 po->r = r;
3083 po->stopped = r > n ? 1 : 0;
3084
3085 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003086
3087error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003088 if (indices != NULL)
3089 PyMem_Free(indices);
3090 if (cycles != NULL)
3091 PyMem_Free(cycles);
3092 Py_XDECREF(pool);
3093 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003094}
3095
3096static void
3097permutations_dealloc(permutationsobject *po)
3098{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003099 PyObject_GC_UnTrack(po);
3100 Py_XDECREF(po->pool);
3101 Py_XDECREF(po->result);
3102 PyMem_Free(po->indices);
3103 PyMem_Free(po->cycles);
3104 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003105}
3106
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003107static PyObject *
3108permutations_sizeof(permutationsobject *po, void *unused)
3109{
3110 Py_ssize_t res;
3111
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003112 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003113 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3114 res += po->r * sizeof(Py_ssize_t);
3115 return PyLong_FromSsize_t(res);
3116}
3117
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003118static int
3119permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3120{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003121 Py_VISIT(po->pool);
3122 Py_VISIT(po->result);
3123 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003124}
3125
3126static PyObject *
3127permutations_next(permutationsobject *po)
3128{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003129 PyObject *elem;
3130 PyObject *oldelem;
3131 PyObject *pool = po->pool;
3132 Py_ssize_t *indices = po->indices;
3133 Py_ssize_t *cycles = po->cycles;
3134 PyObject *result = po->result;
3135 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3136 Py_ssize_t r = po->r;
3137 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003138
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003139 if (po->stopped)
3140 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003141
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003142 if (result == NULL) {
3143 /* On the first pass, initialize result tuple using the indices */
3144 result = PyTuple_New(r);
3145 if (result == NULL)
3146 goto empty;
3147 po->result = result;
3148 for (i=0; i<r ; i++) {
3149 index = indices[i];
3150 elem = PyTuple_GET_ITEM(pool, index);
3151 Py_INCREF(elem);
3152 PyTuple_SET_ITEM(result, i, elem);
3153 }
3154 } else {
3155 if (n == 0)
3156 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003157
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003158 /* Copy the previous result tuple or re-use it if available */
3159 if (Py_REFCNT(result) > 1) {
3160 PyObject *old_result = result;
3161 result = PyTuple_New(r);
3162 if (result == NULL)
3163 goto empty;
3164 po->result = result;
3165 for (i=0; i<r ; i++) {
3166 elem = PyTuple_GET_ITEM(old_result, i);
3167 Py_INCREF(elem);
3168 PyTuple_SET_ITEM(result, i, elem);
3169 }
3170 Py_DECREF(old_result);
3171 }
3172 /* Now, we've got the only copy so we can update it in-place */
3173 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003174
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003175 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3176 for (i=r-1 ; i>=0 ; i--) {
3177 cycles[i] -= 1;
3178 if (cycles[i] == 0) {
3179 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3180 index = indices[i];
3181 for (j=i ; j<n-1 ; j++)
3182 indices[j] = indices[j+1];
3183 indices[n-1] = index;
3184 cycles[i] = n - i;
3185 } else {
3186 j = cycles[i];
3187 index = indices[i];
3188 indices[i] = indices[n-j];
3189 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003190
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003191 for (k=i; k<r ; k++) {
3192 /* start with i, the leftmost element that changed */
3193 /* yield tuple(pool[k] for k in indices[:r]) */
3194 index = indices[k];
3195 elem = PyTuple_GET_ITEM(pool, index);
3196 Py_INCREF(elem);
3197 oldelem = PyTuple_GET_ITEM(result, k);
3198 PyTuple_SET_ITEM(result, k, elem);
3199 Py_DECREF(oldelem);
3200 }
3201 break;
3202 }
3203 }
3204 /* If i is negative, then the cycles have all
3205 rolled-over and we're done. */
3206 if (i < 0)
3207 goto empty;
3208 }
3209 Py_INCREF(result);
3210 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003211
3212empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003213 po->stopped = 1;
3214 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003215}
3216
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003217static PyObject *
3218permutations_reduce(permutationsobject *po)
3219{
3220 if (po->result == NULL) {
3221 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3222 } else if (po->stopped) {
3223 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3224 } else {
3225 PyObject *indices=NULL, *cycles=NULL;
3226 Py_ssize_t n, i;
3227
3228 /* we must pickle the indices and cycles and use them for setstate */
3229 n = PyTuple_GET_SIZE(po->pool);
3230 indices = PyTuple_New(n);
3231 if (indices == NULL)
3232 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003233 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003234 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3235 if (!index)
3236 goto err;
3237 PyTuple_SET_ITEM(indices, i, index);
3238 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003239
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003240 cycles = PyTuple_New(po->r);
3241 if (cycles == NULL)
3242 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003243 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003244 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3245 if (!index)
3246 goto err;
3247 PyTuple_SET_ITEM(cycles, i, index);
3248 }
3249 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3250 po->pool, po->r,
3251 indices, cycles);
3252 err:
3253 Py_XDECREF(indices);
3254 Py_XDECREF(cycles);
3255 return NULL;
3256 }
3257}
3258
3259static PyObject *
3260permutations_setstate(permutationsobject *po, PyObject *state)
3261{
3262 PyObject *indices, *cycles, *result;
3263 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003264
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003265 if (!PyArg_ParseTuple(state, "O!O!",
3266 &PyTuple_Type, &indices,
3267 &PyTuple_Type, &cycles))
3268 return NULL;
3269
3270 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003271 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003272 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3273 return NULL;
3274 }
3275
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003276 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003277 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3278 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3279 if (index < 0 && PyErr_Occurred())
3280 return NULL; /* not an integer */
3281 /* clamp the index */
3282 if (index < 0)
3283 index = 0;
3284 else if (index > n-1)
3285 index = n-1;
3286 po->indices[i] = index;
3287 }
3288
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003289 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003290 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3291 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3292 if (index < 0 && PyErr_Occurred())
3293 return NULL; /* not an integer */
3294 if (index < 1)
3295 index = 1;
3296 else if (index > n-i)
3297 index = n-i;
3298 po->cycles[i] = index;
3299 }
3300 result = PyTuple_New(po->r);
3301 if (result == NULL)
3302 return NULL;
3303 for (i=0; i<po->r; i++) {
3304 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3305 Py_INCREF(element);
3306 PyTuple_SET_ITEM(result, i, element);
3307 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003308 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003309 Py_RETURN_NONE;
3310}
3311
3312static PyMethodDef permuations_methods[] = {
3313 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3314 reduce_doc},
3315 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3316 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003317 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3318 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003319 {NULL, NULL} /* sentinel */
3320};
3321
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003322PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003323"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003324\n\
3325Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003326permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003327
3328static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003329 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003330 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003331 sizeof(permutationsobject), /* tp_basicsize */
3332 0, /* tp_itemsize */
3333 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003334 (destructor)permutations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003335 0, /* tp_print */
3336 0, /* tp_getattr */
3337 0, /* tp_setattr */
3338 0, /* tp_reserved */
3339 0, /* tp_repr */
3340 0, /* tp_as_number */
3341 0, /* tp_as_sequence */
3342 0, /* tp_as_mapping */
3343 0, /* tp_hash */
3344 0, /* tp_call */
3345 0, /* tp_str */
3346 PyObject_GenericGetAttr, /* tp_getattro */
3347 0, /* tp_setattro */
3348 0, /* tp_as_buffer */
3349 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3350 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003351 permutations_doc, /* tp_doc */
3352 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003353 0, /* tp_clear */
3354 0, /* tp_richcompare */
3355 0, /* tp_weaklistoffset */
3356 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003357 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003358 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003359 0, /* tp_members */
3360 0, /* tp_getset */
3361 0, /* tp_base */
3362 0, /* tp_dict */
3363 0, /* tp_descr_get */
3364 0, /* tp_descr_set */
3365 0, /* tp_dictoffset */
3366 0, /* tp_init */
3367 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003368 permutations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003369 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003370};
3371
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003372/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003373
3374typedef struct {
3375 PyObject_HEAD
3376 PyObject *total;
3377 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003378 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003379} accumulateobject;
3380
3381static PyTypeObject accumulate_type;
3382
3383static PyObject *
3384accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3385{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003386 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003387 PyObject *iterable;
3388 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003389 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003390 accumulateobject *lz;
3391
Raymond Hettinger5d446132011-03-27 18:52:10 -07003392 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3393 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003394 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003395
3396 /* Get iterator. */
3397 it = PyObject_GetIter(iterable);
3398 if (it == NULL)
3399 return NULL;
3400
Raymond Hettinger482ba772010-12-01 22:48:00 +00003401 /* create accumulateobject structure */
3402 lz = (accumulateobject *)type->tp_alloc(type, 0);
3403 if (lz == NULL) {
3404 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003405 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003406 }
3407
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003408 if (binop != Py_None) {
3409 Py_XINCREF(binop);
3410 lz->binop = binop;
3411 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003412 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003413 lz->it = it;
3414 return (PyObject *)lz;
3415}
3416
3417static void
3418accumulate_dealloc(accumulateobject *lz)
3419{
3420 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003421 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003422 Py_XDECREF(lz->total);
3423 Py_XDECREF(lz->it);
3424 Py_TYPE(lz)->tp_free(lz);
3425}
3426
3427static int
3428accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3429{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003430 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003431 Py_VISIT(lz->it);
3432 Py_VISIT(lz->total);
3433 return 0;
3434}
3435
3436static PyObject *
3437accumulate_next(accumulateobject *lz)
3438{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003439 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003440
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003441 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003442 if (val == NULL)
3443 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003444
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003445 if (lz->total == NULL) {
3446 Py_INCREF(val);
3447 lz->total = val;
3448 return lz->total;
3449 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003450
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003451 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003452 newtotal = PyNumber_Add(lz->total, val);
3453 else
3454 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003455 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003456 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003457 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003458
Raymond Hettinger482ba772010-12-01 22:48:00 +00003459 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003460 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003461 return newtotal;
3462}
3463
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003464static PyObject *
3465accumulate_reduce(accumulateobject *lz)
3466{
Serhiy Storchakad5516252016-03-06 14:00:45 +02003467 if (lz->total == Py_None) {
3468 PyObject *it;
3469
3470 if (PyType_Ready(&chain_type) < 0)
3471 return NULL;
3472 if (PyType_Ready(&islice_type) < 0)
3473 return NULL;
3474 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3475 lz->total, lz->it);
3476 if (it == NULL)
3477 return NULL;
3478 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3479 it, lz->binop ? lz->binop : Py_None);
3480 if (it == NULL)
3481 return NULL;
3482 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3483 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003484 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3485 lz->it, lz->binop?lz->binop:Py_None,
3486 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003487}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003488
3489static PyObject *
3490accumulate_setstate(accumulateobject *lz, PyObject *state)
3491{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003492 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003493 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003494 Py_RETURN_NONE;
3495}
3496
3497static PyMethodDef accumulate_methods[] = {
3498 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3499 reduce_doc},
3500 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3501 setstate_doc},
3502 {NULL, NULL} /* sentinel */
3503};
3504
Raymond Hettinger482ba772010-12-01 22:48:00 +00003505PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003506"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003507\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003508Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003509
3510static PyTypeObject accumulate_type = {
3511 PyVarObject_HEAD_INIT(NULL, 0)
3512 "itertools.accumulate", /* tp_name */
3513 sizeof(accumulateobject), /* tp_basicsize */
3514 0, /* tp_itemsize */
3515 /* methods */
3516 (destructor)accumulate_dealloc, /* tp_dealloc */
3517 0, /* tp_print */
3518 0, /* tp_getattr */
3519 0, /* tp_setattr */
3520 0, /* tp_reserved */
3521 0, /* tp_repr */
3522 0, /* tp_as_number */
3523 0, /* tp_as_sequence */
3524 0, /* tp_as_mapping */
3525 0, /* tp_hash */
3526 0, /* tp_call */
3527 0, /* tp_str */
3528 PyObject_GenericGetAttr, /* tp_getattro */
3529 0, /* tp_setattro */
3530 0, /* tp_as_buffer */
3531 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3532 Py_TPFLAGS_BASETYPE, /* tp_flags */
3533 accumulate_doc, /* tp_doc */
3534 (traverseproc)accumulate_traverse, /* tp_traverse */
3535 0, /* tp_clear */
3536 0, /* tp_richcompare */
3537 0, /* tp_weaklistoffset */
3538 PyObject_SelfIter, /* tp_iter */
3539 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003540 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003541 0, /* tp_members */
3542 0, /* tp_getset */
3543 0, /* tp_base */
3544 0, /* tp_dict */
3545 0, /* tp_descr_get */
3546 0, /* tp_descr_set */
3547 0, /* tp_dictoffset */
3548 0, /* tp_init */
3549 0, /* tp_alloc */
3550 accumulate_new, /* tp_new */
3551 PyObject_GC_Del, /* tp_free */
3552};
3553
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003554
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003555/* compress object ************************************************************/
3556
3557/* Equivalent to:
3558
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003559 def compress(data, selectors):
3560 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3561 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003562*/
3563
3564typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003565 PyObject_HEAD
3566 PyObject *data;
3567 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003568} compressobject;
3569
3570static PyTypeObject compress_type;
3571
3572static PyObject *
3573compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3574{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003575 PyObject *seq1, *seq2;
3576 PyObject *data=NULL, *selectors=NULL;
3577 compressobject *lz;
3578 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003579
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003580 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3581 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003582
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003583 data = PyObject_GetIter(seq1);
3584 if (data == NULL)
3585 goto fail;
3586 selectors = PyObject_GetIter(seq2);
3587 if (selectors == NULL)
3588 goto fail;
3589
3590 /* create compressobject structure */
3591 lz = (compressobject *)type->tp_alloc(type, 0);
3592 if (lz == NULL)
3593 goto fail;
3594 lz->data = data;
3595 lz->selectors = selectors;
3596 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003597
3598fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003599 Py_XDECREF(data);
3600 Py_XDECREF(selectors);
3601 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003602}
3603
3604static void
3605compress_dealloc(compressobject *lz)
3606{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003607 PyObject_GC_UnTrack(lz);
3608 Py_XDECREF(lz->data);
3609 Py_XDECREF(lz->selectors);
3610 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003611}
3612
3613static int
3614compress_traverse(compressobject *lz, visitproc visit, void *arg)
3615{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003616 Py_VISIT(lz->data);
3617 Py_VISIT(lz->selectors);
3618 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003619}
3620
3621static PyObject *
3622compress_next(compressobject *lz)
3623{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003624 PyObject *data = lz->data, *selectors = lz->selectors;
3625 PyObject *datum, *selector;
3626 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3627 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3628 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003629
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003630 while (1) {
3631 /* Steps: get datum, get selector, evaluate selector.
3632 Order is important (to match the pure python version
3633 in terms of which input gets a chance to raise an
3634 exception first).
3635 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003636
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003637 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003638 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003639 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003640
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003641 selector = selectornext(selectors);
3642 if (selector == NULL) {
3643 Py_DECREF(datum);
3644 return NULL;
3645 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003646
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003647 ok = PyObject_IsTrue(selector);
3648 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003649 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003650 return datum;
3651 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003652 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003653 return NULL;
3654 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003655}
3656
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003657static PyObject *
3658compress_reduce(compressobject *lz)
3659{
3660 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3661 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003662}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003663
3664static PyMethodDef compress_methods[] = {
3665 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3666 reduce_doc},
3667 {NULL, NULL} /* sentinel */
3668};
3669
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003670PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003671"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003672\n\
3673Return data elements corresponding to true selector elements.\n\
3674Forms a shorter iterator from selected data elements using the\n\
3675selectors to choose the data elements.");
3676
3677static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003678 PyVarObject_HEAD_INIT(NULL, 0)
3679 "itertools.compress", /* tp_name */
3680 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003681 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003682 /* methods */
3683 (destructor)compress_dealloc, /* tp_dealloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003684 0, /* tp_print */
3685 0, /* tp_getattr */
3686 0, /* tp_setattr */
3687 0, /* tp_reserved */
3688 0, /* tp_repr */
3689 0, /* tp_as_number */
3690 0, /* tp_as_sequence */
3691 0, /* tp_as_mapping */
3692 0, /* tp_hash */
3693 0, /* tp_call */
3694 0, /* tp_str */
3695 PyObject_GenericGetAttr, /* tp_getattro */
3696 0, /* tp_setattro */
3697 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003698 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003699 Py_TPFLAGS_BASETYPE, /* tp_flags */
3700 compress_doc, /* tp_doc */
3701 (traverseproc)compress_traverse, /* tp_traverse */
3702 0, /* tp_clear */
3703 0, /* tp_richcompare */
3704 0, /* tp_weaklistoffset */
3705 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003706 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003707 compress_methods, /* tp_methods */
3708 0, /* tp_members */
3709 0, /* tp_getset */
3710 0, /* tp_base */
3711 0, /* tp_dict */
3712 0, /* tp_descr_get */
3713 0, /* tp_descr_set */
3714 0, /* tp_dictoffset */
3715 0, /* tp_init */
3716 0, /* tp_alloc */
3717 compress_new, /* tp_new */
3718 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003719};
3720
3721
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003722/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003723
3724typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003725 PyObject_HEAD
3726 PyObject *func;
3727 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003728} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003729
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003730static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003731
3732static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003733filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003734{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003735 PyObject *func, *seq;
3736 PyObject *it;
3737 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003738
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003739 if (type == &filterfalse_type &&
3740 !_PyArg_NoKeywords("filterfalse()", kwds))
3741 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003742
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003743 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3744 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003745
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003746 /* Get iterator. */
3747 it = PyObject_GetIter(seq);
3748 if (it == NULL)
3749 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003750
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003751 /* create filterfalseobject structure */
3752 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3753 if (lz == NULL) {
3754 Py_DECREF(it);
3755 return NULL;
3756 }
3757 Py_INCREF(func);
3758 lz->func = func;
3759 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003760
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003761 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003762}
3763
3764static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003765filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003766{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003767 PyObject_GC_UnTrack(lz);
3768 Py_XDECREF(lz->func);
3769 Py_XDECREF(lz->it);
3770 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003771}
3772
3773static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003774filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003775{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003776 Py_VISIT(lz->it);
3777 Py_VISIT(lz->func);
3778 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003779}
3780
3781static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003782filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003783{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003784 PyObject *item;
3785 PyObject *it = lz->it;
3786 long ok;
3787 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003788
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003789 iternext = *Py_TYPE(it)->tp_iternext;
3790 for (;;) {
3791 item = iternext(it);
3792 if (item == NULL)
3793 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003794
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003795 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3796 ok = PyObject_IsTrue(item);
3797 } else {
3798 PyObject *good;
Raymond Hettingera6a2d442015-08-16 14:38:07 -07003799 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003800 if (good == NULL) {
3801 Py_DECREF(item);
3802 return NULL;
3803 }
3804 ok = PyObject_IsTrue(good);
3805 Py_DECREF(good);
3806 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003807 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003808 return item;
3809 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003810 if (ok < 0)
3811 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003812 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003813}
3814
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003815static PyObject *
3816filterfalse_reduce(filterfalseobject *lz)
3817{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003818 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3819}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003820
3821static PyMethodDef filterfalse_methods[] = {
3822 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3823 reduce_doc},
3824 {NULL, NULL} /* sentinel */
3825};
3826
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003827PyDoc_STRVAR(filterfalse_doc,
3828"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003829\n\
3830Return those items of sequence for which function(item) is false.\n\
3831If function is None, return the items that are false.");
3832
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003833static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003834 PyVarObject_HEAD_INIT(NULL, 0)
3835 "itertools.filterfalse", /* tp_name */
3836 sizeof(filterfalseobject), /* tp_basicsize */
3837 0, /* tp_itemsize */
3838 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003839 (destructor)filterfalse_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003840 0, /* tp_print */
3841 0, /* tp_getattr */
3842 0, /* tp_setattr */
3843 0, /* tp_reserved */
3844 0, /* tp_repr */
3845 0, /* tp_as_number */
3846 0, /* tp_as_sequence */
3847 0, /* tp_as_mapping */
3848 0, /* tp_hash */
3849 0, /* tp_call */
3850 0, /* tp_str */
3851 PyObject_GenericGetAttr, /* tp_getattro */
3852 0, /* tp_setattro */
3853 0, /* tp_as_buffer */
3854 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3855 Py_TPFLAGS_BASETYPE, /* tp_flags */
3856 filterfalse_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003857 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003858 0, /* tp_clear */
3859 0, /* tp_richcompare */
3860 0, /* tp_weaklistoffset */
3861 PyObject_SelfIter, /* tp_iter */
3862 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003863 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003864 0, /* tp_members */
3865 0, /* tp_getset */
3866 0, /* tp_base */
3867 0, /* tp_dict */
3868 0, /* tp_descr_get */
3869 0, /* tp_descr_set */
3870 0, /* tp_dictoffset */
3871 0, /* tp_init */
3872 0, /* tp_alloc */
3873 filterfalse_new, /* tp_new */
3874 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003875};
3876
3877
3878/* count object ************************************************************/
3879
3880typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003881 PyObject_HEAD
3882 Py_ssize_t cnt;
3883 PyObject *long_cnt;
3884 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003885} countobject;
3886
Raymond Hettinger30729212009-02-12 06:28:27 +00003887/* Counting logic and invariants:
3888
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003889fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003890
Serhiy Storchaka483405b2015-02-17 10:14:30 +02003891 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003892 Advances with: cnt += 1
3893 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003894
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003895slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003896
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003897 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3898 All counting is done with python objects (no overflows or underflows).
3899 Advances with: long_cnt += long_step
3900 Step may be zero -- effectively a slow version of repeat(cnt).
3901 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003902*/
3903
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003904static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003905
3906static PyObject *
3907count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3908{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003909 countobject *lz;
3910 int slow_mode = 0;
3911 Py_ssize_t cnt = 0;
3912 PyObject *long_cnt = NULL;
3913 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003914 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003915 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003916
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003917 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3918 kwlist, &long_cnt, &long_step))
3919 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003920
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003921 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3922 (long_step != NULL && !PyNumber_Check(long_step))) {
3923 PyErr_SetString(PyExc_TypeError, "a number is required");
3924 return NULL;
3925 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003926
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003927 if (long_cnt != NULL) {
3928 cnt = PyLong_AsSsize_t(long_cnt);
3929 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3930 PyErr_Clear();
3931 slow_mode = 1;
3932 }
3933 Py_INCREF(long_cnt);
3934 } else {
3935 cnt = 0;
3936 long_cnt = PyLong_FromLong(0);
3937 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003938
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003939 /* If not specified, step defaults to 1 */
3940 if (long_step == NULL) {
3941 long_step = PyLong_FromLong(1);
3942 if (long_step == NULL) {
3943 Py_DECREF(long_cnt);
3944 return NULL;
3945 }
3946 } else
3947 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003948
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003949 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003950
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003951 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003952 step = PyLong_AsLong(long_step);
3953 if (step != 1) {
3954 slow_mode = 1;
3955 if (step == -1 && PyErr_Occurred())
3956 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003957 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003958
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003959 if (slow_mode)
3960 cnt = PY_SSIZE_T_MAX;
3961 else
3962 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003963
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003964 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3965 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3966 assert(slow_mode ||
3967 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003968
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003969 /* create countobject structure */
3970 lz = (countobject *)type->tp_alloc(type, 0);
3971 if (lz == NULL) {
3972 Py_XDECREF(long_cnt);
3973 return NULL;
3974 }
3975 lz->cnt = cnt;
3976 lz->long_cnt = long_cnt;
3977 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003978
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003979 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003980}
3981
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003982static void
3983count_dealloc(countobject *lz)
3984{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003985 PyObject_GC_UnTrack(lz);
3986 Py_XDECREF(lz->long_cnt);
3987 Py_XDECREF(lz->long_step);
3988 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003989}
3990
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003991static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003992count_traverse(countobject *lz, visitproc visit, void *arg)
3993{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003994 Py_VISIT(lz->long_cnt);
3995 Py_VISIT(lz->long_step);
3996 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003997}
3998
3999static PyObject *
4000count_nextlong(countobject *lz)
4001{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004002 PyObject *long_cnt;
4003 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004005 long_cnt = lz->long_cnt;
4006 if (long_cnt == NULL) {
4007 /* Switch to slow_mode */
4008 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4009 if (long_cnt == NULL)
4010 return NULL;
4011 }
4012 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004013
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004014 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4015 if (stepped_up == NULL)
4016 return NULL;
4017 lz->long_cnt = stepped_up;
4018 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004019}
4020
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004021static PyObject *
4022count_next(countobject *lz)
4023{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004024 if (lz->cnt == PY_SSIZE_T_MAX)
4025 return count_nextlong(lz);
4026 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004027}
4028
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004029static PyObject *
4030count_repr(countobject *lz)
4031{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004032 if (lz->cnt != PY_SSIZE_T_MAX)
4033 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004034
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004035 if (PyLong_Check(lz->long_step)) {
4036 long step = PyLong_AsLong(lz->long_step);
4037 if (step == -1 && PyErr_Occurred()) {
4038 PyErr_Clear();
4039 }
4040 if (step == 1) {
4041 /* Don't display step when it is an integer equal to 1 */
4042 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
4043 }
4044 }
4045 return PyUnicode_FromFormat("count(%R, %R)",
4046 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004047}
4048
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004049static PyObject *
4050count_reduce(countobject *lz)
4051{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004052 if (lz->cnt == PY_SSIZE_T_MAX)
4053 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4054 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004055}
4056
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004057static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004058 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004059 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004060 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004061};
4062
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004063PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00004064 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004065\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004066Return a count object whose .__next__() method returns consecutive values.\n\
4067Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00004068 def count(firstval=0, step=1):\n\
Eli Benderskyc554f722013-09-03 06:37:19 -07004069 x = firstval\n\
4070 while 1:\n\
4071 yield x\n\
4072 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004073
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004074static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004075 PyVarObject_HEAD_INIT(NULL, 0)
4076 "itertools.count", /* tp_name */
4077 sizeof(countobject), /* tp_basicsize */
4078 0, /* tp_itemsize */
4079 /* methods */
4080 (destructor)count_dealloc, /* tp_dealloc */
4081 0, /* tp_print */
4082 0, /* tp_getattr */
4083 0, /* tp_setattr */
4084 0, /* tp_reserved */
4085 (reprfunc)count_repr, /* tp_repr */
4086 0, /* tp_as_number */
4087 0, /* tp_as_sequence */
4088 0, /* tp_as_mapping */
4089 0, /* tp_hash */
4090 0, /* tp_call */
4091 0, /* tp_str */
4092 PyObject_GenericGetAttr, /* tp_getattro */
4093 0, /* tp_setattro */
4094 0, /* tp_as_buffer */
4095 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004096 Py_TPFLAGS_BASETYPE, /* tp_flags */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004097 count_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004098 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004099 0, /* tp_clear */
4100 0, /* tp_richcompare */
4101 0, /* tp_weaklistoffset */
4102 PyObject_SelfIter, /* tp_iter */
4103 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004104 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004105 0, /* tp_members */
4106 0, /* tp_getset */
4107 0, /* tp_base */
4108 0, /* tp_dict */
4109 0, /* tp_descr_get */
4110 0, /* tp_descr_set */
4111 0, /* tp_dictoffset */
4112 0, /* tp_init */
4113 0, /* tp_alloc */
4114 count_new, /* tp_new */
4115 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004116};
4117
4118
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004119/* repeat object ************************************************************/
4120
4121typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004122 PyObject_HEAD
4123 PyObject *element;
4124 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004125} repeatobject;
4126
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004127static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004128
4129static PyObject *
4130repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4131{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004132 repeatobject *ro;
4133 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004134 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004135 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004136
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004137 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4138 &element, &cnt))
4139 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004140
Raymond Hettinger97d35552014-06-24 21:36:58 -07004141 if (kwds != NULL)
4142 n_kwds = PyDict_Size(kwds);
4143 /* Does user supply times argument? */
4144 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004145 cnt = 0;
4146
4147 ro = (repeatobject *)type->tp_alloc(type, 0);
4148 if (ro == NULL)
4149 return NULL;
4150 Py_INCREF(element);
4151 ro->element = element;
4152 ro->cnt = cnt;
4153 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004154}
4155
4156static void
4157repeat_dealloc(repeatobject *ro)
4158{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004159 PyObject_GC_UnTrack(ro);
4160 Py_XDECREF(ro->element);
4161 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004162}
4163
4164static int
4165repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4166{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004167 Py_VISIT(ro->element);
4168 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004169}
4170
4171static PyObject *
4172repeat_next(repeatobject *ro)
4173{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004174 if (ro->cnt == 0)
4175 return NULL;
4176 if (ro->cnt > 0)
4177 ro->cnt--;
4178 Py_INCREF(ro->element);
4179 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004180}
4181
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004182static PyObject *
4183repeat_repr(repeatobject *ro)
4184{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004185 if (ro->cnt == -1)
4186 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4187 else
4188 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4189}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004190
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004191static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004192repeat_len(repeatobject *ro)
4193{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004194 if (ro->cnt == -1) {
4195 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4196 return NULL;
4197 }
4198 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004199}
4200
Armin Rigof5b3e362006-02-11 21:32:43 +00004201PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004202
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004203static PyObject *
4204repeat_reduce(repeatobject *ro)
4205{
4206 /* unpickle this so that a new repeat iterator is constructed with an
4207 * object, then call __setstate__ on it to set cnt
4208 */
4209 if (ro->cnt >= 0)
4210 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4211 else
4212 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4213}
4214
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004215static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004216 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004217 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004218 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004219};
4220
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004221PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004222"repeat(object [,times]) -> create an iterator which returns the object\n\
4223for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004224endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004225
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004226static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004227 PyVarObject_HEAD_INIT(NULL, 0)
4228 "itertools.repeat", /* tp_name */
4229 sizeof(repeatobject), /* tp_basicsize */
4230 0, /* tp_itemsize */
4231 /* methods */
4232 (destructor)repeat_dealloc, /* tp_dealloc */
4233 0, /* tp_print */
4234 0, /* tp_getattr */
4235 0, /* tp_setattr */
4236 0, /* tp_reserved */
4237 (reprfunc)repeat_repr, /* tp_repr */
4238 0, /* tp_as_number */
4239 0, /* tp_as_sequence */
4240 0, /* tp_as_mapping */
4241 0, /* tp_hash */
4242 0, /* tp_call */
4243 0, /* tp_str */
4244 PyObject_GenericGetAttr, /* tp_getattro */
4245 0, /* tp_setattro */
4246 0, /* tp_as_buffer */
4247 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4248 Py_TPFLAGS_BASETYPE, /* tp_flags */
4249 repeat_doc, /* tp_doc */
4250 (traverseproc)repeat_traverse, /* tp_traverse */
4251 0, /* tp_clear */
4252 0, /* tp_richcompare */
4253 0, /* tp_weaklistoffset */
4254 PyObject_SelfIter, /* tp_iter */
4255 (iternextfunc)repeat_next, /* tp_iternext */
4256 repeat_methods, /* tp_methods */
4257 0, /* tp_members */
4258 0, /* tp_getset */
4259 0, /* tp_base */
4260 0, /* tp_dict */
4261 0, /* tp_descr_get */
4262 0, /* tp_descr_set */
4263 0, /* tp_dictoffset */
4264 0, /* tp_init */
4265 0, /* tp_alloc */
4266 repeat_new, /* tp_new */
4267 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004268};
4269
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004270/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004271
4272typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004273 PyObject_HEAD
4274 Py_ssize_t tuplesize;
4275 Py_ssize_t numactive;
4276 PyObject *ittuple; /* tuple of iterators */
4277 PyObject *result;
4278 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004279} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004280
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004281static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004282
4283static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004284zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004285{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004286 ziplongestobject *lz;
4287 Py_ssize_t i;
4288 PyObject *ittuple; /* tuple of iterators */
4289 PyObject *result;
4290 PyObject *fillvalue = Py_None;
4291 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004292
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004293 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
4294 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4295 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
4296 PyErr_SetString(PyExc_TypeError,
4297 "zip_longest() got an unexpected keyword argument");
4298 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004299 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004300 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004301
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004302 /* args must be a tuple */
4303 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00004304
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004305 /* obtain iterators */
4306 ittuple = PyTuple_New(tuplesize);
4307 if (ittuple == NULL)
4308 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004309 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004310 PyObject *item = PyTuple_GET_ITEM(args, i);
4311 PyObject *it = PyObject_GetIter(item);
4312 if (it == NULL) {
4313 if (PyErr_ExceptionMatches(PyExc_TypeError))
4314 PyErr_Format(PyExc_TypeError,
4315 "zip_longest argument #%zd must support iteration",
4316 i+1);
4317 Py_DECREF(ittuple);
4318 return NULL;
4319 }
4320 PyTuple_SET_ITEM(ittuple, i, it);
4321 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004322
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004323 /* create a result holder */
4324 result = PyTuple_New(tuplesize);
4325 if (result == NULL) {
4326 Py_DECREF(ittuple);
4327 return NULL;
4328 }
4329 for (i=0 ; i < tuplesize ; i++) {
4330 Py_INCREF(Py_None);
4331 PyTuple_SET_ITEM(result, i, Py_None);
4332 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004333
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004334 /* create ziplongestobject structure */
4335 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4336 if (lz == NULL) {
4337 Py_DECREF(ittuple);
4338 Py_DECREF(result);
4339 return NULL;
4340 }
4341 lz->ittuple = ittuple;
4342 lz->tuplesize = tuplesize;
4343 lz->numactive = tuplesize;
4344 lz->result = result;
4345 Py_INCREF(fillvalue);
4346 lz->fillvalue = fillvalue;
4347 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004348}
4349
4350static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004351zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004352{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004353 PyObject_GC_UnTrack(lz);
4354 Py_XDECREF(lz->ittuple);
4355 Py_XDECREF(lz->result);
4356 Py_XDECREF(lz->fillvalue);
4357 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004358}
4359
4360static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004361zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004362{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004363 Py_VISIT(lz->ittuple);
4364 Py_VISIT(lz->result);
4365 Py_VISIT(lz->fillvalue);
4366 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004367}
4368
4369static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004370zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004371{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004372 Py_ssize_t i;
4373 Py_ssize_t tuplesize = lz->tuplesize;
4374 PyObject *result = lz->result;
4375 PyObject *it;
4376 PyObject *item;
4377 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004378
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004379 if (tuplesize == 0)
4380 return NULL;
4381 if (lz->numactive == 0)
4382 return NULL;
4383 if (Py_REFCNT(result) == 1) {
4384 Py_INCREF(result);
4385 for (i=0 ; i < tuplesize ; i++) {
4386 it = PyTuple_GET_ITEM(lz->ittuple, i);
4387 if (it == NULL) {
4388 Py_INCREF(lz->fillvalue);
4389 item = lz->fillvalue;
4390 } else {
4391 item = PyIter_Next(it);
4392 if (item == NULL) {
4393 lz->numactive -= 1;
4394 if (lz->numactive == 0 || PyErr_Occurred()) {
4395 lz->numactive = 0;
4396 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004397 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004398 } else {
4399 Py_INCREF(lz->fillvalue);
4400 item = lz->fillvalue;
4401 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4402 Py_DECREF(it);
4403 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004404 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004405 }
4406 olditem = PyTuple_GET_ITEM(result, i);
4407 PyTuple_SET_ITEM(result, i, item);
4408 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004409 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004410 } else {
4411 result = PyTuple_New(tuplesize);
4412 if (result == NULL)
4413 return NULL;
4414 for (i=0 ; i < tuplesize ; i++) {
4415 it = PyTuple_GET_ITEM(lz->ittuple, i);
4416 if (it == NULL) {
4417 Py_INCREF(lz->fillvalue);
4418 item = lz->fillvalue;
4419 } else {
4420 item = PyIter_Next(it);
4421 if (item == NULL) {
4422 lz->numactive -= 1;
4423 if (lz->numactive == 0 || PyErr_Occurred()) {
4424 lz->numactive = 0;
4425 Py_DECREF(result);
4426 return NULL;
4427 } else {
4428 Py_INCREF(lz->fillvalue);
4429 item = lz->fillvalue;
4430 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4431 Py_DECREF(it);
4432 }
4433 }
4434 }
4435 PyTuple_SET_ITEM(result, i, item);
4436 }
4437 }
4438 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004439}
4440
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004441static PyObject *
4442zip_longest_reduce(ziplongestobject *lz)
4443{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004444
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004445 /* Create a new tuple with empty sequences where appropriate to pickle.
4446 * Then use setstate to set the fillvalue
4447 */
4448 int i;
4449 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004450
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004451 if (args == NULL)
4452 return NULL;
4453 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4454 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4455 if (elem == NULL) {
4456 elem = PyTuple_New(0);
4457 if (elem == NULL) {
4458 Py_DECREF(args);
4459 return NULL;
4460 }
4461 } else
4462 Py_INCREF(elem);
4463 PyTuple_SET_ITEM(args, i, elem);
4464 }
4465 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4466}
4467
4468static PyObject *
4469zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4470{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004471 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004472 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004473 Py_RETURN_NONE;
4474}
4475
4476static PyMethodDef zip_longest_methods[] = {
4477 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4478 reduce_doc},
4479 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4480 setstate_doc},
4481 {NULL, NULL} /* sentinel */
4482};
4483
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004484PyDoc_STRVAR(zip_longest_doc,
4485"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004486\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004487Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004488the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004489method continues until the longest iterable in the argument sequence\n\
4490is exhausted and then it raises StopIteration. When the shorter iterables\n\
4491are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4492defaults to None or can be specified by a keyword argument.\n\
4493");
4494
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004495static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004496 PyVarObject_HEAD_INIT(NULL, 0)
4497 "itertools.zip_longest", /* tp_name */
4498 sizeof(ziplongestobject), /* tp_basicsize */
4499 0, /* tp_itemsize */
4500 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004501 (destructor)zip_longest_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004502 0, /* tp_print */
4503 0, /* tp_getattr */
4504 0, /* tp_setattr */
4505 0, /* tp_reserved */
4506 0, /* tp_repr */
4507 0, /* tp_as_number */
4508 0, /* tp_as_sequence */
4509 0, /* tp_as_mapping */
4510 0, /* tp_hash */
4511 0, /* tp_call */
4512 0, /* tp_str */
4513 PyObject_GenericGetAttr, /* tp_getattro */
4514 0, /* tp_setattro */
4515 0, /* tp_as_buffer */
4516 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4517 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004518 zip_longest_doc, /* tp_doc */
4519 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004520 0, /* tp_clear */
4521 0, /* tp_richcompare */
4522 0, /* tp_weaklistoffset */
4523 PyObject_SelfIter, /* tp_iter */
4524 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004525 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004526 0, /* tp_members */
4527 0, /* tp_getset */
4528 0, /* tp_base */
4529 0, /* tp_dict */
4530 0, /* tp_descr_get */
4531 0, /* tp_descr_set */
4532 0, /* tp_dictoffset */
4533 0, /* tp_init */
4534 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004535 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004536 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004537};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004538
4539/* module level code ********************************************************/
4540
4541PyDoc_STRVAR(module_doc,
4542"Functional tools for creating and using iterators.\n\
4543\n\
4544Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004545count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004546cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004547repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004548\n\
4549Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004550accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004551chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger8df58f72013-09-09 01:29:40 -05004552chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004553compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4554dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4555groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004556filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004557islice(seq, [start,] stop [, step]) --> elements from\n\
4558 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004559starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004560tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004561takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004562zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004563\n\
4564Combinatoric generators:\n\
4565product(p, q, ... [repeat=1]) --> cartesian product\n\
4566permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004567combinations(p, r)\n\
4568combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004569");
4570
4571
Raymond Hettingerad983e72003-11-12 14:32:26 +00004572static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004573 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4574 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004575};
4576
Martin v. Löwis1a214512008-06-11 05:26:20 +00004577
4578static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004579 PyModuleDef_HEAD_INIT,
4580 "itertools",
4581 module_doc,
4582 -1,
4583 module_methods,
4584 NULL,
4585 NULL,
4586 NULL,
4587 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004588};
4589
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004590PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004591PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004592{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004593 int i;
4594 PyObject *m;
4595 char *name;
4596 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004597 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004598 &combinations_type,
4599 &cwr_type,
4600 &cycle_type,
4601 &dropwhile_type,
4602 &takewhile_type,
4603 &islice_type,
4604 &starmap_type,
4605 &chain_type,
4606 &compress_type,
4607 &filterfalse_type,
4608 &count_type,
4609 &ziplongest_type,
4610 &permutations_type,
4611 &product_type,
4612 &repeat_type,
4613 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004614 &_grouper_type,
4615 &tee_type,
4616 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004617 NULL
4618 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004619
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004620 Py_TYPE(&teedataobject_type) = &PyType_Type;
4621 m = PyModule_Create(&itertoolsmodule);
4622 if (m == NULL)
4623 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004624
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004625 for (i=0 ; typelist[i] != NULL ; i++) {
4626 if (PyType_Ready(typelist[i]) < 0)
4627 return NULL;
4628 name = strchr(typelist[i]->tp_name, '.');
4629 assert (name != NULL);
4630 Py_INCREF(typelist[i]);
4631 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4632 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004633
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004634 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004635}