blob: afff7e4aa87125e4cd1f9847dff49734297e608d [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
2#include "Python.h"
Fred Drake08ebfec2004-10-17 19:36:57 +00003#include "structmember.h"
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00005/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00006 by Raymond D. Hettinger <python@rcn.com>
Raymond Hettinger8df58f72013-09-09 01:29:40 -05007 Copyright (c) 2003-2013 Python Software Foundation.
Raymond Hettinger96ef8112003-02-01 00:10:11 +00008 All rights reserved.
9*/
10
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000011
12/* groupby object ***********************************************************/
13
14typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000015 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000021} groupbyobject;
22
23static PyTypeObject groupby_type;
24static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26static PyObject *
27groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000029 static char *kwargs[] = {"iterable", "key", NULL};
30 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000032
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000033 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
34 &it, &keyfunc))
35 return NULL;
36
37 gbo = (groupbyobject *)type->tp_alloc(type, 0);
38 if (gbo == NULL)
39 return NULL;
40 gbo->tgtkey = NULL;
41 gbo->currkey = NULL;
42 gbo->currvalue = NULL;
43 gbo->keyfunc = keyfunc;
44 Py_INCREF(keyfunc);
45 gbo->it = PyObject_GetIter(it);
46 if (gbo->it == NULL) {
47 Py_DECREF(gbo);
48 return NULL;
49 }
50 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000051}
52
53static void
54groupby_dealloc(groupbyobject *gbo)
55{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000056 PyObject_GC_UnTrack(gbo);
57 Py_XDECREF(gbo->it);
58 Py_XDECREF(gbo->keyfunc);
59 Py_XDECREF(gbo->tgtkey);
60 Py_XDECREF(gbo->currkey);
61 Py_XDECREF(gbo->currvalue);
62 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000063}
64
65static int
66groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
67{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000068 Py_VISIT(gbo->it);
69 Py_VISIT(gbo->keyfunc);
70 Py_VISIT(gbo->tgtkey);
71 Py_VISIT(gbo->currkey);
72 Py_VISIT(gbo->currvalue);
73 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000074}
75
76static PyObject *
77groupby_next(groupbyobject *gbo)
78{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000081 /* skip to next iteration group */
82 for (;;) {
83 if (gbo->currkey == NULL)
84 /* pass */;
85 else if (gbo->tgtkey == NULL)
86 break;
87 else {
88 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000089
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000090 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
91 gbo->currkey, Py_EQ);
92 if (rcmp == -1)
93 return NULL;
94 else if (rcmp == 0)
95 break;
96 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000097
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000098 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000101
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000102 if (gbo->keyfunc == Py_None) {
103 newkey = newvalue;
104 Py_INCREF(newvalue);
105 } else {
106 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
107 newvalue, NULL);
108 if (newkey == NULL) {
109 Py_DECREF(newvalue);
110 return NULL;
111 }
112 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000113
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000114 tmp = gbo->currkey;
115 gbo->currkey = newkey;
116 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000117
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000118 tmp = gbo->currvalue;
119 gbo->currvalue = newvalue;
120 Py_XDECREF(tmp);
121 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000123 Py_INCREF(gbo->currkey);
124 tmp = gbo->tgtkey;
125 gbo->tgtkey = gbo->currkey;
126 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000127
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000131
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000135}
136
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000137static PyObject *
138groupby_reduce(groupbyobject *lz)
139{
140 /* reduce as a 'new' call with an optional 'setstate' if groupby
141 * has started
142 */
143 PyObject *value;
144 if (lz->tgtkey && lz->currkey && lz->currvalue)
145 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
146 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
147 else
148 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
149 lz->it, lz->keyfunc);
150
151 return value;
152}
153
154PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
155
156static PyObject *
157groupby_setstate(groupbyobject *lz, PyObject *state)
158{
159 PyObject *currkey, *currvalue, *tgtkey;
160 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey))
161 return NULL;
162 Py_CLEAR(lz->currkey);
163 lz->currkey = currkey;
164 Py_INCREF(lz->currkey);
165 Py_CLEAR(lz->currvalue);
166 lz->currvalue = currvalue;
167 Py_INCREF(lz->currvalue);
168 Py_CLEAR(lz->tgtkey);
169 lz->tgtkey = tgtkey;
170 Py_INCREF(lz->tgtkey);
171 Py_RETURN_NONE;
172}
173
174PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
175
176static PyMethodDef groupby_methods[] = {
177 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
178 reduce_doc},
179 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
180 setstate_doc},
181 {NULL, NULL} /* sentinel */
182};
183
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000184PyDoc_STRVAR(groupby_doc,
185"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
186(key, sub-iterator) grouped by each value of key(value).\n");
187
188static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000189 PyVarObject_HEAD_INIT(NULL, 0)
190 "itertools.groupby", /* tp_name */
191 sizeof(groupbyobject), /* tp_basicsize */
192 0, /* tp_itemsize */
193 /* methods */
194 (destructor)groupby_dealloc, /* tp_dealloc */
195 0, /* tp_print */
196 0, /* tp_getattr */
197 0, /* tp_setattr */
198 0, /* tp_reserved */
199 0, /* tp_repr */
200 0, /* tp_as_number */
201 0, /* tp_as_sequence */
202 0, /* tp_as_mapping */
203 0, /* tp_hash */
204 0, /* tp_call */
205 0, /* tp_str */
206 PyObject_GenericGetAttr, /* tp_getattro */
207 0, /* tp_setattro */
208 0, /* tp_as_buffer */
209 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
210 Py_TPFLAGS_BASETYPE, /* tp_flags */
211 groupby_doc, /* tp_doc */
212 (traverseproc)groupby_traverse, /* tp_traverse */
213 0, /* tp_clear */
214 0, /* tp_richcompare */
215 0, /* tp_weaklistoffset */
216 PyObject_SelfIter, /* tp_iter */
217 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000218 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000219 0, /* tp_members */
220 0, /* tp_getset */
221 0, /* tp_base */
222 0, /* tp_dict */
223 0, /* tp_descr_get */
224 0, /* tp_descr_set */
225 0, /* tp_dictoffset */
226 0, /* tp_init */
227 0, /* tp_alloc */
228 groupby_new, /* tp_new */
229 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000230};
231
232
233/* _grouper object (internal) ************************************************/
234
235typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000236 PyObject_HEAD
237 PyObject *parent;
238 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000239} _grouperobject;
240
241static PyTypeObject _grouper_type;
242
243static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000244_grouper_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
245{
246 PyObject *parent, *tgtkey;
247
248 if (!PyArg_ParseTuple(args, "O!O", &groupby_type, &parent, &tgtkey))
249 return NULL;
250
251 return _grouper_create((groupbyobject*) parent, tgtkey);
252}
253
254static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000255_grouper_create(groupbyobject *parent, PyObject *tgtkey)
256{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000257 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000259 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
260 if (igo == NULL)
261 return NULL;
262 igo->parent = (PyObject *)parent;
263 Py_INCREF(parent);
264 igo->tgtkey = tgtkey;
265 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000266
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000267 PyObject_GC_Track(igo);
268 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000269}
270
271static void
272_grouper_dealloc(_grouperobject *igo)
273{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000274 PyObject_GC_UnTrack(igo);
275 Py_DECREF(igo->parent);
276 Py_DECREF(igo->tgtkey);
277 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000278}
279
280static int
281_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
282{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000283 Py_VISIT(igo->parent);
284 Py_VISIT(igo->tgtkey);
285 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000286}
287
288static PyObject *
289_grouper_next(_grouperobject *igo)
290{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000291 groupbyobject *gbo = (groupbyobject *)igo->parent;
292 PyObject *newvalue, *newkey, *r;
293 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000294
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000295 if (gbo->currvalue == NULL) {
296 newvalue = PyIter_Next(gbo->it);
297 if (newvalue == NULL)
298 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000299
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000300 if (gbo->keyfunc == Py_None) {
301 newkey = newvalue;
302 Py_INCREF(newvalue);
303 } else {
304 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
305 newvalue, NULL);
306 if (newkey == NULL) {
307 Py_DECREF(newvalue);
308 return NULL;
309 }
310 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000311
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000312 assert(gbo->currkey == NULL);
313 gbo->currkey = newkey;
314 gbo->currvalue = newvalue;
315 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000316
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000317 assert(gbo->currkey != NULL);
318 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
319 if (rcmp <= 0)
320 /* got any error or current group is end */
321 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000322
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000323 r = gbo->currvalue;
324 gbo->currvalue = NULL;
325 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000326
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000327 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000328}
329
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000330static PyObject *
331_grouper_reduce(_grouperobject *lz)
332{
333 return Py_BuildValue("O(OO)", Py_TYPE(lz),
334 lz->parent, lz->tgtkey);
335}
336
337static PyMethodDef _grouper_methods[] = {
338 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
339 reduce_doc},
340 {NULL, NULL} /* sentinel */
341};
342
343
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000344static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000345 PyVarObject_HEAD_INIT(NULL, 0)
346 "itertools._grouper", /* tp_name */
347 sizeof(_grouperobject), /* tp_basicsize */
348 0, /* tp_itemsize */
349 /* methods */
350 (destructor)_grouper_dealloc, /* tp_dealloc */
351 0, /* tp_print */
352 0, /* tp_getattr */
353 0, /* tp_setattr */
354 0, /* tp_reserved */
355 0, /* tp_repr */
356 0, /* tp_as_number */
357 0, /* tp_as_sequence */
358 0, /* tp_as_mapping */
359 0, /* tp_hash */
360 0, /* tp_call */
361 0, /* tp_str */
362 PyObject_GenericGetAttr, /* tp_getattro */
363 0, /* tp_setattro */
364 0, /* tp_as_buffer */
365 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
366 0, /* tp_doc */
367 (traverseproc)_grouper_traverse,/* tp_traverse */
368 0, /* tp_clear */
369 0, /* tp_richcompare */
370 0, /* tp_weaklistoffset */
371 PyObject_SelfIter, /* tp_iter */
372 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000373 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000374 0, /* tp_members */
375 0, /* tp_getset */
376 0, /* tp_base */
377 0, /* tp_dict */
378 0, /* tp_descr_get */
379 0, /* tp_descr_set */
380 0, /* tp_dictoffset */
381 0, /* tp_init */
382 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000383 _grouper_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000384 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000385};
386
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000387
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000388
Raymond Hettingerad983e72003-11-12 14:32:26 +0000389/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000390
Raymond Hettingerad983e72003-11-12 14:32:26 +0000391/* The teedataobject pre-allocates space for LINKCELLS number of objects.
392 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000393 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000394 the other structure members including PyHEAD overhead. The larger the
395 value, the less memory overhead per object and the less time spent
396 allocating/deallocating new links. The smaller the number, the less
397 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000398*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000399#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000400
401typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000402 PyObject_HEAD
403 PyObject *it;
Antoine Pitroub4a46cb2013-09-20 22:19:22 +0200404 int numread; /* 0 <= numread <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000405 PyObject *nextlink;
406 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000407} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000408
409typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000410 PyObject_HEAD
411 teedataobject *dataobj;
Antoine Pitroub4a46cb2013-09-20 22:19:22 +0200412 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000413 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000414} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415
Raymond Hettingerad983e72003-11-12 14:32:26 +0000416static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417
418static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000419teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000420{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000421 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000422
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000423 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
424 if (tdo == NULL)
425 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000426
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000427 tdo->numread = 0;
428 tdo->nextlink = NULL;
429 Py_INCREF(it);
430 tdo->it = it;
431 PyObject_GC_Track(tdo);
432 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000433}
434
435static PyObject *
436teedataobject_jumplink(teedataobject *tdo)
437{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000438 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000439 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000440 Py_XINCREF(tdo->nextlink);
441 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000442}
443
444static PyObject *
445teedataobject_getitem(teedataobject *tdo, int i)
446{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000447 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000448
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000449 assert(i < LINKCELLS);
450 if (i < tdo->numread)
451 value = tdo->values[i];
452 else {
453 /* this is the lead iterator, so fetch more data */
454 assert(i == tdo->numread);
455 value = PyIter_Next(tdo->it);
456 if (value == NULL)
457 return NULL;
458 tdo->numread++;
459 tdo->values[i] = value;
460 }
461 Py_INCREF(value);
462 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000463}
464
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000465static int
466teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
467{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000468 int i;
469 Py_VISIT(tdo->it);
470 for (i = 0; i < tdo->numread; i++)
471 Py_VISIT(tdo->values[i]);
472 Py_VISIT(tdo->nextlink);
473 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000474}
475
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200476static void
477teedataobject_safe_decref(PyObject *obj)
478{
479 while (obj && Py_TYPE(obj) == &teedataobject_type &&
480 Py_REFCNT(obj) == 1) {
481 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
482 ((teedataobject *)obj)->nextlink = NULL;
483 Py_DECREF(obj);
484 obj = nextlink;
485 }
486 Py_XDECREF(obj);
487}
488
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000489static int
490teedataobject_clear(teedataobject *tdo)
491{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000492 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200493 PyObject *tmp;
494
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000495 Py_CLEAR(tdo->it);
496 for (i=0 ; i<tdo->numread ; i++)
497 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200498 tmp = tdo->nextlink;
499 tdo->nextlink = NULL;
500 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000501 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000502}
503
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000504static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000505teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000506{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000507 PyObject_GC_UnTrack(tdo);
508 teedataobject_clear(tdo);
509 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000510}
511
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000512static PyObject *
513teedataobject_reduce(teedataobject *tdo)
514{
515 int i;
516 /* create a temporary list of already iterated values */
517 PyObject *values = PyList_New(tdo->numread);
518 if (!values)
519 return NULL;
520 for (i=0 ; i<tdo->numread ; i++) {
521 Py_INCREF(tdo->values[i]);
522 PyList_SET_ITEM(values, i, tdo->values[i]);
523 }
524 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
525 values,
526 tdo->nextlink ? tdo->nextlink : Py_None);
527}
528
529static PyTypeObject teedataobject_type;
530
531static PyObject *
532teedataobject_new(PyTypeObject *type, PyObject *args, PyObject *kw)
533{
534 teedataobject *tdo;
535 PyObject *it, *values, *next;
536 Py_ssize_t i, len;
537
538 assert(type == &teedataobject_type);
539 if (!PyArg_ParseTuple(args, "OO!O", &it, &PyList_Type, &values, &next))
540 return NULL;
541
542 tdo = (teedataobject *)teedataobject_newinternal(it);
543 if (!tdo)
544 return NULL;
545
546 len = PyList_GET_SIZE(values);
547 if (len > LINKCELLS)
548 goto err;
549 for (i=0; i<len; i++) {
550 tdo->values[i] = PyList_GET_ITEM(values, i);
551 Py_INCREF(tdo->values[i]);
552 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200553 /* len <= LINKCELLS < INT_MAX */
554 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000555
556 if (len == LINKCELLS) {
557 if (next != Py_None) {
558 if (Py_TYPE(next) != &teedataobject_type)
559 goto err;
560 assert(tdo->nextlink == NULL);
561 Py_INCREF(next);
562 tdo->nextlink = next;
563 }
564 } else {
565 if (next != Py_None)
566 goto err; /* shouldn't have a next if we are not full */
567 }
568 return (PyObject*)tdo;
569
570err:
571 Py_XDECREF(tdo);
572 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
573 return NULL;
574}
575
576static PyMethodDef teedataobject_methods[] = {
577 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
578 reduce_doc},
579 {NULL, NULL} /* sentinel */
580};
581
Raymond Hettingerad983e72003-11-12 14:32:26 +0000582PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000583
Raymond Hettingerad983e72003-11-12 14:32:26 +0000584static PyTypeObject teedataobject_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000585 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000586 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000587 sizeof(teedataobject), /* tp_basicsize */
588 0, /* tp_itemsize */
589 /* methods */
590 (destructor)teedataobject_dealloc, /* tp_dealloc */
591 0, /* tp_print */
592 0, /* tp_getattr */
593 0, /* tp_setattr */
594 0, /* tp_reserved */
595 0, /* tp_repr */
596 0, /* tp_as_number */
597 0, /* tp_as_sequence */
598 0, /* tp_as_mapping */
599 0, /* tp_hash */
600 0, /* tp_call */
601 0, /* tp_str */
602 PyObject_GenericGetAttr, /* tp_getattro */
603 0, /* tp_setattro */
604 0, /* tp_as_buffer */
605 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
606 teedataobject_doc, /* tp_doc */
607 (traverseproc)teedataobject_traverse, /* tp_traverse */
608 (inquiry)teedataobject_clear, /* tp_clear */
609 0, /* tp_richcompare */
610 0, /* tp_weaklistoffset */
611 0, /* tp_iter */
612 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000613 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000614 0, /* tp_members */
615 0, /* tp_getset */
616 0, /* tp_base */
617 0, /* tp_dict */
618 0, /* tp_descr_get */
619 0, /* tp_descr_set */
620 0, /* tp_dictoffset */
621 0, /* tp_init */
622 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000623 teedataobject_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000624 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000625};
626
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000627
628static PyTypeObject tee_type;
629
630static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000631tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000632{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000633 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000634
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000635 if (to->index >= LINKCELLS) {
636 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200637 if (link == NULL)
638 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000639 Py_DECREF(to->dataobj);
640 to->dataobj = (teedataobject *)link;
641 to->index = 0;
642 }
643 value = teedataobject_getitem(to->dataobj, to->index);
644 if (value == NULL)
645 return NULL;
646 to->index++;
647 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000648}
649
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000650static int
651tee_traverse(teeobject *to, visitproc visit, void *arg)
652{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000653 Py_VISIT((PyObject *)to->dataobj);
654 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000655}
656
Raymond Hettingerad983e72003-11-12 14:32:26 +0000657static PyObject *
658tee_copy(teeobject *to)
659{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000660 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000661
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000662 newto = PyObject_GC_New(teeobject, &tee_type);
663 if (newto == NULL)
664 return NULL;
665 Py_INCREF(to->dataobj);
666 newto->dataobj = to->dataobj;
667 newto->index = to->index;
668 newto->weakreflist = NULL;
669 PyObject_GC_Track(newto);
670 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000671}
672
673PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
674
675static PyObject *
676tee_fromiterable(PyObject *iterable)
677{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000678 teeobject *to;
679 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000680
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000681 it = PyObject_GetIter(iterable);
682 if (it == NULL)
683 return NULL;
684 if (PyObject_TypeCheck(it, &tee_type)) {
685 to = (teeobject *)tee_copy((teeobject *)it);
686 goto done;
687 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000688
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000689 to = PyObject_GC_New(teeobject, &tee_type);
690 if (to == NULL)
691 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000692 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000693 if (!to->dataobj) {
694 PyObject_GC_Del(to);
695 to = NULL;
696 goto done;
697 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000698
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000699 to->index = 0;
700 to->weakreflist = NULL;
701 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000702done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000703 Py_XDECREF(it);
704 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000705}
706
707static PyObject *
708tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
709{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000710 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000711
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000712 if (!PyArg_UnpackTuple(args, "_tee", 1, 1, &iterable))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000713 return NULL;
714 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000715}
716
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000717static int
718tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000719{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000720 if (to->weakreflist != NULL)
721 PyObject_ClearWeakRefs((PyObject *) to);
722 Py_CLEAR(to->dataobj);
723 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000724}
725
726static void
727tee_dealloc(teeobject *to)
728{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000729 PyObject_GC_UnTrack(to);
730 tee_clear(to);
731 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000732}
733
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000734static PyObject *
735tee_reduce(teeobject *to)
736{
737 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
738}
739
740static PyObject *
741tee_setstate(teeobject *to, PyObject *state)
742{
743 teedataobject *tdo;
744 int index;
745 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index))
746 return NULL;
747 if (index < 0 || index > LINKCELLS) {
748 PyErr_SetString(PyExc_ValueError, "Index out of range");
749 return NULL;
750 }
751 Py_CLEAR(to->dataobj);
752 to->dataobj = tdo;
753 Py_INCREF(to->dataobj);
754 to->index = index;
755 Py_RETURN_NONE;
756}
757
Raymond Hettingerad983e72003-11-12 14:32:26 +0000758PyDoc_STRVAR(teeobject_doc,
759"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000760
Raymond Hettingerad983e72003-11-12 14:32:26 +0000761static PyMethodDef tee_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000762 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000763 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
764 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000765 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000766};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000767
768static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000769 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000770 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000771 sizeof(teeobject), /* tp_basicsize */
772 0, /* tp_itemsize */
773 /* methods */
774 (destructor)tee_dealloc, /* tp_dealloc */
775 0, /* tp_print */
776 0, /* tp_getattr */
777 0, /* tp_setattr */
778 0, /* tp_reserved */
779 0, /* tp_repr */
780 0, /* tp_as_number */
781 0, /* tp_as_sequence */
782 0, /* tp_as_mapping */
783 0, /* tp_hash */
784 0, /* tp_call */
785 0, /* tp_str */
786 0, /* tp_getattro */
787 0, /* tp_setattro */
788 0, /* tp_as_buffer */
789 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
790 teeobject_doc, /* tp_doc */
791 (traverseproc)tee_traverse, /* tp_traverse */
792 (inquiry)tee_clear, /* tp_clear */
793 0, /* tp_richcompare */
794 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
795 PyObject_SelfIter, /* tp_iter */
796 (iternextfunc)tee_next, /* tp_iternext */
797 tee_methods, /* tp_methods */
798 0, /* tp_members */
799 0, /* tp_getset */
800 0, /* tp_base */
801 0, /* tp_dict */
802 0, /* tp_descr_get */
803 0, /* tp_descr_set */
804 0, /* tp_dictoffset */
805 0, /* tp_init */
806 0, /* tp_alloc */
807 tee_new, /* tp_new */
808 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000809};
810
Raymond Hettingerad983e72003-11-12 14:32:26 +0000811static PyObject *
812tee(PyObject *self, PyObject *args)
813{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000814 Py_ssize_t i, n=2;
815 PyObject *it, *iterable, *copyable, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200816 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000817
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000818 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
819 return NULL;
820 if (n < 0) {
821 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
822 return NULL;
823 }
824 result = PyTuple_New(n);
825 if (result == NULL)
826 return NULL;
827 if (n == 0)
828 return result;
829 it = PyObject_GetIter(iterable);
830 if (it == NULL) {
831 Py_DECREF(result);
832 return NULL;
833 }
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200834 if (!_PyObject_HasAttrId(it, &PyId___copy__)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000835 copyable = tee_fromiterable(it);
836 Py_DECREF(it);
837 if (copyable == NULL) {
838 Py_DECREF(result);
839 return NULL;
840 }
841 } else
842 copyable = it;
843 PyTuple_SET_ITEM(result, 0, copyable);
844 for (i=1 ; i<n ; i++) {
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200845
846 copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000847 if (copyable == NULL) {
848 Py_DECREF(result);
849 return NULL;
850 }
851 PyTuple_SET_ITEM(result, i, copyable);
852 }
853 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000854}
855
856PyDoc_STRVAR(tee_doc,
857"tee(iterable, n=2) --> tuple of n independent iterators.");
858
859
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000860/* cycle object **********************************************************/
861
862typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000863 PyObject_HEAD
864 PyObject *it;
865 PyObject *saved;
866 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000867} cycleobject;
868
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000869static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000870
871static PyObject *
872cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
873{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000874 PyObject *it;
875 PyObject *iterable;
876 PyObject *saved;
877 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000878
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000879 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
880 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000881
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000882 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
883 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000884
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000885 /* Get iterator. */
886 it = PyObject_GetIter(iterable);
887 if (it == NULL)
888 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000889
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000890 saved = PyList_New(0);
891 if (saved == NULL) {
892 Py_DECREF(it);
893 return NULL;
894 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000895
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000896 /* create cycleobject structure */
897 lz = (cycleobject *)type->tp_alloc(type, 0);
898 if (lz == NULL) {
899 Py_DECREF(it);
900 Py_DECREF(saved);
901 return NULL;
902 }
903 lz->it = it;
904 lz->saved = saved;
905 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000906
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000907 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000908}
909
910static void
911cycle_dealloc(cycleobject *lz)
912{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000913 PyObject_GC_UnTrack(lz);
914 Py_XDECREF(lz->saved);
915 Py_XDECREF(lz->it);
916 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000917}
918
919static int
920cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
921{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000922 Py_VISIT(lz->it);
923 Py_VISIT(lz->saved);
924 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000925}
926
927static PyObject *
928cycle_next(cycleobject *lz)
929{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000930 PyObject *item;
931 PyObject *it;
932 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000933
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000934 while (1) {
935 item = PyIter_Next(lz->it);
936 if (item != NULL) {
937 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
938 Py_DECREF(item);
939 return NULL;
940 }
941 return item;
942 }
943 if (PyErr_Occurred()) {
944 if (PyErr_ExceptionMatches(PyExc_StopIteration))
945 PyErr_Clear();
946 else
947 return NULL;
948 }
949 if (PyList_Size(lz->saved) == 0)
950 return NULL;
951 it = PyObject_GetIter(lz->saved);
952 if (it == NULL)
953 return NULL;
954 tmp = lz->it;
955 lz->it = it;
956 lz->firstpass = 1;
957 Py_DECREF(tmp);
958 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000959}
960
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000961static PyObject *
962cycle_reduce(cycleobject *lz)
963{
964 /* Create a new cycle with the iterator tuple, then set
965 * the saved state on it.
966 */
Serhiy Storchakad269b5e2013-01-25 13:38:56 +0200967 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz),
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000968 lz->it, lz->saved, lz->firstpass);
969 }
970
971static PyObject *
972cycle_setstate(cycleobject *lz, PyObject *state)
973{
974 PyObject *saved=NULL;
975 int firstpass;
976 if (!PyArg_ParseTuple(state, "Oi", &saved, &firstpass))
977 return NULL;
978 Py_CLEAR(lz->saved);
979 lz->saved = saved;
980 Py_XINCREF(lz->saved);
981 lz->firstpass = firstpass != 0;
982 Py_RETURN_NONE;
983}
984
985static PyMethodDef cycle_methods[] = {
986 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
987 reduce_doc},
988 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
989 setstate_doc},
990 {NULL, NULL} /* sentinel */
991};
992
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000993PyDoc_STRVAR(cycle_doc,
994"cycle(iterable) --> cycle object\n\
995\n\
996Return elements from the iterable until it is exhausted.\n\
997Then repeat the sequence indefinitely.");
998
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000999static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001000 PyVarObject_HEAD_INIT(NULL, 0)
1001 "itertools.cycle", /* tp_name */
1002 sizeof(cycleobject), /* tp_basicsize */
1003 0, /* tp_itemsize */
1004 /* methods */
1005 (destructor)cycle_dealloc, /* tp_dealloc */
1006 0, /* tp_print */
1007 0, /* tp_getattr */
1008 0, /* tp_setattr */
1009 0, /* tp_reserved */
1010 0, /* tp_repr */
1011 0, /* tp_as_number */
1012 0, /* tp_as_sequence */
1013 0, /* tp_as_mapping */
1014 0, /* tp_hash */
1015 0, /* tp_call */
1016 0, /* tp_str */
1017 PyObject_GenericGetAttr, /* tp_getattro */
1018 0, /* tp_setattro */
1019 0, /* tp_as_buffer */
1020 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1021 Py_TPFLAGS_BASETYPE, /* tp_flags */
1022 cycle_doc, /* tp_doc */
1023 (traverseproc)cycle_traverse, /* tp_traverse */
1024 0, /* tp_clear */
1025 0, /* tp_richcompare */
1026 0, /* tp_weaklistoffset */
1027 PyObject_SelfIter, /* tp_iter */
1028 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001029 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001030 0, /* tp_members */
1031 0, /* tp_getset */
1032 0, /* tp_base */
1033 0, /* tp_dict */
1034 0, /* tp_descr_get */
1035 0, /* tp_descr_set */
1036 0, /* tp_dictoffset */
1037 0, /* tp_init */
1038 0, /* tp_alloc */
1039 cycle_new, /* tp_new */
1040 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001041};
1042
1043
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001044/* dropwhile object **********************************************************/
1045
1046typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001047 PyObject_HEAD
1048 PyObject *func;
1049 PyObject *it;
1050 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001051} dropwhileobject;
1052
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001053static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001054
1055static PyObject *
1056dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1057{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001058 PyObject *func, *seq;
1059 PyObject *it;
1060 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001061
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001062 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
1063 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001064
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001065 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
1066 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001067
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001068 /* Get iterator. */
1069 it = PyObject_GetIter(seq);
1070 if (it == NULL)
1071 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001072
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001073 /* create dropwhileobject structure */
1074 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1075 if (lz == NULL) {
1076 Py_DECREF(it);
1077 return NULL;
1078 }
1079 Py_INCREF(func);
1080 lz->func = func;
1081 lz->it = it;
1082 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001083
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001084 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001085}
1086
1087static void
1088dropwhile_dealloc(dropwhileobject *lz)
1089{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001090 PyObject_GC_UnTrack(lz);
1091 Py_XDECREF(lz->func);
1092 Py_XDECREF(lz->it);
1093 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001094}
1095
1096static int
1097dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1098{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001099 Py_VISIT(lz->it);
1100 Py_VISIT(lz->func);
1101 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001102}
1103
1104static PyObject *
1105dropwhile_next(dropwhileobject *lz)
1106{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001107 PyObject *item, *good;
1108 PyObject *it = lz->it;
1109 long ok;
1110 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001111
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001112 iternext = *Py_TYPE(it)->tp_iternext;
1113 for (;;) {
1114 item = iternext(it);
1115 if (item == NULL)
1116 return NULL;
1117 if (lz->start == 1)
1118 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001119
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001120 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1121 if (good == NULL) {
1122 Py_DECREF(item);
1123 return NULL;
1124 }
1125 ok = PyObject_IsTrue(good);
1126 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001127 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001128 lz->start = 1;
1129 return item;
1130 }
1131 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001132 if (ok < 0)
1133 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001134 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001135}
1136
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001137static PyObject *
1138dropwhile_reduce(dropwhileobject *lz)
1139{
1140 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1141 lz->func, lz->it, lz->start);
1142}
1143
1144static PyObject *
1145dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1146{
1147 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001148 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001149 return NULL;
1150 lz->start = start;
1151 Py_RETURN_NONE;
1152}
1153
1154static PyMethodDef dropwhile_methods[] = {
1155 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1156 reduce_doc},
1157 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1158 setstate_doc},
1159 {NULL, NULL} /* sentinel */
1160};
1161
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001162PyDoc_STRVAR(dropwhile_doc,
1163"dropwhile(predicate, iterable) --> dropwhile object\n\
1164\n\
1165Drop items from the iterable while predicate(item) is true.\n\
1166Afterwards, return every element until the iterable is exhausted.");
1167
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001168static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001169 PyVarObject_HEAD_INIT(NULL, 0)
1170 "itertools.dropwhile", /* tp_name */
1171 sizeof(dropwhileobject), /* tp_basicsize */
1172 0, /* tp_itemsize */
1173 /* methods */
1174 (destructor)dropwhile_dealloc, /* tp_dealloc */
1175 0, /* tp_print */
1176 0, /* tp_getattr */
1177 0, /* tp_setattr */
1178 0, /* tp_reserved */
1179 0, /* tp_repr */
1180 0, /* tp_as_number */
1181 0, /* tp_as_sequence */
1182 0, /* tp_as_mapping */
1183 0, /* tp_hash */
1184 0, /* tp_call */
1185 0, /* tp_str */
1186 PyObject_GenericGetAttr, /* tp_getattro */
1187 0, /* tp_setattro */
1188 0, /* tp_as_buffer */
1189 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1190 Py_TPFLAGS_BASETYPE, /* tp_flags */
1191 dropwhile_doc, /* tp_doc */
1192 (traverseproc)dropwhile_traverse, /* tp_traverse */
1193 0, /* tp_clear */
1194 0, /* tp_richcompare */
1195 0, /* tp_weaklistoffset */
1196 PyObject_SelfIter, /* tp_iter */
1197 (iternextfunc)dropwhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001198 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001199 0, /* tp_members */
1200 0, /* tp_getset */
1201 0, /* tp_base */
1202 0, /* tp_dict */
1203 0, /* tp_descr_get */
1204 0, /* tp_descr_set */
1205 0, /* tp_dictoffset */
1206 0, /* tp_init */
1207 0, /* tp_alloc */
1208 dropwhile_new, /* tp_new */
1209 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001210};
1211
1212
1213/* takewhile object **********************************************************/
1214
1215typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001216 PyObject_HEAD
1217 PyObject *func;
1218 PyObject *it;
1219 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001220} takewhileobject;
1221
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001222static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001223
1224static PyObject *
1225takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1226{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001227 PyObject *func, *seq;
1228 PyObject *it;
1229 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001230
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001231 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1232 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001233
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001234 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1235 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001236
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001237 /* Get iterator. */
1238 it = PyObject_GetIter(seq);
1239 if (it == NULL)
1240 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001241
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001242 /* create takewhileobject structure */
1243 lz = (takewhileobject *)type->tp_alloc(type, 0);
1244 if (lz == NULL) {
1245 Py_DECREF(it);
1246 return NULL;
1247 }
1248 Py_INCREF(func);
1249 lz->func = func;
1250 lz->it = it;
1251 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001252
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001253 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001254}
1255
1256static void
1257takewhile_dealloc(takewhileobject *lz)
1258{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001259 PyObject_GC_UnTrack(lz);
1260 Py_XDECREF(lz->func);
1261 Py_XDECREF(lz->it);
1262 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001263}
1264
1265static int
1266takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1267{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001268 Py_VISIT(lz->it);
1269 Py_VISIT(lz->func);
1270 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001271}
1272
1273static PyObject *
1274takewhile_next(takewhileobject *lz)
1275{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001276 PyObject *item, *good;
1277 PyObject *it = lz->it;
1278 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001279
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001280 if (lz->stop == 1)
1281 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001282
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001283 item = (*Py_TYPE(it)->tp_iternext)(it);
1284 if (item == NULL)
1285 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001286
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001287 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1288 if (good == NULL) {
1289 Py_DECREF(item);
1290 return NULL;
1291 }
1292 ok = PyObject_IsTrue(good);
1293 Py_DECREF(good);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001294 if (ok == 1)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001295 return item;
1296 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001297 if (ok == 0)
1298 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001299 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001300}
1301
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001302static PyObject *
1303takewhile_reduce(takewhileobject *lz)
1304{
1305 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1306 lz->func, lz->it, lz->stop);
1307}
1308
1309static PyObject *
1310takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1311{
1312 int stop = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001313 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001314 return NULL;
1315 lz->stop = stop;
1316 Py_RETURN_NONE;
1317}
1318
1319static PyMethodDef takewhile_reduce_methods[] = {
1320 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1321 reduce_doc},
1322 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1323 setstate_doc},
1324 {NULL, NULL} /* sentinel */
1325};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001326PyDoc_STRVAR(takewhile_doc,
1327"takewhile(predicate, iterable) --> takewhile object\n\
1328\n\
1329Return successive entries from an iterable as long as the \n\
1330predicate evaluates to true for each entry.");
1331
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001332static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001333 PyVarObject_HEAD_INIT(NULL, 0)
1334 "itertools.takewhile", /* tp_name */
1335 sizeof(takewhileobject), /* tp_basicsize */
1336 0, /* tp_itemsize */
1337 /* methods */
1338 (destructor)takewhile_dealloc, /* tp_dealloc */
1339 0, /* tp_print */
1340 0, /* tp_getattr */
1341 0, /* tp_setattr */
1342 0, /* tp_reserved */
1343 0, /* tp_repr */
1344 0, /* tp_as_number */
1345 0, /* tp_as_sequence */
1346 0, /* tp_as_mapping */
1347 0, /* tp_hash */
1348 0, /* tp_call */
1349 0, /* tp_str */
1350 PyObject_GenericGetAttr, /* tp_getattro */
1351 0, /* tp_setattro */
1352 0, /* tp_as_buffer */
1353 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1354 Py_TPFLAGS_BASETYPE, /* tp_flags */
1355 takewhile_doc, /* tp_doc */
1356 (traverseproc)takewhile_traverse, /* tp_traverse */
1357 0, /* tp_clear */
1358 0, /* tp_richcompare */
1359 0, /* tp_weaklistoffset */
1360 PyObject_SelfIter, /* tp_iter */
1361 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001362 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001363 0, /* tp_members */
1364 0, /* tp_getset */
1365 0, /* tp_base */
1366 0, /* tp_dict */
1367 0, /* tp_descr_get */
1368 0, /* tp_descr_set */
1369 0, /* tp_dictoffset */
1370 0, /* tp_init */
1371 0, /* tp_alloc */
1372 takewhile_new, /* tp_new */
1373 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001374};
1375
1376
1377/* islice object ************************************************************/
1378
1379typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001380 PyObject_HEAD
1381 PyObject *it;
1382 Py_ssize_t next;
1383 Py_ssize_t stop;
1384 Py_ssize_t step;
1385 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001386} isliceobject;
1387
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001388static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001389
1390static PyObject *
1391islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1392{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001393 PyObject *seq;
1394 Py_ssize_t start=0, stop=-1, step=1;
1395 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1396 Py_ssize_t numargs;
1397 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001398
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001399 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1400 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001401
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001402 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1403 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001404
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001405 numargs = PyTuple_Size(args);
1406 if (numargs == 2) {
1407 if (a1 != Py_None) {
1408 stop = PyLong_AsSsize_t(a1);
1409 if (stop == -1) {
1410 if (PyErr_Occurred())
1411 PyErr_Clear();
1412 PyErr_SetString(PyExc_ValueError,
1413 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1414 return NULL;
1415 }
1416 }
1417 } else {
1418 if (a1 != Py_None)
1419 start = PyLong_AsSsize_t(a1);
1420 if (start == -1 && PyErr_Occurred())
1421 PyErr_Clear();
1422 if (a2 != Py_None) {
1423 stop = PyLong_AsSsize_t(a2);
1424 if (stop == -1) {
1425 if (PyErr_Occurred())
1426 PyErr_Clear();
1427 PyErr_SetString(PyExc_ValueError,
1428 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1429 return NULL;
1430 }
1431 }
1432 }
1433 if (start<0 || stop<-1) {
1434 PyErr_SetString(PyExc_ValueError,
1435 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1436 return NULL;
1437 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001438
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001439 if (a3 != NULL) {
1440 if (a3 != Py_None)
1441 step = PyLong_AsSsize_t(a3);
1442 if (step == -1 && PyErr_Occurred())
1443 PyErr_Clear();
1444 }
1445 if (step<1) {
1446 PyErr_SetString(PyExc_ValueError,
1447 "Step for islice() must be a positive integer or None.");
1448 return NULL;
1449 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001450
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001451 /* Get iterator. */
1452 it = PyObject_GetIter(seq);
1453 if (it == NULL)
1454 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001455
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001456 /* create isliceobject structure */
1457 lz = (isliceobject *)type->tp_alloc(type, 0);
1458 if (lz == NULL) {
1459 Py_DECREF(it);
1460 return NULL;
1461 }
1462 lz->it = it;
1463 lz->next = start;
1464 lz->stop = stop;
1465 lz->step = step;
1466 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001467
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001468 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469}
1470
1471static void
1472islice_dealloc(isliceobject *lz)
1473{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001474 PyObject_GC_UnTrack(lz);
1475 Py_XDECREF(lz->it);
1476 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001477}
1478
1479static int
1480islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1481{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001482 Py_VISIT(lz->it);
1483 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001484}
1485
1486static PyObject *
1487islice_next(isliceobject *lz)
1488{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001489 PyObject *item;
1490 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001491 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001492 Py_ssize_t oldnext;
1493 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001494
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001495 if (it == NULL)
1496 return NULL;
1497
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001498 iternext = *Py_TYPE(it)->tp_iternext;
1499 while (lz->cnt < lz->next) {
1500 item = iternext(it);
1501 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001502 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001503 Py_DECREF(item);
1504 lz->cnt++;
1505 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001506 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001507 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001508 item = iternext(it);
1509 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001510 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001511 lz->cnt++;
1512 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001513 /* The (size_t) cast below avoids the danger of undefined
1514 behaviour from signed integer overflow. */
1515 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001516 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1517 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001518 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001519
1520empty:
1521 Py_CLEAR(lz->it);
1522 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001523}
1524
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001525static PyObject *
1526islice_reduce(isliceobject *lz)
1527{
1528 /* When unpickled, generate a new object with the same bounds,
1529 * then 'setstate' with the next and count
1530 */
1531 PyObject *stop;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001532 if (lz->it == NULL) {
1533 PyObject *empty_list;
1534 PyObject *empty_it;
1535 empty_list = PyList_New(0);
1536 if (empty_list == NULL)
1537 return NULL;
1538 empty_it = PyObject_GetIter(empty_list);
1539 Py_DECREF(empty_list);
1540 if (empty_it == NULL)
1541 return NULL;
1542 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1543 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001544 if (lz->stop == -1) {
1545 stop = Py_None;
1546 Py_INCREF(stop);
1547 } else {
1548 stop = PyLong_FromSsize_t(lz->stop);
1549 if (stop == NULL)
1550 return NULL;
1551 }
1552 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1553 lz->it, lz->next, stop, lz->step,
1554 lz->cnt);
1555}
1556
1557static PyObject *
1558islice_setstate(isliceobject *lz, PyObject *state)
1559{
1560 Py_ssize_t cnt = PyLong_AsSsize_t(state);
1561 if (cnt == -1 && PyErr_Occurred())
1562 return NULL;
1563 lz->cnt = cnt;
1564 Py_RETURN_NONE;
1565}
1566
1567static PyMethodDef islice_methods[] = {
1568 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1569 reduce_doc},
1570 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1571 setstate_doc},
1572 {NULL, NULL} /* sentinel */
1573};
1574
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001575PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001576"islice(iterable, stop) --> islice object\n\
1577islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001578\n\
1579Return an iterator whose next() method returns selected values from an\n\
1580iterable. If start is specified, will skip all preceding elements;\n\
1581otherwise, start defaults to zero. Step defaults to one. If\n\
1582specified as another value, step determines how many values are \n\
1583skipped between successive calls. Works like a slice() on a list\n\
1584but returns an iterator.");
1585
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001586static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001587 PyVarObject_HEAD_INIT(NULL, 0)
1588 "itertools.islice", /* tp_name */
1589 sizeof(isliceobject), /* tp_basicsize */
1590 0, /* tp_itemsize */
1591 /* methods */
1592 (destructor)islice_dealloc, /* tp_dealloc */
1593 0, /* tp_print */
1594 0, /* tp_getattr */
1595 0, /* tp_setattr */
1596 0, /* tp_reserved */
1597 0, /* tp_repr */
1598 0, /* tp_as_number */
1599 0, /* tp_as_sequence */
1600 0, /* tp_as_mapping */
1601 0, /* tp_hash */
1602 0, /* tp_call */
1603 0, /* tp_str */
1604 PyObject_GenericGetAttr, /* tp_getattro */
1605 0, /* tp_setattro */
1606 0, /* tp_as_buffer */
1607 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1608 Py_TPFLAGS_BASETYPE, /* tp_flags */
1609 islice_doc, /* tp_doc */
1610 (traverseproc)islice_traverse, /* tp_traverse */
1611 0, /* tp_clear */
1612 0, /* tp_richcompare */
1613 0, /* tp_weaklistoffset */
1614 PyObject_SelfIter, /* tp_iter */
1615 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001616 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001617 0, /* tp_members */
1618 0, /* tp_getset */
1619 0, /* tp_base */
1620 0, /* tp_dict */
1621 0, /* tp_descr_get */
1622 0, /* tp_descr_set */
1623 0, /* tp_dictoffset */
1624 0, /* tp_init */
1625 0, /* tp_alloc */
1626 islice_new, /* tp_new */
1627 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001628};
1629
1630
1631/* starmap object ************************************************************/
1632
1633typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001634 PyObject_HEAD
1635 PyObject *func;
1636 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001637} starmapobject;
1638
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001639static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001640
1641static PyObject *
1642starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1643{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001644 PyObject *func, *seq;
1645 PyObject *it;
1646 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001647
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001648 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1649 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001650
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001651 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1652 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001653
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001654 /* Get iterator. */
1655 it = PyObject_GetIter(seq);
1656 if (it == NULL)
1657 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001658
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001659 /* create starmapobject structure */
1660 lz = (starmapobject *)type->tp_alloc(type, 0);
1661 if (lz == NULL) {
1662 Py_DECREF(it);
1663 return NULL;
1664 }
1665 Py_INCREF(func);
1666 lz->func = func;
1667 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001668
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001669 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001670}
1671
1672static void
1673starmap_dealloc(starmapobject *lz)
1674{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001675 PyObject_GC_UnTrack(lz);
1676 Py_XDECREF(lz->func);
1677 Py_XDECREF(lz->it);
1678 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001679}
1680
1681static int
1682starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1683{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001684 Py_VISIT(lz->it);
1685 Py_VISIT(lz->func);
1686 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001687}
1688
1689static PyObject *
1690starmap_next(starmapobject *lz)
1691{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001692 PyObject *args;
1693 PyObject *result;
1694 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001695
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001696 args = (*Py_TYPE(it)->tp_iternext)(it);
1697 if (args == NULL)
1698 return NULL;
1699 if (!PyTuple_CheckExact(args)) {
1700 PyObject *newargs = PySequence_Tuple(args);
1701 Py_DECREF(args);
1702 if (newargs == NULL)
1703 return NULL;
1704 args = newargs;
1705 }
1706 result = PyObject_Call(lz->func, args, NULL);
1707 Py_DECREF(args);
1708 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001709}
1710
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001711static PyObject *
1712starmap_reduce(starmapobject *lz)
1713{
1714 /* Just pickle the iterator */
1715 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1716}
1717
1718static PyMethodDef starmap_methods[] = {
1719 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1720 reduce_doc},
1721 {NULL, NULL} /* sentinel */
1722};
1723
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001724PyDoc_STRVAR(starmap_doc,
1725"starmap(function, sequence) --> starmap object\n\
1726\n\
1727Return an iterator whose values are returned from the function evaluated\n\
1728with a argument tuple taken from the given sequence.");
1729
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001730static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001731 PyVarObject_HEAD_INIT(NULL, 0)
1732 "itertools.starmap", /* tp_name */
1733 sizeof(starmapobject), /* tp_basicsize */
1734 0, /* tp_itemsize */
1735 /* methods */
1736 (destructor)starmap_dealloc, /* tp_dealloc */
1737 0, /* tp_print */
1738 0, /* tp_getattr */
1739 0, /* tp_setattr */
1740 0, /* tp_reserved */
1741 0, /* tp_repr */
1742 0, /* tp_as_number */
1743 0, /* tp_as_sequence */
1744 0, /* tp_as_mapping */
1745 0, /* tp_hash */
1746 0, /* tp_call */
1747 0, /* tp_str */
1748 PyObject_GenericGetAttr, /* tp_getattro */
1749 0, /* tp_setattro */
1750 0, /* tp_as_buffer */
1751 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1752 Py_TPFLAGS_BASETYPE, /* tp_flags */
1753 starmap_doc, /* tp_doc */
1754 (traverseproc)starmap_traverse, /* tp_traverse */
1755 0, /* tp_clear */
1756 0, /* tp_richcompare */
1757 0, /* tp_weaklistoffset */
1758 PyObject_SelfIter, /* tp_iter */
1759 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001760 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001761 0, /* tp_members */
1762 0, /* tp_getset */
1763 0, /* tp_base */
1764 0, /* tp_dict */
1765 0, /* tp_descr_get */
1766 0, /* tp_descr_set */
1767 0, /* tp_dictoffset */
1768 0, /* tp_init */
1769 0, /* tp_alloc */
1770 starmap_new, /* tp_new */
1771 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001772};
1773
1774
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001775/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001776
1777typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001778 PyObject_HEAD
1779 PyObject *source; /* Iterator over input iterables */
1780 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001781} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001782
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001783static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001784
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001785static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001786chain_new_internal(PyTypeObject *type, PyObject *source)
1787{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001788 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001789
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001790 lz = (chainobject *)type->tp_alloc(type, 0);
1791 if (lz == NULL) {
1792 Py_DECREF(source);
1793 return NULL;
1794 }
1795
1796 lz->source = source;
1797 lz->active = NULL;
1798 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001799}
1800
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001801static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001802chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001803{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001804 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001805
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001806 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1807 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001808
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001809 source = PyObject_GetIter(args);
1810 if (source == NULL)
1811 return NULL;
1812
1813 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001814}
1815
1816static PyObject *
1817chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1818{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001819 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001820
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001821 source = PyObject_GetIter(arg);
1822 if (source == NULL)
1823 return NULL;
1824
1825 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001826}
1827
1828static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001829chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001830{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001831 PyObject_GC_UnTrack(lz);
1832 Py_XDECREF(lz->active);
1833 Py_XDECREF(lz->source);
1834 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001835}
1836
Raymond Hettinger2012f172003-02-07 05:32:58 +00001837static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001838chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001839{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001840 Py_VISIT(lz->source);
1841 Py_VISIT(lz->active);
1842 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001843}
1844
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001845static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001846chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001847{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001848 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001849
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001850 if (lz->source == NULL)
1851 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001852
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001853 if (lz->active == NULL) {
1854 PyObject *iterable = PyIter_Next(lz->source);
1855 if (iterable == NULL) {
1856 Py_CLEAR(lz->source);
1857 return NULL; /* no more input sources */
1858 }
1859 lz->active = PyObject_GetIter(iterable);
1860 Py_DECREF(iterable);
1861 if (lz->active == NULL) {
1862 Py_CLEAR(lz->source);
1863 return NULL; /* input not iterable */
1864 }
1865 }
1866 item = PyIter_Next(lz->active);
1867 if (item != NULL)
1868 return item;
1869 if (PyErr_Occurred()) {
1870 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1871 PyErr_Clear();
1872 else
1873 return NULL; /* input raised an exception */
1874 }
1875 Py_CLEAR(lz->active);
1876 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001877}
1878
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001879static PyObject *
1880chain_reduce(chainobject *lz)
1881{
1882 if (lz->source) {
1883 /* we can't pickle function objects (itertools.from_iterable) so
1884 * we must use setstate to replace the iterable. One day we
1885 * will fix pickling of functions
1886 */
1887 if (lz->active) {
1888 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1889 } else {
1890 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1891 }
1892 } else {
1893 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1894 }
1895 return NULL;
1896}
1897
1898static PyObject *
1899chain_setstate(chainobject *lz, PyObject *state)
1900{
1901 PyObject *source, *active=NULL;
1902 if (! PyArg_ParseTuple(state, "O|O", &source, &active))
1903 return NULL;
1904
1905 Py_CLEAR(lz->source);
1906 lz->source = source;
1907 Py_INCREF(lz->source);
1908 Py_CLEAR(lz->active);
1909 lz->active = active;
1910 Py_XINCREF(lz->active);
1911 Py_RETURN_NONE;
1912}
1913
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001914PyDoc_STRVAR(chain_doc,
1915"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001916\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001917Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001918first iterable until it is exhausted, then elements from the next\n\
1919iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001920
Christian Heimesf16baeb2008-02-29 14:57:44 +00001921PyDoc_STRVAR(chain_from_iterable_doc,
1922"chain.from_iterable(iterable) --> chain object\n\
1923\n\
1924Alternate chain() contructor taking a single iterable argument\n\
1925that evaluates lazily.");
1926
1927static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001928 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1929 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001930 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1931 reduce_doc},
1932 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1933 setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001934 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001935};
1936
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001937static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001938 PyVarObject_HEAD_INIT(NULL, 0)
1939 "itertools.chain", /* tp_name */
1940 sizeof(chainobject), /* tp_basicsize */
1941 0, /* tp_itemsize */
1942 /* methods */
1943 (destructor)chain_dealloc, /* tp_dealloc */
1944 0, /* tp_print */
1945 0, /* tp_getattr */
1946 0, /* tp_setattr */
1947 0, /* tp_reserved */
1948 0, /* tp_repr */
1949 0, /* tp_as_number */
1950 0, /* tp_as_sequence */
1951 0, /* tp_as_mapping */
1952 0, /* tp_hash */
1953 0, /* tp_call */
1954 0, /* tp_str */
1955 PyObject_GenericGetAttr, /* tp_getattro */
1956 0, /* tp_setattro */
1957 0, /* tp_as_buffer */
1958 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1959 Py_TPFLAGS_BASETYPE, /* tp_flags */
1960 chain_doc, /* tp_doc */
1961 (traverseproc)chain_traverse, /* tp_traverse */
1962 0, /* tp_clear */
1963 0, /* tp_richcompare */
1964 0, /* tp_weaklistoffset */
1965 PyObject_SelfIter, /* tp_iter */
1966 (iternextfunc)chain_next, /* tp_iternext */
1967 chain_methods, /* tp_methods */
1968 0, /* tp_members */
1969 0, /* tp_getset */
1970 0, /* tp_base */
1971 0, /* tp_dict */
1972 0, /* tp_descr_get */
1973 0, /* tp_descr_set */
1974 0, /* tp_dictoffset */
1975 0, /* tp_init */
1976 0, /* tp_alloc */
1977 chain_new, /* tp_new */
1978 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001979};
1980
1981
Christian Heimesc3f30c42008-02-22 16:37:40 +00001982/* product object ************************************************************/
1983
1984typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001985 PyObject_HEAD
1986 PyObject *pools; /* tuple of pool tuples */
1987 Py_ssize_t *indices; /* one index per pool */
1988 PyObject *result; /* most recently returned result tuple */
1989 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001990} productobject;
1991
1992static PyTypeObject product_type;
1993
1994static PyObject *
1995product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1996{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001997 productobject *lz;
1998 Py_ssize_t nargs, npools, repeat=1;
1999 PyObject *pools = NULL;
2000 Py_ssize_t *indices = NULL;
2001 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002002
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002003 if (kwds != NULL) {
2004 char *kwlist[] = {"repeat", 0};
2005 PyObject *tmpargs = PyTuple_New(0);
2006 if (tmpargs == NULL)
2007 return NULL;
2008 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
2009 Py_DECREF(tmpargs);
2010 return NULL;
2011 }
2012 Py_DECREF(tmpargs);
2013 if (repeat < 0) {
2014 PyErr_SetString(PyExc_ValueError,
2015 "repeat argument cannot be negative");
2016 return NULL;
2017 }
2018 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002019
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002020 assert(PyTuple_CheckExact(args));
2021 if (repeat == 0) {
2022 nargs = 0;
2023 } else {
2024 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002025 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002026 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2027 return NULL;
2028 }
2029 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002030 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002031
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002032 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002033 if (indices == NULL) {
2034 PyErr_NoMemory();
2035 goto error;
2036 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002037
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002038 pools = PyTuple_New(npools);
2039 if (pools == NULL)
2040 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002041
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002042 for (i=0; i < nargs ; ++i) {
2043 PyObject *item = PyTuple_GET_ITEM(args, i);
2044 PyObject *pool = PySequence_Tuple(item);
2045 if (pool == NULL)
2046 goto error;
2047 PyTuple_SET_ITEM(pools, i, pool);
2048 indices[i] = 0;
2049 }
2050 for ( ; i < npools; ++i) {
2051 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2052 Py_INCREF(pool);
2053 PyTuple_SET_ITEM(pools, i, pool);
2054 indices[i] = 0;
2055 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002056
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002057 /* create productobject structure */
2058 lz = (productobject *)type->tp_alloc(type, 0);
2059 if (lz == NULL)
2060 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002061
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002062 lz->pools = pools;
2063 lz->indices = indices;
2064 lz->result = NULL;
2065 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002066
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002067 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002068
2069error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002070 if (indices != NULL)
2071 PyMem_Free(indices);
2072 Py_XDECREF(pools);
2073 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002074}
2075
2076static void
2077product_dealloc(productobject *lz)
2078{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002079 PyObject_GC_UnTrack(lz);
2080 Py_XDECREF(lz->pools);
2081 Py_XDECREF(lz->result);
2082 if (lz->indices != NULL)
2083 PyMem_Free(lz->indices);
2084 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002085}
2086
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002087static PyObject *
2088product_sizeof(productobject *lz, void *unused)
2089{
2090 Py_ssize_t res;
2091
2092 res = sizeof(productobject);
2093 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2094 return PyLong_FromSsize_t(res);
2095}
2096
2097PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2098
Christian Heimesc3f30c42008-02-22 16:37:40 +00002099static int
2100product_traverse(productobject *lz, visitproc visit, void *arg)
2101{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002102 Py_VISIT(lz->pools);
2103 Py_VISIT(lz->result);
2104 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002105}
2106
2107static PyObject *
2108product_next(productobject *lz)
2109{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002110 PyObject *pool;
2111 PyObject *elem;
2112 PyObject *oldelem;
2113 PyObject *pools = lz->pools;
2114 PyObject *result = lz->result;
2115 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2116 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002117
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002118 if (lz->stopped)
2119 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002120
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002121 if (result == NULL) {
2122 /* On the first pass, return an initial tuple filled with the
2123 first element from each pool. */
2124 result = PyTuple_New(npools);
2125 if (result == NULL)
2126 goto empty;
2127 lz->result = result;
2128 for (i=0; i < npools; i++) {
2129 pool = PyTuple_GET_ITEM(pools, i);
2130 if (PyTuple_GET_SIZE(pool) == 0)
2131 goto empty;
2132 elem = PyTuple_GET_ITEM(pool, 0);
2133 Py_INCREF(elem);
2134 PyTuple_SET_ITEM(result, i, elem);
2135 }
2136 } else {
2137 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002138
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002139 /* Copy the previous result tuple or re-use it if available */
2140 if (Py_REFCNT(result) > 1) {
2141 PyObject *old_result = result;
2142 result = PyTuple_New(npools);
2143 if (result == NULL)
2144 goto empty;
2145 lz->result = result;
2146 for (i=0; i < npools; i++) {
2147 elem = PyTuple_GET_ITEM(old_result, i);
2148 Py_INCREF(elem);
2149 PyTuple_SET_ITEM(result, i, elem);
2150 }
2151 Py_DECREF(old_result);
2152 }
2153 /* Now, we've got the only copy so we can update it in-place */
2154 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002155
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002156 /* Update the pool indices right-to-left. Only advance to the
2157 next pool when the previous one rolls-over */
2158 for (i=npools-1 ; i >= 0 ; i--) {
2159 pool = PyTuple_GET_ITEM(pools, i);
2160 indices[i]++;
2161 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2162 /* Roll-over and advance to next pool */
2163 indices[i] = 0;
2164 elem = PyTuple_GET_ITEM(pool, 0);
2165 Py_INCREF(elem);
2166 oldelem = PyTuple_GET_ITEM(result, i);
2167 PyTuple_SET_ITEM(result, i, elem);
2168 Py_DECREF(oldelem);
2169 } else {
2170 /* No rollover. Just increment and stop here. */
2171 elem = PyTuple_GET_ITEM(pool, indices[i]);
2172 Py_INCREF(elem);
2173 oldelem = PyTuple_GET_ITEM(result, i);
2174 PyTuple_SET_ITEM(result, i, elem);
2175 Py_DECREF(oldelem);
2176 break;
2177 }
2178 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002179
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002180 /* If i is negative, then the indices have all rolled-over
2181 and we're done. */
2182 if (i < 0)
2183 goto empty;
2184 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002185
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002186 Py_INCREF(result);
2187 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002188
2189empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002190 lz->stopped = 1;
2191 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002192}
2193
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002194static PyObject *
2195product_reduce(productobject *lz)
2196{
2197 if (lz->stopped) {
2198 return Py_BuildValue("O(())", Py_TYPE(lz));
2199 } else if (lz->result == NULL) {
2200 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2201 } else {
2202 PyObject *indices;
2203 Py_ssize_t n, i;
2204
2205 /* we must pickle the indices use them for setstate, and
2206 * additionally indicate that the iterator has started
2207 */
2208 n = PyTuple_GET_SIZE(lz->pools);
2209 indices = PyTuple_New(n);
2210 if (indices == NULL)
2211 return NULL;
2212 for (i=0; i<n; i++){
2213 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2214 if (!index) {
2215 Py_DECREF(indices);
2216 return NULL;
2217 }
2218 PyTuple_SET_ITEM(indices, i, index);
2219 }
2220 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2221 }
2222}
2223
2224static PyObject *
2225product_setstate(productobject *lz, PyObject *state)
2226{
2227 PyObject *result;
2228 Py_ssize_t n, i;
2229
2230 n = PyTuple_GET_SIZE(lz->pools);
2231 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2232 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2233 return NULL;
2234 }
2235 for (i=0; i<n; i++)
2236 {
2237 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2238 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2239 if (index < 0 && PyErr_Occurred())
2240 return NULL; /* not an integer */
2241 /* clamp the index */
2242 if (index < 0)
2243 index = 0;
2244 else if (index > n-1)
2245 index = n-1;
2246 lz->indices[i] = index;
2247 }
2248
2249 result = PyTuple_New(n);
2250 if (!result)
2251 return NULL;
2252 for (i=0; i<n; i++) {
2253 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2254 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2255 Py_INCREF(element);
2256 PyTuple_SET_ITEM(result, i, element);
2257 }
2258 Py_CLEAR(lz->result);
2259 lz->result = result;
2260 Py_RETURN_NONE;
2261}
2262
2263static PyMethodDef product_methods[] = {
2264 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2265 reduce_doc},
2266 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2267 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002268 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2269 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002270 {NULL, NULL} /* sentinel */
2271};
2272
Christian Heimesc3f30c42008-02-22 16:37:40 +00002273PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002274"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002275\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002276Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002277For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2278The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2279cycle in a manner similar to an odometer (with the rightmost element changing\n\
2280on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002281To compute the product of an iterable with itself, specify the number\n\
2282of repetitions with the optional repeat keyword argument. For example,\n\
2283product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002284product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2285product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2286
2287static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002288 PyVarObject_HEAD_INIT(NULL, 0)
2289 "itertools.product", /* tp_name */
2290 sizeof(productobject), /* tp_basicsize */
2291 0, /* tp_itemsize */
2292 /* methods */
2293 (destructor)product_dealloc, /* tp_dealloc */
2294 0, /* tp_print */
2295 0, /* tp_getattr */
2296 0, /* tp_setattr */
2297 0, /* tp_reserved */
2298 0, /* tp_repr */
2299 0, /* tp_as_number */
2300 0, /* tp_as_sequence */
2301 0, /* tp_as_mapping */
2302 0, /* tp_hash */
2303 0, /* tp_call */
2304 0, /* tp_str */
2305 PyObject_GenericGetAttr, /* tp_getattro */
2306 0, /* tp_setattro */
2307 0, /* tp_as_buffer */
2308 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2309 Py_TPFLAGS_BASETYPE, /* tp_flags */
2310 product_doc, /* tp_doc */
2311 (traverseproc)product_traverse, /* tp_traverse */
2312 0, /* tp_clear */
2313 0, /* tp_richcompare */
2314 0, /* tp_weaklistoffset */
2315 PyObject_SelfIter, /* tp_iter */
2316 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002317 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002318 0, /* tp_members */
2319 0, /* tp_getset */
2320 0, /* tp_base */
2321 0, /* tp_dict */
2322 0, /* tp_descr_get */
2323 0, /* tp_descr_set */
2324 0, /* tp_dictoffset */
2325 0, /* tp_init */
2326 0, /* tp_alloc */
2327 product_new, /* tp_new */
2328 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002329};
2330
2331
Christian Heimes380f7f22008-02-28 11:19:05 +00002332/* combinations object ************************************************************/
2333
2334typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002335 PyObject_HEAD
2336 PyObject *pool; /* input converted to a tuple */
2337 Py_ssize_t *indices; /* one index per result element */
2338 PyObject *result; /* most recently returned result tuple */
2339 Py_ssize_t r; /* size of result tuple */
2340 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002341} combinationsobject;
2342
2343static PyTypeObject combinations_type;
2344
2345static PyObject *
2346combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2347{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002348 combinationsobject *co;
2349 Py_ssize_t n;
2350 Py_ssize_t r;
2351 PyObject *pool = NULL;
2352 PyObject *iterable = NULL;
2353 Py_ssize_t *indices = NULL;
2354 Py_ssize_t i;
2355 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002356
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002357 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2358 &iterable, &r))
2359 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002360
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002361 pool = PySequence_Tuple(iterable);
2362 if (pool == NULL)
2363 goto error;
2364 n = PyTuple_GET_SIZE(pool);
2365 if (r < 0) {
2366 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2367 goto error;
2368 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002369
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002370 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002371 if (indices == NULL) {
2372 PyErr_NoMemory();
2373 goto error;
2374 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002375
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002376 for (i=0 ; i<r ; i++)
2377 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002378
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002379 /* create combinationsobject structure */
2380 co = (combinationsobject *)type->tp_alloc(type, 0);
2381 if (co == NULL)
2382 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002383
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002384 co->pool = pool;
2385 co->indices = indices;
2386 co->result = NULL;
2387 co->r = r;
2388 co->stopped = r > n ? 1 : 0;
2389
2390 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002391
2392error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002393 if (indices != NULL)
2394 PyMem_Free(indices);
2395 Py_XDECREF(pool);
2396 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002397}
2398
2399static void
2400combinations_dealloc(combinationsobject *co)
2401{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002402 PyObject_GC_UnTrack(co);
2403 Py_XDECREF(co->pool);
2404 Py_XDECREF(co->result);
2405 if (co->indices != NULL)
2406 PyMem_Free(co->indices);
2407 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002408}
2409
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002410static PyObject *
2411combinations_sizeof(combinationsobject *co, void *unused)
2412{
2413 Py_ssize_t res;
2414
2415 res = sizeof(combinationsobject);
2416 res += co->r * sizeof(Py_ssize_t);
2417 return PyLong_FromSsize_t(res);
2418}
2419
Christian Heimes380f7f22008-02-28 11:19:05 +00002420static int
2421combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2422{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002423 Py_VISIT(co->pool);
2424 Py_VISIT(co->result);
2425 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002426}
2427
2428static PyObject *
2429combinations_next(combinationsobject *co)
2430{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002431 PyObject *elem;
2432 PyObject *oldelem;
2433 PyObject *pool = co->pool;
2434 Py_ssize_t *indices = co->indices;
2435 PyObject *result = co->result;
2436 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2437 Py_ssize_t r = co->r;
2438 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002439
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002440 if (co->stopped)
2441 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002442
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002443 if (result == NULL) {
2444 /* On the first pass, initialize result tuple using the indices */
2445 result = PyTuple_New(r);
2446 if (result == NULL)
2447 goto empty;
2448 co->result = result;
2449 for (i=0; i<r ; i++) {
2450 index = indices[i];
2451 elem = PyTuple_GET_ITEM(pool, index);
2452 Py_INCREF(elem);
2453 PyTuple_SET_ITEM(result, i, elem);
2454 }
2455 } else {
2456 /* Copy the previous result tuple or re-use it if available */
2457 if (Py_REFCNT(result) > 1) {
2458 PyObject *old_result = result;
2459 result = PyTuple_New(r);
2460 if (result == NULL)
2461 goto empty;
2462 co->result = result;
2463 for (i=0; i<r ; i++) {
2464 elem = PyTuple_GET_ITEM(old_result, i);
2465 Py_INCREF(elem);
2466 PyTuple_SET_ITEM(result, i, elem);
2467 }
2468 Py_DECREF(old_result);
2469 }
2470 /* Now, we've got the only copy so we can update it in-place
2471 * CPython's empty tuple is a singleton and cached in
2472 * PyTuple's freelist.
2473 */
2474 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002475
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002476 /* Scan indices right-to-left until finding one that is not
2477 at its maximum (i + n - r). */
2478 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2479 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002480
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002481 /* If i is negative, then the indices are all at
2482 their maximum value and we're done. */
2483 if (i < 0)
2484 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002485
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002486 /* Increment the current index which we know is not at its
2487 maximum. Then move back to the right setting each index
2488 to its lowest possible value (one higher than the index
2489 to its left -- this maintains the sort order invariant). */
2490 indices[i]++;
2491 for (j=i+1 ; j<r ; j++)
2492 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002493
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002494 /* Update the result tuple for the new indices
2495 starting with i, the leftmost index that changed */
2496 for ( ; i<r ; i++) {
2497 index = indices[i];
2498 elem = PyTuple_GET_ITEM(pool, index);
2499 Py_INCREF(elem);
2500 oldelem = PyTuple_GET_ITEM(result, i);
2501 PyTuple_SET_ITEM(result, i, elem);
2502 Py_DECREF(oldelem);
2503 }
2504 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002505
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002506 Py_INCREF(result);
2507 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002508
2509empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002510 co->stopped = 1;
2511 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002512}
2513
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002514static PyObject *
2515combinations_reduce(combinationsobject *lz)
2516{
2517 if (lz->result == NULL) {
2518 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2519 } else if (lz->stopped) {
2520 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2521 } else {
2522 PyObject *indices;
2523 Py_ssize_t i;
2524
2525 /* we must pickle the indices and use them for setstate */
2526 indices = PyTuple_New(lz->r);
2527 if (!indices)
2528 return NULL;
2529 for (i=0; i<lz->r; i++)
2530 {
2531 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2532 if (!index) {
2533 Py_DECREF(indices);
2534 return NULL;
2535 }
2536 PyTuple_SET_ITEM(indices, i, index);
2537 }
2538
2539 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2540 }
2541}
2542
2543static PyObject *
2544combinations_setstate(combinationsobject *lz, PyObject *state)
2545{
2546 PyObject *result;
2547 Py_ssize_t i;
2548 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2549
2550 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2551 {
2552 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2553 return NULL;
2554 }
2555
2556 for (i=0; i<lz->r; i++)
2557 {
2558 Py_ssize_t max;
2559 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2560 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2561 if (index == -1 && PyErr_Occurred())
2562 return NULL; /* not an integer */
2563 max = i + n - lz->r;
2564 /* clamp the index (beware of negative max) */
2565 if (index > max)
2566 index = max;
2567 if (index < 0)
2568 index = 0;
2569 lz->indices[i] = index;
2570 }
2571
2572 result = PyTuple_New(lz->r);
2573 if (result == NULL)
2574 return NULL;
2575 for (i=0; i<lz->r; i++) {
2576 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2577 Py_INCREF(element);
2578 PyTuple_SET_ITEM(result, i, element);
2579 }
2580
2581 Py_CLEAR(lz->result);
2582 lz->result = result;
2583 Py_RETURN_NONE;
2584}
2585
2586static PyMethodDef combinations_methods[] = {
2587 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2588 reduce_doc},
2589 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2590 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002591 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2592 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002593 {NULL, NULL} /* sentinel */
2594};
2595
Christian Heimes380f7f22008-02-28 11:19:05 +00002596PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002597"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002598\n\
2599Return successive r-length combinations of elements in the iterable.\n\n\
2600combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2601
2602static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002603 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002604 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002605 sizeof(combinationsobject), /* tp_basicsize */
2606 0, /* tp_itemsize */
2607 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002608 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002609 0, /* tp_print */
2610 0, /* tp_getattr */
2611 0, /* tp_setattr */
2612 0, /* tp_reserved */
2613 0, /* tp_repr */
2614 0, /* tp_as_number */
2615 0, /* tp_as_sequence */
2616 0, /* tp_as_mapping */
2617 0, /* tp_hash */
2618 0, /* tp_call */
2619 0, /* tp_str */
2620 PyObject_GenericGetAttr, /* tp_getattro */
2621 0, /* tp_setattro */
2622 0, /* tp_as_buffer */
2623 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2624 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002625 combinations_doc, /* tp_doc */
2626 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002627 0, /* tp_clear */
2628 0, /* tp_richcompare */
2629 0, /* tp_weaklistoffset */
2630 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002631 (iternextfunc)combinations_next, /* tp_iternext */
2632 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002633 0, /* tp_members */
2634 0, /* tp_getset */
2635 0, /* tp_base */
2636 0, /* tp_dict */
2637 0, /* tp_descr_get */
2638 0, /* tp_descr_set */
2639 0, /* tp_dictoffset */
2640 0, /* tp_init */
2641 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002642 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002643 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002644};
2645
2646
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002647/* combinations with replacement object *******************************************/
2648
2649/* Equivalent to:
2650
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002651 def combinations_with_replacement(iterable, r):
2652 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2653 # number items returned: (n+r-1)! / r! / (n-1)!
2654 pool = tuple(iterable)
2655 n = len(pool)
2656 indices = [0] * r
2657 yield tuple(pool[i] for i in indices)
2658 while 1:
2659 for i in reversed(range(r)):
2660 if indices[i] != n - 1:
2661 break
2662 else:
2663 return
2664 indices[i:] = [indices[i] + 1] * (r - i)
2665 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002666
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002667 def combinations_with_replacement2(iterable, r):
2668 'Alternate version that filters from product()'
2669 pool = tuple(iterable)
2670 n = len(pool)
2671 for indices in product(range(n), repeat=r):
2672 if sorted(indices) == list(indices):
2673 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002674*/
2675typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002676 PyObject_HEAD
2677 PyObject *pool; /* input converted to a tuple */
2678 Py_ssize_t *indices; /* one index per result element */
2679 PyObject *result; /* most recently returned result tuple */
2680 Py_ssize_t r; /* size of result tuple */
2681 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002682} cwrobject;
2683
2684static PyTypeObject cwr_type;
2685
2686static PyObject *
2687cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2688{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002689 cwrobject *co;
2690 Py_ssize_t n;
2691 Py_ssize_t r;
2692 PyObject *pool = NULL;
2693 PyObject *iterable = NULL;
2694 Py_ssize_t *indices = NULL;
2695 Py_ssize_t i;
2696 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002697
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002698 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2699 &iterable, &r))
2700 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002701
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002702 pool = PySequence_Tuple(iterable);
2703 if (pool == NULL)
2704 goto error;
2705 n = PyTuple_GET_SIZE(pool);
2706 if (r < 0) {
2707 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2708 goto error;
2709 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002710
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002711 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002712 if (indices == NULL) {
2713 PyErr_NoMemory();
2714 goto error;
2715 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002716
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002717 for (i=0 ; i<r ; i++)
2718 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002719
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002720 /* create cwrobject structure */
2721 co = (cwrobject *)type->tp_alloc(type, 0);
2722 if (co == NULL)
2723 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002724
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002725 co->pool = pool;
2726 co->indices = indices;
2727 co->result = NULL;
2728 co->r = r;
2729 co->stopped = !n && r;
2730
2731 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002732
2733error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002734 if (indices != NULL)
2735 PyMem_Free(indices);
2736 Py_XDECREF(pool);
2737 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002738}
2739
2740static void
2741cwr_dealloc(cwrobject *co)
2742{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002743 PyObject_GC_UnTrack(co);
2744 Py_XDECREF(co->pool);
2745 Py_XDECREF(co->result);
2746 if (co->indices != NULL)
2747 PyMem_Free(co->indices);
2748 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002749}
2750
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002751static PyObject *
2752cwr_sizeof(cwrobject *co, void *unused)
2753{
2754 Py_ssize_t res;
2755
2756 res = sizeof(cwrobject);
2757 res += co->r * sizeof(Py_ssize_t);
2758 return PyLong_FromSsize_t(res);
2759}
2760
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002761static int
2762cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2763{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002764 Py_VISIT(co->pool);
2765 Py_VISIT(co->result);
2766 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002767}
2768
2769static PyObject *
2770cwr_next(cwrobject *co)
2771{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002772 PyObject *elem;
2773 PyObject *oldelem;
2774 PyObject *pool = co->pool;
2775 Py_ssize_t *indices = co->indices;
2776 PyObject *result = co->result;
2777 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2778 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002779 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002780
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002781 if (co->stopped)
2782 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002783
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002784 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002785 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002786 result = PyTuple_New(r);
2787 if (result == NULL)
2788 goto empty;
2789 co->result = result;
Tim Peters9edb1682013-09-03 11:49:31 -05002790 elem = PyTuple_GET_ITEM(pool, 0);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002791 for (i=0; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002792 assert(indices[i] == 0);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002793 Py_INCREF(elem);
2794 PyTuple_SET_ITEM(result, i, elem);
2795 }
2796 } else {
2797 /* Copy the previous result tuple or re-use it if available */
2798 if (Py_REFCNT(result) > 1) {
2799 PyObject *old_result = result;
2800 result = PyTuple_New(r);
2801 if (result == NULL)
2802 goto empty;
2803 co->result = result;
2804 for (i=0; i<r ; i++) {
2805 elem = PyTuple_GET_ITEM(old_result, i);
2806 Py_INCREF(elem);
2807 PyTuple_SET_ITEM(result, i, elem);
2808 }
2809 Py_DECREF(old_result);
2810 }
2811 /* Now, we've got the only copy so we can update it in-place CPython's
2812 empty tuple is a singleton and cached in PyTuple's freelist. */
2813 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002814
Tim Peters9edb1682013-09-03 11:49:31 -05002815 /* Scan indices right-to-left until finding one that is not
2816 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002817 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2818 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002819
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002820 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002821 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002822 if (i < 0)
2823 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002824
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002825 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002826 maximum. Then set all to the right to the same value. */
2827 index = indices[i] + 1;
2828 assert(index < n);
2829 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002830 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002831 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002832 Py_INCREF(elem);
2833 oldelem = PyTuple_GET_ITEM(result, i);
2834 PyTuple_SET_ITEM(result, i, elem);
2835 Py_DECREF(oldelem);
2836 }
2837 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002838
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002839 Py_INCREF(result);
2840 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002841
2842empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002843 co->stopped = 1;
2844 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002845}
2846
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002847static PyObject *
2848cwr_reduce(cwrobject *lz)
2849{
2850 if (lz->result == NULL) {
2851 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2852 } else if (lz->stopped) {
2853 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2854 } else {
2855 PyObject *indices;
2856 Py_ssize_t i;
2857
2858 /* we must pickle the indices and use them for setstate */
2859 indices = PyTuple_New(lz->r);
2860 if (!indices)
2861 return NULL;
2862 for (i=0; i<lz->r; i++)
2863 {
2864 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2865 if (!index) {
2866 Py_DECREF(indices);
2867 return NULL;
2868 }
2869 PyTuple_SET_ITEM(indices, i, index);
2870 }
2871
2872 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2873 }
2874}
2875
2876static PyObject *
2877cwr_setstate(cwrobject *lz, PyObject *state)
2878{
2879 PyObject *result;
2880 Py_ssize_t n, i;
2881
2882 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2883 {
2884 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2885 return NULL;
2886 }
2887
2888 n = PyTuple_GET_SIZE(lz->pool);
2889 for (i=0; i<lz->r; i++)
2890 {
2891 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2892 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2893 if (index < 0 && PyErr_Occurred())
2894 return NULL; /* not an integer */
2895 /* clamp the index */
2896 if (index < 0)
2897 index = 0;
2898 else if (index > n-1)
2899 index = n-1;
2900 lz->indices[i] = index;
2901 }
2902 result = PyTuple_New(lz->r);
2903 if (result == NULL)
2904 return NULL;
2905 for (i=0; i<lz->r; i++) {
2906 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2907 Py_INCREF(element);
2908 PyTuple_SET_ITEM(result, i, element);
2909 }
2910 Py_CLEAR(lz->result);
2911 lz->result = result;
2912 Py_RETURN_NONE;
2913}
2914
2915static PyMethodDef cwr_methods[] = {
2916 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2917 reduce_doc},
2918 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2919 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002920 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
2921 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002922 {NULL, NULL} /* sentinel */
2923};
2924
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002925PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002926"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002927\n\
2928Return successive r-length combinations of elements in the iterable\n\
2929allowing individual elements to have successive repeats.\n\
2930combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2931
2932static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002933 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002934 "itertools.combinations_with_replacement", /* tp_name */
2935 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002936 0, /* tp_itemsize */
2937 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002938 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002939 0, /* tp_print */
2940 0, /* tp_getattr */
2941 0, /* tp_setattr */
2942 0, /* tp_reserved */
2943 0, /* tp_repr */
2944 0, /* tp_as_number */
2945 0, /* tp_as_sequence */
2946 0, /* tp_as_mapping */
2947 0, /* tp_hash */
2948 0, /* tp_call */
2949 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002950 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002951 0, /* tp_setattro */
2952 0, /* tp_as_buffer */
2953 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002954 Py_TPFLAGS_BASETYPE, /* tp_flags */
2955 cwr_doc, /* tp_doc */
2956 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002957 0, /* tp_clear */
2958 0, /* tp_richcompare */
2959 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002960 PyObject_SelfIter, /* tp_iter */
2961 (iternextfunc)cwr_next, /* tp_iternext */
2962 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002963 0, /* tp_members */
2964 0, /* tp_getset */
2965 0, /* tp_base */
2966 0, /* tp_dict */
2967 0, /* tp_descr_get */
2968 0, /* tp_descr_set */
2969 0, /* tp_dictoffset */
2970 0, /* tp_init */
2971 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002972 cwr_new, /* tp_new */
2973 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002974};
2975
2976
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002977/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002978
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002979def permutations(iterable, r=None):
2980 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2981 pool = tuple(iterable)
2982 n = len(pool)
2983 r = n if r is None else r
2984 indices = range(n)
2985 cycles = range(n-r+1, n+1)[::-1]
2986 yield tuple(pool[i] for i in indices[:r])
2987 while n:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002988 for i in reversed(range(r)):
2989 cycles[i] -= 1
2990 if cycles[i] == 0:
2991 indices[i:] = indices[i+1:] + indices[i:i+1]
2992 cycles[i] = n - i
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002993 else:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002994 j = cycles[i]
2995 indices[i], indices[-j] = indices[-j], indices[i]
2996 yield tuple(pool[i] for i in indices[:r])
2997 break
2998 else:
2999 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003000*/
3001
3002typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003003 PyObject_HEAD
3004 PyObject *pool; /* input converted to a tuple */
3005 Py_ssize_t *indices; /* one index per element in the pool */
3006 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3007 PyObject *result; /* most recently returned result tuple */
3008 Py_ssize_t r; /* size of result tuple */
3009 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003010} permutationsobject;
3011
3012static PyTypeObject permutations_type;
3013
3014static PyObject *
3015permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3016{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003017 permutationsobject *po;
3018 Py_ssize_t n;
3019 Py_ssize_t r;
3020 PyObject *robj = Py_None;
3021 PyObject *pool = NULL;
3022 PyObject *iterable = NULL;
3023 Py_ssize_t *indices = NULL;
3024 Py_ssize_t *cycles = NULL;
3025 Py_ssize_t i;
3026 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003027
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003028 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
3029 &iterable, &robj))
3030 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003031
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003032 pool = PySequence_Tuple(iterable);
3033 if (pool == NULL)
3034 goto error;
3035 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003036
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003037 r = n;
3038 if (robj != Py_None) {
3039 if (!PyLong_Check(robj)) {
3040 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3041 goto error;
3042 }
3043 r = PyLong_AsSsize_t(robj);
3044 if (r == -1 && PyErr_Occurred())
3045 goto error;
3046 }
3047 if (r < 0) {
3048 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3049 goto error;
3050 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003051
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003052 indices = PyMem_New(Py_ssize_t, n);
3053 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003054 if (indices == NULL || cycles == NULL) {
3055 PyErr_NoMemory();
3056 goto error;
3057 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003058
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003059 for (i=0 ; i<n ; i++)
3060 indices[i] = i;
3061 for (i=0 ; i<r ; i++)
3062 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003063
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003064 /* create permutationsobject structure */
3065 po = (permutationsobject *)type->tp_alloc(type, 0);
3066 if (po == NULL)
3067 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003068
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003069 po->pool = pool;
3070 po->indices = indices;
3071 po->cycles = cycles;
3072 po->result = NULL;
3073 po->r = r;
3074 po->stopped = r > n ? 1 : 0;
3075
3076 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003077
3078error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003079 if (indices != NULL)
3080 PyMem_Free(indices);
3081 if (cycles != NULL)
3082 PyMem_Free(cycles);
3083 Py_XDECREF(pool);
3084 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003085}
3086
3087static void
3088permutations_dealloc(permutationsobject *po)
3089{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003090 PyObject_GC_UnTrack(po);
3091 Py_XDECREF(po->pool);
3092 Py_XDECREF(po->result);
3093 PyMem_Free(po->indices);
3094 PyMem_Free(po->cycles);
3095 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003096}
3097
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003098static PyObject *
3099permutations_sizeof(permutationsobject *po, void *unused)
3100{
3101 Py_ssize_t res;
3102
3103 res = sizeof(permutationsobject);
3104 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3105 res += po->r * sizeof(Py_ssize_t);
3106 return PyLong_FromSsize_t(res);
3107}
3108
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003109static int
3110permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3111{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003112 Py_VISIT(po->pool);
3113 Py_VISIT(po->result);
3114 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003115}
3116
3117static PyObject *
3118permutations_next(permutationsobject *po)
3119{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003120 PyObject *elem;
3121 PyObject *oldelem;
3122 PyObject *pool = po->pool;
3123 Py_ssize_t *indices = po->indices;
3124 Py_ssize_t *cycles = po->cycles;
3125 PyObject *result = po->result;
3126 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3127 Py_ssize_t r = po->r;
3128 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003129
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003130 if (po->stopped)
3131 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003132
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003133 if (result == NULL) {
3134 /* On the first pass, initialize result tuple using the indices */
3135 result = PyTuple_New(r);
3136 if (result == NULL)
3137 goto empty;
3138 po->result = result;
3139 for (i=0; i<r ; i++) {
3140 index = indices[i];
3141 elem = PyTuple_GET_ITEM(pool, index);
3142 Py_INCREF(elem);
3143 PyTuple_SET_ITEM(result, i, elem);
3144 }
3145 } else {
3146 if (n == 0)
3147 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003148
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003149 /* Copy the previous result tuple or re-use it if available */
3150 if (Py_REFCNT(result) > 1) {
3151 PyObject *old_result = result;
3152 result = PyTuple_New(r);
3153 if (result == NULL)
3154 goto empty;
3155 po->result = result;
3156 for (i=0; i<r ; i++) {
3157 elem = PyTuple_GET_ITEM(old_result, i);
3158 Py_INCREF(elem);
3159 PyTuple_SET_ITEM(result, i, elem);
3160 }
3161 Py_DECREF(old_result);
3162 }
3163 /* Now, we've got the only copy so we can update it in-place */
3164 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003165
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003166 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3167 for (i=r-1 ; i>=0 ; i--) {
3168 cycles[i] -= 1;
3169 if (cycles[i] == 0) {
3170 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3171 index = indices[i];
3172 for (j=i ; j<n-1 ; j++)
3173 indices[j] = indices[j+1];
3174 indices[n-1] = index;
3175 cycles[i] = n - i;
3176 } else {
3177 j = cycles[i];
3178 index = indices[i];
3179 indices[i] = indices[n-j];
3180 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003181
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003182 for (k=i; k<r ; k++) {
3183 /* start with i, the leftmost element that changed */
3184 /* yield tuple(pool[k] for k in indices[:r]) */
3185 index = indices[k];
3186 elem = PyTuple_GET_ITEM(pool, index);
3187 Py_INCREF(elem);
3188 oldelem = PyTuple_GET_ITEM(result, k);
3189 PyTuple_SET_ITEM(result, k, elem);
3190 Py_DECREF(oldelem);
3191 }
3192 break;
3193 }
3194 }
3195 /* If i is negative, then the cycles have all
3196 rolled-over and we're done. */
3197 if (i < 0)
3198 goto empty;
3199 }
3200 Py_INCREF(result);
3201 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003202
3203empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003204 po->stopped = 1;
3205 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003206}
3207
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003208static PyObject *
3209permutations_reduce(permutationsobject *po)
3210{
3211 if (po->result == NULL) {
3212 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3213 } else if (po->stopped) {
3214 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3215 } else {
3216 PyObject *indices=NULL, *cycles=NULL;
3217 Py_ssize_t n, i;
3218
3219 /* we must pickle the indices and cycles and use them for setstate */
3220 n = PyTuple_GET_SIZE(po->pool);
3221 indices = PyTuple_New(n);
3222 if (indices == NULL)
3223 goto err;
3224 for (i=0; i<n; i++){
3225 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3226 if (!index)
3227 goto err;
3228 PyTuple_SET_ITEM(indices, i, index);
3229 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003230
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003231 cycles = PyTuple_New(po->r);
3232 if (cycles == NULL)
3233 goto err;
3234 for (i=0; i<po->r; i++)
3235 {
3236 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3237 if (!index)
3238 goto err;
3239 PyTuple_SET_ITEM(cycles, i, index);
3240 }
3241 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3242 po->pool, po->r,
3243 indices, cycles);
3244 err:
3245 Py_XDECREF(indices);
3246 Py_XDECREF(cycles);
3247 return NULL;
3248 }
3249}
3250
3251static PyObject *
3252permutations_setstate(permutationsobject *po, PyObject *state)
3253{
3254 PyObject *indices, *cycles, *result;
3255 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003256
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003257 if (!PyArg_ParseTuple(state, "O!O!",
3258 &PyTuple_Type, &indices,
3259 &PyTuple_Type, &cycles))
3260 return NULL;
3261
3262 n = PyTuple_GET_SIZE(po->pool);
3263 if (PyTuple_GET_SIZE(indices) != n ||
3264 PyTuple_GET_SIZE(cycles) != po->r)
3265 {
3266 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3267 return NULL;
3268 }
3269
3270 for (i=0; i<n; i++)
3271 {
3272 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3273 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3274 if (index < 0 && PyErr_Occurred())
3275 return NULL; /* not an integer */
3276 /* clamp the index */
3277 if (index < 0)
3278 index = 0;
3279 else if (index > n-1)
3280 index = n-1;
3281 po->indices[i] = index;
3282 }
3283
3284 for (i=0; i<po->r; i++)
3285 {
3286 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3287 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3288 if (index < 0 && PyErr_Occurred())
3289 return NULL; /* not an integer */
3290 if (index < 1)
3291 index = 1;
3292 else if (index > n-i)
3293 index = n-i;
3294 po->cycles[i] = index;
3295 }
3296 result = PyTuple_New(po->r);
3297 if (result == NULL)
3298 return NULL;
3299 for (i=0; i<po->r; i++) {
3300 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3301 Py_INCREF(element);
3302 PyTuple_SET_ITEM(result, i, element);
3303 }
3304 Py_CLEAR(po->result);
3305 po->result = result;
3306 Py_RETURN_NONE;
3307}
3308
3309static PyMethodDef permuations_methods[] = {
3310 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3311 reduce_doc},
3312 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3313 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003314 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3315 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003316 {NULL, NULL} /* sentinel */
3317};
3318
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003319PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003320"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003321\n\
3322Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003323permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003324
3325static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003326 PyVarObject_HEAD_INIT(NULL, 0)
3327 "itertools.permutations", /* tp_name */
3328 sizeof(permutationsobject), /* tp_basicsize */
3329 0, /* tp_itemsize */
3330 /* methods */
3331 (destructor)permutations_dealloc, /* tp_dealloc */
3332 0, /* tp_print */
3333 0, /* tp_getattr */
3334 0, /* tp_setattr */
3335 0, /* tp_reserved */
3336 0, /* tp_repr */
3337 0, /* tp_as_number */
3338 0, /* tp_as_sequence */
3339 0, /* tp_as_mapping */
3340 0, /* tp_hash */
3341 0, /* tp_call */
3342 0, /* tp_str */
3343 PyObject_GenericGetAttr, /* tp_getattro */
3344 0, /* tp_setattro */
3345 0, /* tp_as_buffer */
3346 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3347 Py_TPFLAGS_BASETYPE, /* tp_flags */
3348 permutations_doc, /* tp_doc */
3349 (traverseproc)permutations_traverse, /* tp_traverse */
3350 0, /* tp_clear */
3351 0, /* tp_richcompare */
3352 0, /* tp_weaklistoffset */
3353 PyObject_SelfIter, /* tp_iter */
3354 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003355 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003356 0, /* tp_members */
3357 0, /* tp_getset */
3358 0, /* tp_base */
3359 0, /* tp_dict */
3360 0, /* tp_descr_get */
3361 0, /* tp_descr_set */
3362 0, /* tp_dictoffset */
3363 0, /* tp_init */
3364 0, /* tp_alloc */
3365 permutations_new, /* tp_new */
3366 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003367};
3368
Raymond Hettinger482ba772010-12-01 22:48:00 +00003369/* accumulate object ************************************************************/
3370
3371typedef struct {
3372 PyObject_HEAD
3373 PyObject *total;
3374 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003375 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003376} accumulateobject;
3377
3378static PyTypeObject accumulate_type;
3379
3380static PyObject *
3381accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3382{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003383 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003384 PyObject *iterable;
3385 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003386 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003387 accumulateobject *lz;
3388
Raymond Hettinger5d446132011-03-27 18:52:10 -07003389 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3390 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003391 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003392
3393 /* Get iterator. */
3394 it = PyObject_GetIter(iterable);
3395 if (it == NULL)
3396 return NULL;
3397
Raymond Hettinger482ba772010-12-01 22:48:00 +00003398 /* create accumulateobject structure */
3399 lz = (accumulateobject *)type->tp_alloc(type, 0);
3400 if (lz == NULL) {
3401 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003402 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003403 }
3404
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003405 if (binop != Py_None) {
3406 Py_XINCREF(binop);
3407 lz->binop = binop;
3408 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003409 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003410 lz->it = it;
3411 return (PyObject *)lz;
3412}
3413
3414static void
3415accumulate_dealloc(accumulateobject *lz)
3416{
3417 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003418 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003419 Py_XDECREF(lz->total);
3420 Py_XDECREF(lz->it);
3421 Py_TYPE(lz)->tp_free(lz);
3422}
3423
3424static int
3425accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3426{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003427 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003428 Py_VISIT(lz->it);
3429 Py_VISIT(lz->total);
3430 return 0;
3431}
3432
3433static PyObject *
3434accumulate_next(accumulateobject *lz)
3435{
3436 PyObject *val, *oldtotal, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003437
Raymond Hettinger482ba772010-12-01 22:48:00 +00003438 val = PyIter_Next(lz->it);
3439 if (val == NULL)
3440 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003441
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003442 if (lz->total == NULL) {
3443 Py_INCREF(val);
3444 lz->total = val;
3445 return lz->total;
3446 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003447
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003448 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003449 newtotal = PyNumber_Add(lz->total, val);
3450 else
3451 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003452 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003453 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003454 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003455
3456 oldtotal = lz->total;
3457 lz->total = newtotal;
3458 Py_DECREF(oldtotal);
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003459
Raymond Hettinger482ba772010-12-01 22:48:00 +00003460 Py_INCREF(newtotal);
3461 return newtotal;
3462}
3463
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003464static PyObject *
3465accumulate_reduce(accumulateobject *lz)
3466{
3467 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3468 lz->it, lz->binop?lz->binop:Py_None,
3469 lz->total?lz->total:Py_None);
3470 }
3471
3472static PyObject *
3473accumulate_setstate(accumulateobject *lz, PyObject *state)
3474{
3475 Py_CLEAR(lz->total);
3476 lz->total = state;
3477 Py_INCREF(lz->total);
3478 Py_RETURN_NONE;
3479}
3480
3481static PyMethodDef accumulate_methods[] = {
3482 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3483 reduce_doc},
3484 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3485 setstate_doc},
3486 {NULL, NULL} /* sentinel */
3487};
3488
Raymond Hettinger482ba772010-12-01 22:48:00 +00003489PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003490"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003491\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003492Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003493
3494static PyTypeObject accumulate_type = {
3495 PyVarObject_HEAD_INIT(NULL, 0)
3496 "itertools.accumulate", /* tp_name */
3497 sizeof(accumulateobject), /* tp_basicsize */
3498 0, /* tp_itemsize */
3499 /* methods */
3500 (destructor)accumulate_dealloc, /* tp_dealloc */
3501 0, /* tp_print */
3502 0, /* tp_getattr */
3503 0, /* tp_setattr */
3504 0, /* tp_reserved */
3505 0, /* tp_repr */
3506 0, /* tp_as_number */
3507 0, /* tp_as_sequence */
3508 0, /* tp_as_mapping */
3509 0, /* tp_hash */
3510 0, /* tp_call */
3511 0, /* tp_str */
3512 PyObject_GenericGetAttr, /* tp_getattro */
3513 0, /* tp_setattro */
3514 0, /* tp_as_buffer */
3515 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3516 Py_TPFLAGS_BASETYPE, /* tp_flags */
3517 accumulate_doc, /* tp_doc */
3518 (traverseproc)accumulate_traverse, /* tp_traverse */
3519 0, /* tp_clear */
3520 0, /* tp_richcompare */
3521 0, /* tp_weaklistoffset */
3522 PyObject_SelfIter, /* tp_iter */
3523 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003524 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003525 0, /* tp_members */
3526 0, /* tp_getset */
3527 0, /* tp_base */
3528 0, /* tp_dict */
3529 0, /* tp_descr_get */
3530 0, /* tp_descr_set */
3531 0, /* tp_dictoffset */
3532 0, /* tp_init */
3533 0, /* tp_alloc */
3534 accumulate_new, /* tp_new */
3535 PyObject_GC_Del, /* tp_free */
3536};
3537
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003538
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003539/* compress object ************************************************************/
3540
3541/* Equivalent to:
3542
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003543 def compress(data, selectors):
3544 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3545 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003546*/
3547
3548typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003549 PyObject_HEAD
3550 PyObject *data;
3551 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003552} compressobject;
3553
3554static PyTypeObject compress_type;
3555
3556static PyObject *
3557compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3558{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003559 PyObject *seq1, *seq2;
3560 PyObject *data=NULL, *selectors=NULL;
3561 compressobject *lz;
3562 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003563
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003564 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3565 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003566
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003567 data = PyObject_GetIter(seq1);
3568 if (data == NULL)
3569 goto fail;
3570 selectors = PyObject_GetIter(seq2);
3571 if (selectors == NULL)
3572 goto fail;
3573
3574 /* create compressobject structure */
3575 lz = (compressobject *)type->tp_alloc(type, 0);
3576 if (lz == NULL)
3577 goto fail;
3578 lz->data = data;
3579 lz->selectors = selectors;
3580 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003581
3582fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003583 Py_XDECREF(data);
3584 Py_XDECREF(selectors);
3585 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003586}
3587
3588static void
3589compress_dealloc(compressobject *lz)
3590{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003591 PyObject_GC_UnTrack(lz);
3592 Py_XDECREF(lz->data);
3593 Py_XDECREF(lz->selectors);
3594 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003595}
3596
3597static int
3598compress_traverse(compressobject *lz, visitproc visit, void *arg)
3599{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003600 Py_VISIT(lz->data);
3601 Py_VISIT(lz->selectors);
3602 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003603}
3604
3605static PyObject *
3606compress_next(compressobject *lz)
3607{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003608 PyObject *data = lz->data, *selectors = lz->selectors;
3609 PyObject *datum, *selector;
3610 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3611 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3612 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003613
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003614 while (1) {
3615 /* Steps: get datum, get selector, evaluate selector.
3616 Order is important (to match the pure python version
3617 in terms of which input gets a chance to raise an
3618 exception first).
3619 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003620
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003621 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003622 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003623 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003624
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003625 selector = selectornext(selectors);
3626 if (selector == NULL) {
3627 Py_DECREF(datum);
3628 return NULL;
3629 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003630
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003631 ok = PyObject_IsTrue(selector);
3632 Py_DECREF(selector);
3633 if (ok == 1)
3634 return datum;
3635 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003636 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003637 return NULL;
3638 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003639}
3640
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003641static PyObject *
3642compress_reduce(compressobject *lz)
3643{
3644 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3645 lz->data, lz->selectors);
3646 }
3647
3648static PyMethodDef compress_methods[] = {
3649 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3650 reduce_doc},
3651 {NULL, NULL} /* sentinel */
3652};
3653
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003654PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003655"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003656\n\
3657Return data elements corresponding to true selector elements.\n\
3658Forms a shorter iterator from selected data elements using the\n\
3659selectors to choose the data elements.");
3660
3661static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003662 PyVarObject_HEAD_INIT(NULL, 0)
3663 "itertools.compress", /* tp_name */
3664 sizeof(compressobject), /* tp_basicsize */
3665 0, /* tp_itemsize */
3666 /* methods */
3667 (destructor)compress_dealloc, /* tp_dealloc */
3668 0, /* tp_print */
3669 0, /* tp_getattr */
3670 0, /* tp_setattr */
3671 0, /* tp_reserved */
3672 0, /* tp_repr */
3673 0, /* tp_as_number */
3674 0, /* tp_as_sequence */
3675 0, /* tp_as_mapping */
3676 0, /* tp_hash */
3677 0, /* tp_call */
3678 0, /* tp_str */
3679 PyObject_GenericGetAttr, /* tp_getattro */
3680 0, /* tp_setattro */
3681 0, /* tp_as_buffer */
3682 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3683 Py_TPFLAGS_BASETYPE, /* tp_flags */
3684 compress_doc, /* tp_doc */
3685 (traverseproc)compress_traverse, /* tp_traverse */
3686 0, /* tp_clear */
3687 0, /* tp_richcompare */
3688 0, /* tp_weaklistoffset */
3689 PyObject_SelfIter, /* tp_iter */
3690 (iternextfunc)compress_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003691 compress_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003692 0, /* tp_members */
3693 0, /* tp_getset */
3694 0, /* tp_base */
3695 0, /* tp_dict */
3696 0, /* tp_descr_get */
3697 0, /* tp_descr_set */
3698 0, /* tp_dictoffset */
3699 0, /* tp_init */
3700 0, /* tp_alloc */
3701 compress_new, /* tp_new */
3702 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003703};
3704
3705
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003706/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003707
3708typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003709 PyObject_HEAD
3710 PyObject *func;
3711 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003712} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003713
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003714static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003715
3716static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003717filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003718{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003719 PyObject *func, *seq;
3720 PyObject *it;
3721 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003722
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003723 if (type == &filterfalse_type &&
3724 !_PyArg_NoKeywords("filterfalse()", kwds))
3725 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003726
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003727 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3728 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003729
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003730 /* Get iterator. */
3731 it = PyObject_GetIter(seq);
3732 if (it == NULL)
3733 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003734
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003735 /* create filterfalseobject structure */
3736 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3737 if (lz == NULL) {
3738 Py_DECREF(it);
3739 return NULL;
3740 }
3741 Py_INCREF(func);
3742 lz->func = func;
3743 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003744
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003745 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003746}
3747
3748static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003749filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003750{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003751 PyObject_GC_UnTrack(lz);
3752 Py_XDECREF(lz->func);
3753 Py_XDECREF(lz->it);
3754 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003755}
3756
3757static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003758filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003759{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003760 Py_VISIT(lz->it);
3761 Py_VISIT(lz->func);
3762 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003763}
3764
3765static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003766filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003767{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003768 PyObject *item;
3769 PyObject *it = lz->it;
3770 long ok;
3771 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003772
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003773 iternext = *Py_TYPE(it)->tp_iternext;
3774 for (;;) {
3775 item = iternext(it);
3776 if (item == NULL)
3777 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003778
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003779 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3780 ok = PyObject_IsTrue(item);
3781 } else {
3782 PyObject *good;
3783 good = PyObject_CallFunctionObjArgs(lz->func,
3784 item, NULL);
3785 if (good == NULL) {
3786 Py_DECREF(item);
3787 return NULL;
3788 }
3789 ok = PyObject_IsTrue(good);
3790 Py_DECREF(good);
3791 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003792 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003793 return item;
3794 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003795 if (ok < 0)
3796 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003797 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003798}
3799
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003800static PyObject *
3801filterfalse_reduce(filterfalseobject *lz)
3802{
3803 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3804 lz->func, lz->it);
3805 }
3806
3807static PyMethodDef filterfalse_methods[] = {
3808 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3809 reduce_doc},
3810 {NULL, NULL} /* sentinel */
3811};
3812
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003813PyDoc_STRVAR(filterfalse_doc,
3814"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003815\n\
3816Return those items of sequence for which function(item) is false.\n\
3817If function is None, return the items that are false.");
3818
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003819static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003820 PyVarObject_HEAD_INIT(NULL, 0)
3821 "itertools.filterfalse", /* tp_name */
3822 sizeof(filterfalseobject), /* tp_basicsize */
3823 0, /* tp_itemsize */
3824 /* methods */
3825 (destructor)filterfalse_dealloc, /* tp_dealloc */
3826 0, /* tp_print */
3827 0, /* tp_getattr */
3828 0, /* tp_setattr */
3829 0, /* tp_reserved */
3830 0, /* tp_repr */
3831 0, /* tp_as_number */
3832 0, /* tp_as_sequence */
3833 0, /* tp_as_mapping */
3834 0, /* tp_hash */
3835 0, /* tp_call */
3836 0, /* tp_str */
3837 PyObject_GenericGetAttr, /* tp_getattro */
3838 0, /* tp_setattro */
3839 0, /* tp_as_buffer */
3840 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3841 Py_TPFLAGS_BASETYPE, /* tp_flags */
3842 filterfalse_doc, /* tp_doc */
3843 (traverseproc)filterfalse_traverse, /* tp_traverse */
3844 0, /* tp_clear */
3845 0, /* tp_richcompare */
3846 0, /* tp_weaklistoffset */
3847 PyObject_SelfIter, /* tp_iter */
3848 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003849 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003850 0, /* tp_members */
3851 0, /* tp_getset */
3852 0, /* tp_base */
3853 0, /* tp_dict */
3854 0, /* tp_descr_get */
3855 0, /* tp_descr_set */
3856 0, /* tp_dictoffset */
3857 0, /* tp_init */
3858 0, /* tp_alloc */
3859 filterfalse_new, /* tp_new */
3860 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003861};
3862
3863
3864/* count object ************************************************************/
3865
3866typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003867 PyObject_HEAD
3868 Py_ssize_t cnt;
3869 PyObject *long_cnt;
3870 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003871} countobject;
3872
Raymond Hettinger30729212009-02-12 06:28:27 +00003873/* Counting logic and invariants:
3874
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003875fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003876
Serhiy Storchaka483405b2015-02-17 10:14:30 +02003877 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003878 Advances with: cnt += 1
3879 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003880
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003881slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003882
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003883 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3884 All counting is done with python objects (no overflows or underflows).
3885 Advances with: long_cnt += long_step
3886 Step may be zero -- effectively a slow version of repeat(cnt).
3887 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003888*/
3889
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003890static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003891
3892static PyObject *
3893count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3894{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003895 countobject *lz;
3896 int slow_mode = 0;
3897 Py_ssize_t cnt = 0;
3898 PyObject *long_cnt = NULL;
3899 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003900 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003901 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003902
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003903 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3904 kwlist, &long_cnt, &long_step))
3905 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003906
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003907 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3908 (long_step != NULL && !PyNumber_Check(long_step))) {
3909 PyErr_SetString(PyExc_TypeError, "a number is required");
3910 return NULL;
3911 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003912
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003913 if (long_cnt != NULL) {
3914 cnt = PyLong_AsSsize_t(long_cnt);
3915 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3916 PyErr_Clear();
3917 slow_mode = 1;
3918 }
3919 Py_INCREF(long_cnt);
3920 } else {
3921 cnt = 0;
3922 long_cnt = PyLong_FromLong(0);
3923 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003924
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003925 /* If not specified, step defaults to 1 */
3926 if (long_step == NULL) {
3927 long_step = PyLong_FromLong(1);
3928 if (long_step == NULL) {
3929 Py_DECREF(long_cnt);
3930 return NULL;
3931 }
3932 } else
3933 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003934
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003935 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003936
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003937 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003938 step = PyLong_AsLong(long_step);
3939 if (step != 1) {
3940 slow_mode = 1;
3941 if (step == -1 && PyErr_Occurred())
3942 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003943 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003944
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003945 if (slow_mode)
3946 cnt = PY_SSIZE_T_MAX;
3947 else
3948 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003949
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003950 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3951 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3952 assert(slow_mode ||
3953 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003954
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003955 /* create countobject structure */
3956 lz = (countobject *)type->tp_alloc(type, 0);
3957 if (lz == NULL) {
3958 Py_XDECREF(long_cnt);
3959 return NULL;
3960 }
3961 lz->cnt = cnt;
3962 lz->long_cnt = long_cnt;
3963 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003964
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003965 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003966}
3967
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003968static void
3969count_dealloc(countobject *lz)
3970{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003971 PyObject_GC_UnTrack(lz);
3972 Py_XDECREF(lz->long_cnt);
3973 Py_XDECREF(lz->long_step);
3974 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003975}
3976
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003977static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003978count_traverse(countobject *lz, visitproc visit, void *arg)
3979{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003980 Py_VISIT(lz->long_cnt);
3981 Py_VISIT(lz->long_step);
3982 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003983}
3984
3985static PyObject *
3986count_nextlong(countobject *lz)
3987{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003988 PyObject *long_cnt;
3989 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003990
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003991 long_cnt = lz->long_cnt;
3992 if (long_cnt == NULL) {
3993 /* Switch to slow_mode */
3994 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3995 if (long_cnt == NULL)
3996 return NULL;
3997 }
3998 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00003999
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004000 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4001 if (stepped_up == NULL)
4002 return NULL;
4003 lz->long_cnt = stepped_up;
4004 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004005}
4006
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004007static PyObject *
4008count_next(countobject *lz)
4009{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004010 if (lz->cnt == PY_SSIZE_T_MAX)
4011 return count_nextlong(lz);
4012 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004013}
4014
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004015static PyObject *
4016count_repr(countobject *lz)
4017{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004018 if (lz->cnt != PY_SSIZE_T_MAX)
4019 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004020
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004021 if (PyLong_Check(lz->long_step)) {
4022 long step = PyLong_AsLong(lz->long_step);
4023 if (step == -1 && PyErr_Occurred()) {
4024 PyErr_Clear();
4025 }
4026 if (step == 1) {
4027 /* Don't display step when it is an integer equal to 1 */
4028 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
4029 }
4030 }
4031 return PyUnicode_FromFormat("count(%R, %R)",
4032 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004033}
4034
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004035static PyObject *
4036count_reduce(countobject *lz)
4037{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004038 if (lz->cnt == PY_SSIZE_T_MAX)
4039 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4040 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004041}
4042
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004043static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004044 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004045 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004046 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004047};
4048
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004049PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00004050 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004051\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004052Return a count object whose .__next__() method returns consecutive values.\n\
4053Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00004054 def count(firstval=0, step=1):\n\
Eli Benderskyc554f722013-09-03 06:37:19 -07004055 x = firstval\n\
4056 while 1:\n\
4057 yield x\n\
4058 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004059
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004060static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004061 PyVarObject_HEAD_INIT(NULL, 0)
4062 "itertools.count", /* tp_name */
4063 sizeof(countobject), /* tp_basicsize */
4064 0, /* tp_itemsize */
4065 /* methods */
4066 (destructor)count_dealloc, /* tp_dealloc */
4067 0, /* tp_print */
4068 0, /* tp_getattr */
4069 0, /* tp_setattr */
4070 0, /* tp_reserved */
4071 (reprfunc)count_repr, /* tp_repr */
4072 0, /* tp_as_number */
4073 0, /* tp_as_sequence */
4074 0, /* tp_as_mapping */
4075 0, /* tp_hash */
4076 0, /* tp_call */
4077 0, /* tp_str */
4078 PyObject_GenericGetAttr, /* tp_getattro */
4079 0, /* tp_setattro */
4080 0, /* tp_as_buffer */
4081 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4082 Py_TPFLAGS_BASETYPE, /* tp_flags */
4083 count_doc, /* tp_doc */
4084 (traverseproc)count_traverse, /* tp_traverse */
4085 0, /* tp_clear */
4086 0, /* tp_richcompare */
4087 0, /* tp_weaklistoffset */
4088 PyObject_SelfIter, /* tp_iter */
4089 (iternextfunc)count_next, /* tp_iternext */
4090 count_methods, /* tp_methods */
4091 0, /* tp_members */
4092 0, /* tp_getset */
4093 0, /* tp_base */
4094 0, /* tp_dict */
4095 0, /* tp_descr_get */
4096 0, /* tp_descr_set */
4097 0, /* tp_dictoffset */
4098 0, /* tp_init */
4099 0, /* tp_alloc */
4100 count_new, /* tp_new */
4101 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004102};
4103
4104
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004105/* repeat object ************************************************************/
4106
4107typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004108 PyObject_HEAD
4109 PyObject *element;
4110 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004111} repeatobject;
4112
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004113static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004114
4115static PyObject *
4116repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4117{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004118 repeatobject *ro;
4119 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004120 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004121 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004122
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004123 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4124 &element, &cnt))
4125 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004126
Raymond Hettinger97d35552014-06-24 21:36:58 -07004127 if (kwds != NULL)
4128 n_kwds = PyDict_Size(kwds);
4129 /* Does user supply times argument? */
4130 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004131 cnt = 0;
4132
4133 ro = (repeatobject *)type->tp_alloc(type, 0);
4134 if (ro == NULL)
4135 return NULL;
4136 Py_INCREF(element);
4137 ro->element = element;
4138 ro->cnt = cnt;
4139 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004140}
4141
4142static void
4143repeat_dealloc(repeatobject *ro)
4144{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004145 PyObject_GC_UnTrack(ro);
4146 Py_XDECREF(ro->element);
4147 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004148}
4149
4150static int
4151repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4152{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004153 Py_VISIT(ro->element);
4154 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004155}
4156
4157static PyObject *
4158repeat_next(repeatobject *ro)
4159{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004160 if (ro->cnt == 0)
4161 return NULL;
4162 if (ro->cnt > 0)
4163 ro->cnt--;
4164 Py_INCREF(ro->element);
4165 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004166}
4167
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004168static PyObject *
4169repeat_repr(repeatobject *ro)
4170{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004171 if (ro->cnt == -1)
4172 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4173 else
4174 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4175}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004176
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004177static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004178repeat_len(repeatobject *ro)
4179{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004180 if (ro->cnt == -1) {
4181 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4182 return NULL;
4183 }
4184 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004185}
4186
Armin Rigof5b3e362006-02-11 21:32:43 +00004187PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004188
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004189static PyObject *
4190repeat_reduce(repeatobject *ro)
4191{
4192 /* unpickle this so that a new repeat iterator is constructed with an
4193 * object, then call __setstate__ on it to set cnt
4194 */
4195 if (ro->cnt >= 0)
4196 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4197 else
4198 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4199}
4200
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004201static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004202 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004203 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004204 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004205};
4206
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004207PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004208"repeat(object [,times]) -> create an iterator which returns the object\n\
4209for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004210endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004211
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004212static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004213 PyVarObject_HEAD_INIT(NULL, 0)
4214 "itertools.repeat", /* tp_name */
4215 sizeof(repeatobject), /* tp_basicsize */
4216 0, /* tp_itemsize */
4217 /* methods */
4218 (destructor)repeat_dealloc, /* tp_dealloc */
4219 0, /* tp_print */
4220 0, /* tp_getattr */
4221 0, /* tp_setattr */
4222 0, /* tp_reserved */
4223 (reprfunc)repeat_repr, /* tp_repr */
4224 0, /* tp_as_number */
4225 0, /* tp_as_sequence */
4226 0, /* tp_as_mapping */
4227 0, /* tp_hash */
4228 0, /* tp_call */
4229 0, /* tp_str */
4230 PyObject_GenericGetAttr, /* tp_getattro */
4231 0, /* tp_setattro */
4232 0, /* tp_as_buffer */
4233 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4234 Py_TPFLAGS_BASETYPE, /* tp_flags */
4235 repeat_doc, /* tp_doc */
4236 (traverseproc)repeat_traverse, /* tp_traverse */
4237 0, /* tp_clear */
4238 0, /* tp_richcompare */
4239 0, /* tp_weaklistoffset */
4240 PyObject_SelfIter, /* tp_iter */
4241 (iternextfunc)repeat_next, /* tp_iternext */
4242 repeat_methods, /* tp_methods */
4243 0, /* tp_members */
4244 0, /* tp_getset */
4245 0, /* tp_base */
4246 0, /* tp_dict */
4247 0, /* tp_descr_get */
4248 0, /* tp_descr_set */
4249 0, /* tp_dictoffset */
4250 0, /* tp_init */
4251 0, /* tp_alloc */
4252 repeat_new, /* tp_new */
4253 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004254};
4255
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004256/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004257
4258#include "Python.h"
4259
4260typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004261 PyObject_HEAD
4262 Py_ssize_t tuplesize;
4263 Py_ssize_t numactive;
4264 PyObject *ittuple; /* tuple of iterators */
4265 PyObject *result;
4266 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004267} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004268
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004269static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004270
4271static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004272zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004273{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004274 ziplongestobject *lz;
4275 Py_ssize_t i;
4276 PyObject *ittuple; /* tuple of iterators */
4277 PyObject *result;
4278 PyObject *fillvalue = Py_None;
4279 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004280
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004281 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
4282 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4283 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
4284 PyErr_SetString(PyExc_TypeError,
4285 "zip_longest() got an unexpected keyword argument");
4286 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004287 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004288 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004289
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004290 /* args must be a tuple */
4291 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00004292
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004293 /* obtain iterators */
4294 ittuple = PyTuple_New(tuplesize);
4295 if (ittuple == NULL)
4296 return NULL;
4297 for (i=0; i < tuplesize; ++i) {
4298 PyObject *item = PyTuple_GET_ITEM(args, i);
4299 PyObject *it = PyObject_GetIter(item);
4300 if (it == NULL) {
4301 if (PyErr_ExceptionMatches(PyExc_TypeError))
4302 PyErr_Format(PyExc_TypeError,
4303 "zip_longest argument #%zd must support iteration",
4304 i+1);
4305 Py_DECREF(ittuple);
4306 return NULL;
4307 }
4308 PyTuple_SET_ITEM(ittuple, i, it);
4309 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004310
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004311 /* create a result holder */
4312 result = PyTuple_New(tuplesize);
4313 if (result == NULL) {
4314 Py_DECREF(ittuple);
4315 return NULL;
4316 }
4317 for (i=0 ; i < tuplesize ; i++) {
4318 Py_INCREF(Py_None);
4319 PyTuple_SET_ITEM(result, i, Py_None);
4320 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004321
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004322 /* create ziplongestobject structure */
4323 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4324 if (lz == NULL) {
4325 Py_DECREF(ittuple);
4326 Py_DECREF(result);
4327 return NULL;
4328 }
4329 lz->ittuple = ittuple;
4330 lz->tuplesize = tuplesize;
4331 lz->numactive = tuplesize;
4332 lz->result = result;
4333 Py_INCREF(fillvalue);
4334 lz->fillvalue = fillvalue;
4335 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004336}
4337
4338static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004339zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004340{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004341 PyObject_GC_UnTrack(lz);
4342 Py_XDECREF(lz->ittuple);
4343 Py_XDECREF(lz->result);
4344 Py_XDECREF(lz->fillvalue);
4345 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004346}
4347
4348static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004349zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004350{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004351 Py_VISIT(lz->ittuple);
4352 Py_VISIT(lz->result);
4353 Py_VISIT(lz->fillvalue);
4354 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004355}
4356
4357static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004358zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004359{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004360 Py_ssize_t i;
4361 Py_ssize_t tuplesize = lz->tuplesize;
4362 PyObject *result = lz->result;
4363 PyObject *it;
4364 PyObject *item;
4365 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004366
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004367 if (tuplesize == 0)
4368 return NULL;
4369 if (lz->numactive == 0)
4370 return NULL;
4371 if (Py_REFCNT(result) == 1) {
4372 Py_INCREF(result);
4373 for (i=0 ; i < tuplesize ; i++) {
4374 it = PyTuple_GET_ITEM(lz->ittuple, i);
4375 if (it == NULL) {
4376 Py_INCREF(lz->fillvalue);
4377 item = lz->fillvalue;
4378 } else {
4379 item = PyIter_Next(it);
4380 if (item == NULL) {
4381 lz->numactive -= 1;
4382 if (lz->numactive == 0 || PyErr_Occurred()) {
4383 lz->numactive = 0;
4384 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004385 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004386 } else {
4387 Py_INCREF(lz->fillvalue);
4388 item = lz->fillvalue;
4389 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4390 Py_DECREF(it);
4391 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004392 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004393 }
4394 olditem = PyTuple_GET_ITEM(result, i);
4395 PyTuple_SET_ITEM(result, i, item);
4396 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004397 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004398 } else {
4399 result = PyTuple_New(tuplesize);
4400 if (result == NULL)
4401 return NULL;
4402 for (i=0 ; i < tuplesize ; i++) {
4403 it = PyTuple_GET_ITEM(lz->ittuple, i);
4404 if (it == NULL) {
4405 Py_INCREF(lz->fillvalue);
4406 item = lz->fillvalue;
4407 } else {
4408 item = PyIter_Next(it);
4409 if (item == NULL) {
4410 lz->numactive -= 1;
4411 if (lz->numactive == 0 || PyErr_Occurred()) {
4412 lz->numactive = 0;
4413 Py_DECREF(result);
4414 return NULL;
4415 } else {
4416 Py_INCREF(lz->fillvalue);
4417 item = lz->fillvalue;
4418 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4419 Py_DECREF(it);
4420 }
4421 }
4422 }
4423 PyTuple_SET_ITEM(result, i, item);
4424 }
4425 }
4426 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004427}
4428
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004429static PyObject *
4430zip_longest_reduce(ziplongestobject *lz)
4431{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004432
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004433 /* Create a new tuple with empty sequences where appropriate to pickle.
4434 * Then use setstate to set the fillvalue
4435 */
4436 int i;
4437 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
4438 if (args == NULL)
4439 return NULL;
4440 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4441 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4442 if (elem == NULL) {
4443 elem = PyTuple_New(0);
4444 if (elem == NULL) {
4445 Py_DECREF(args);
4446 return NULL;
4447 }
4448 } else
4449 Py_INCREF(elem);
4450 PyTuple_SET_ITEM(args, i, elem);
4451 }
4452 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4453}
4454
4455static PyObject *
4456zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4457{
4458 Py_CLEAR(lz->fillvalue);
4459 lz->fillvalue = state;
4460 Py_INCREF(lz->fillvalue);
4461 Py_RETURN_NONE;
4462}
4463
4464static PyMethodDef zip_longest_methods[] = {
4465 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4466 reduce_doc},
4467 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4468 setstate_doc},
4469 {NULL, NULL} /* sentinel */
4470};
4471
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004472PyDoc_STRVAR(zip_longest_doc,
4473"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004474\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004475Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004476the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004477method continues until the longest iterable in the argument sequence\n\
4478is exhausted and then it raises StopIteration. When the shorter iterables\n\
4479are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4480defaults to None or can be specified by a keyword argument.\n\
4481");
4482
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004483static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004484 PyVarObject_HEAD_INIT(NULL, 0)
4485 "itertools.zip_longest", /* tp_name */
4486 sizeof(ziplongestobject), /* tp_basicsize */
4487 0, /* tp_itemsize */
4488 /* methods */
4489 (destructor)zip_longest_dealloc, /* tp_dealloc */
4490 0, /* tp_print */
4491 0, /* tp_getattr */
4492 0, /* tp_setattr */
4493 0, /* tp_reserved */
4494 0, /* tp_repr */
4495 0, /* tp_as_number */
4496 0, /* tp_as_sequence */
4497 0, /* tp_as_mapping */
4498 0, /* tp_hash */
4499 0, /* tp_call */
4500 0, /* tp_str */
4501 PyObject_GenericGetAttr, /* tp_getattro */
4502 0, /* tp_setattro */
4503 0, /* tp_as_buffer */
4504 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4505 Py_TPFLAGS_BASETYPE, /* tp_flags */
4506 zip_longest_doc, /* tp_doc */
4507 (traverseproc)zip_longest_traverse, /* tp_traverse */
4508 0, /* tp_clear */
4509 0, /* tp_richcompare */
4510 0, /* tp_weaklistoffset */
4511 PyObject_SelfIter, /* tp_iter */
4512 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004513 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004514 0, /* tp_members */
4515 0, /* tp_getset */
4516 0, /* tp_base */
4517 0, /* tp_dict */
4518 0, /* tp_descr_get */
4519 0, /* tp_descr_set */
4520 0, /* tp_dictoffset */
4521 0, /* tp_init */
4522 0, /* tp_alloc */
4523 zip_longest_new, /* tp_new */
4524 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004525};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004526
4527/* module level code ********************************************************/
4528
4529PyDoc_STRVAR(module_doc,
4530"Functional tools for creating and using iterators.\n\
4531\n\
4532Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004533count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004534cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004535repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004536\n\
4537Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004538accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004539chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger8df58f72013-09-09 01:29:40 -05004540chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004541compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4542dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4543groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004544filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004545islice(seq, [start,] stop [, step]) --> elements from\n\
4546 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004547starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004548tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004549takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004550zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004551\n\
4552Combinatoric generators:\n\
4553product(p, q, ... [repeat=1]) --> cartesian product\n\
4554permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004555combinations(p, r)\n\
4556combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004557");
4558
4559
Raymond Hettingerad983e72003-11-12 14:32:26 +00004560static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004561 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4562 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004563};
4564
Martin v. Löwis1a214512008-06-11 05:26:20 +00004565
4566static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004567 PyModuleDef_HEAD_INIT,
4568 "itertools",
4569 module_doc,
4570 -1,
4571 module_methods,
4572 NULL,
4573 NULL,
4574 NULL,
4575 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004576};
4577
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004578PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004579PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004580{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004581 int i;
4582 PyObject *m;
4583 char *name;
4584 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004585 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004586 &combinations_type,
4587 &cwr_type,
4588 &cycle_type,
4589 &dropwhile_type,
4590 &takewhile_type,
4591 &islice_type,
4592 &starmap_type,
4593 &chain_type,
4594 &compress_type,
4595 &filterfalse_type,
4596 &count_type,
4597 &ziplongest_type,
4598 &permutations_type,
4599 &product_type,
4600 &repeat_type,
4601 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004602 &_grouper_type,
4603 &tee_type,
4604 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004605 NULL
4606 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004607
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004608 Py_TYPE(&teedataobject_type) = &PyType_Type;
4609 m = PyModule_Create(&itertoolsmodule);
4610 if (m == NULL)
4611 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004612
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004613 for (i=0 ; typelist[i] != NULL ; i++) {
4614 if (PyType_Ready(typelist[i]) < 0)
4615 return NULL;
4616 name = strchr(typelist[i]->tp_name, '.');
4617 assert (name != NULL);
4618 Py_INCREF(typelist[i]);
4619 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4620 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004621
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004622 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004623}