blob: 16ac2fb5f4ec11d6d34527fdae5f9017a28da883 [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);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002239 PyObject* pool;
2240 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002241 if (index < 0 && PyErr_Occurred())
2242 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002243 pool = PyTuple_GET_ITEM(lz->pools, i);
2244 poolsize = PyTuple_GET_SIZE(pool);
2245 if (poolsize == 0) {
2246 lz->stopped = 1;
2247 Py_RETURN_NONE;
2248 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002249 /* clamp the index */
2250 if (index < 0)
2251 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002252 else if (index > poolsize-1)
2253 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002254 lz->indices[i] = index;
2255 }
2256
2257 result = PyTuple_New(n);
2258 if (!result)
2259 return NULL;
2260 for (i=0; i<n; i++) {
2261 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2262 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2263 Py_INCREF(element);
2264 PyTuple_SET_ITEM(result, i, element);
2265 }
2266 Py_CLEAR(lz->result);
2267 lz->result = result;
2268 Py_RETURN_NONE;
2269}
2270
2271static PyMethodDef product_methods[] = {
2272 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2273 reduce_doc},
2274 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2275 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002276 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2277 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002278 {NULL, NULL} /* sentinel */
2279};
2280
Christian Heimesc3f30c42008-02-22 16:37:40 +00002281PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002282"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002283\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002284Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002285For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2286The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2287cycle in a manner similar to an odometer (with the rightmost element changing\n\
2288on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002289To compute the product of an iterable with itself, specify the number\n\
2290of repetitions with the optional repeat keyword argument. For example,\n\
2291product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002292product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2293product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2294
2295static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002296 PyVarObject_HEAD_INIT(NULL, 0)
2297 "itertools.product", /* tp_name */
2298 sizeof(productobject), /* tp_basicsize */
2299 0, /* tp_itemsize */
2300 /* methods */
2301 (destructor)product_dealloc, /* tp_dealloc */
2302 0, /* tp_print */
2303 0, /* tp_getattr */
2304 0, /* tp_setattr */
2305 0, /* tp_reserved */
2306 0, /* tp_repr */
2307 0, /* tp_as_number */
2308 0, /* tp_as_sequence */
2309 0, /* tp_as_mapping */
2310 0, /* tp_hash */
2311 0, /* tp_call */
2312 0, /* tp_str */
2313 PyObject_GenericGetAttr, /* tp_getattro */
2314 0, /* tp_setattro */
2315 0, /* tp_as_buffer */
2316 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2317 Py_TPFLAGS_BASETYPE, /* tp_flags */
2318 product_doc, /* tp_doc */
2319 (traverseproc)product_traverse, /* tp_traverse */
2320 0, /* tp_clear */
2321 0, /* tp_richcompare */
2322 0, /* tp_weaklistoffset */
2323 PyObject_SelfIter, /* tp_iter */
2324 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002325 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002326 0, /* tp_members */
2327 0, /* tp_getset */
2328 0, /* tp_base */
2329 0, /* tp_dict */
2330 0, /* tp_descr_get */
2331 0, /* tp_descr_set */
2332 0, /* tp_dictoffset */
2333 0, /* tp_init */
2334 0, /* tp_alloc */
2335 product_new, /* tp_new */
2336 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002337};
2338
2339
Christian Heimes380f7f22008-02-28 11:19:05 +00002340/* combinations object ************************************************************/
2341
2342typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002343 PyObject_HEAD
2344 PyObject *pool; /* input converted to a tuple */
2345 Py_ssize_t *indices; /* one index per result element */
2346 PyObject *result; /* most recently returned result tuple */
2347 Py_ssize_t r; /* size of result tuple */
2348 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002349} combinationsobject;
2350
2351static PyTypeObject combinations_type;
2352
2353static PyObject *
2354combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2355{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002356 combinationsobject *co;
2357 Py_ssize_t n;
2358 Py_ssize_t r;
2359 PyObject *pool = NULL;
2360 PyObject *iterable = NULL;
2361 Py_ssize_t *indices = NULL;
2362 Py_ssize_t i;
2363 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002364
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002365 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2366 &iterable, &r))
2367 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002368
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002369 pool = PySequence_Tuple(iterable);
2370 if (pool == NULL)
2371 goto error;
2372 n = PyTuple_GET_SIZE(pool);
2373 if (r < 0) {
2374 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2375 goto error;
2376 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002377
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002378 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002379 if (indices == NULL) {
2380 PyErr_NoMemory();
2381 goto error;
2382 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002383
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002384 for (i=0 ; i<r ; i++)
2385 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002386
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002387 /* create combinationsobject structure */
2388 co = (combinationsobject *)type->tp_alloc(type, 0);
2389 if (co == NULL)
2390 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002391
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002392 co->pool = pool;
2393 co->indices = indices;
2394 co->result = NULL;
2395 co->r = r;
2396 co->stopped = r > n ? 1 : 0;
2397
2398 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002399
2400error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002401 if (indices != NULL)
2402 PyMem_Free(indices);
2403 Py_XDECREF(pool);
2404 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002405}
2406
2407static void
2408combinations_dealloc(combinationsobject *co)
2409{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002410 PyObject_GC_UnTrack(co);
2411 Py_XDECREF(co->pool);
2412 Py_XDECREF(co->result);
2413 if (co->indices != NULL)
2414 PyMem_Free(co->indices);
2415 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002416}
2417
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002418static PyObject *
2419combinations_sizeof(combinationsobject *co, void *unused)
2420{
2421 Py_ssize_t res;
2422
2423 res = sizeof(combinationsobject);
2424 res += co->r * sizeof(Py_ssize_t);
2425 return PyLong_FromSsize_t(res);
2426}
2427
Christian Heimes380f7f22008-02-28 11:19:05 +00002428static int
2429combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2430{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002431 Py_VISIT(co->pool);
2432 Py_VISIT(co->result);
2433 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002434}
2435
2436static PyObject *
2437combinations_next(combinationsobject *co)
2438{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002439 PyObject *elem;
2440 PyObject *oldelem;
2441 PyObject *pool = co->pool;
2442 Py_ssize_t *indices = co->indices;
2443 PyObject *result = co->result;
2444 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2445 Py_ssize_t r = co->r;
2446 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002447
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002448 if (co->stopped)
2449 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002450
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002451 if (result == NULL) {
2452 /* On the first pass, initialize result tuple using the indices */
2453 result = PyTuple_New(r);
2454 if (result == NULL)
2455 goto empty;
2456 co->result = result;
2457 for (i=0; i<r ; i++) {
2458 index = indices[i];
2459 elem = PyTuple_GET_ITEM(pool, index);
2460 Py_INCREF(elem);
2461 PyTuple_SET_ITEM(result, i, elem);
2462 }
2463 } else {
2464 /* Copy the previous result tuple or re-use it if available */
2465 if (Py_REFCNT(result) > 1) {
2466 PyObject *old_result = result;
2467 result = PyTuple_New(r);
2468 if (result == NULL)
2469 goto empty;
2470 co->result = result;
2471 for (i=0; i<r ; i++) {
2472 elem = PyTuple_GET_ITEM(old_result, i);
2473 Py_INCREF(elem);
2474 PyTuple_SET_ITEM(result, i, elem);
2475 }
2476 Py_DECREF(old_result);
2477 }
2478 /* Now, we've got the only copy so we can update it in-place
2479 * CPython's empty tuple is a singleton and cached in
2480 * PyTuple's freelist.
2481 */
2482 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002483
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002484 /* Scan indices right-to-left until finding one that is not
2485 at its maximum (i + n - r). */
2486 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2487 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002488
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002489 /* If i is negative, then the indices are all at
2490 their maximum value and we're done. */
2491 if (i < 0)
2492 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002493
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002494 /* Increment the current index which we know is not at its
2495 maximum. Then move back to the right setting each index
2496 to its lowest possible value (one higher than the index
2497 to its left -- this maintains the sort order invariant). */
2498 indices[i]++;
2499 for (j=i+1 ; j<r ; j++)
2500 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002501
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002502 /* Update the result tuple for the new indices
2503 starting with i, the leftmost index that changed */
2504 for ( ; i<r ; i++) {
2505 index = indices[i];
2506 elem = PyTuple_GET_ITEM(pool, index);
2507 Py_INCREF(elem);
2508 oldelem = PyTuple_GET_ITEM(result, i);
2509 PyTuple_SET_ITEM(result, i, elem);
2510 Py_DECREF(oldelem);
2511 }
2512 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002513
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002514 Py_INCREF(result);
2515 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002516
2517empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002518 co->stopped = 1;
2519 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002520}
2521
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002522static PyObject *
2523combinations_reduce(combinationsobject *lz)
2524{
2525 if (lz->result == NULL) {
2526 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2527 } else if (lz->stopped) {
2528 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2529 } else {
2530 PyObject *indices;
2531 Py_ssize_t i;
2532
2533 /* we must pickle the indices and use them for setstate */
2534 indices = PyTuple_New(lz->r);
2535 if (!indices)
2536 return NULL;
2537 for (i=0; i<lz->r; i++)
2538 {
2539 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2540 if (!index) {
2541 Py_DECREF(indices);
2542 return NULL;
2543 }
2544 PyTuple_SET_ITEM(indices, i, index);
2545 }
2546
2547 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2548 }
2549}
2550
2551static PyObject *
2552combinations_setstate(combinationsobject *lz, PyObject *state)
2553{
2554 PyObject *result;
2555 Py_ssize_t i;
2556 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2557
2558 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2559 {
2560 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2561 return NULL;
2562 }
2563
2564 for (i=0; i<lz->r; i++)
2565 {
2566 Py_ssize_t max;
2567 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2568 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2569 if (index == -1 && PyErr_Occurred())
2570 return NULL; /* not an integer */
2571 max = i + n - lz->r;
2572 /* clamp the index (beware of negative max) */
2573 if (index > max)
2574 index = max;
2575 if (index < 0)
2576 index = 0;
2577 lz->indices[i] = index;
2578 }
2579
2580 result = PyTuple_New(lz->r);
2581 if (result == NULL)
2582 return NULL;
2583 for (i=0; i<lz->r; i++) {
2584 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2585 Py_INCREF(element);
2586 PyTuple_SET_ITEM(result, i, element);
2587 }
2588
2589 Py_CLEAR(lz->result);
2590 lz->result = result;
2591 Py_RETURN_NONE;
2592}
2593
2594static PyMethodDef combinations_methods[] = {
2595 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2596 reduce_doc},
2597 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2598 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002599 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2600 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002601 {NULL, NULL} /* sentinel */
2602};
2603
Christian Heimes380f7f22008-02-28 11:19:05 +00002604PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002605"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002606\n\
2607Return successive r-length combinations of elements in the iterable.\n\n\
2608combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2609
2610static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002611 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002612 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002613 sizeof(combinationsobject), /* tp_basicsize */
2614 0, /* tp_itemsize */
2615 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002616 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002617 0, /* tp_print */
2618 0, /* tp_getattr */
2619 0, /* tp_setattr */
2620 0, /* tp_reserved */
2621 0, /* tp_repr */
2622 0, /* tp_as_number */
2623 0, /* tp_as_sequence */
2624 0, /* tp_as_mapping */
2625 0, /* tp_hash */
2626 0, /* tp_call */
2627 0, /* tp_str */
2628 PyObject_GenericGetAttr, /* tp_getattro */
2629 0, /* tp_setattro */
2630 0, /* tp_as_buffer */
2631 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2632 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002633 combinations_doc, /* tp_doc */
2634 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002635 0, /* tp_clear */
2636 0, /* tp_richcompare */
2637 0, /* tp_weaklistoffset */
2638 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002639 (iternextfunc)combinations_next, /* tp_iternext */
2640 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002641 0, /* tp_members */
2642 0, /* tp_getset */
2643 0, /* tp_base */
2644 0, /* tp_dict */
2645 0, /* tp_descr_get */
2646 0, /* tp_descr_set */
2647 0, /* tp_dictoffset */
2648 0, /* tp_init */
2649 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002650 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002651 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002652};
2653
2654
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002655/* combinations with replacement object *******************************************/
2656
2657/* Equivalent to:
2658
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002659 def combinations_with_replacement(iterable, r):
2660 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2661 # number items returned: (n+r-1)! / r! / (n-1)!
2662 pool = tuple(iterable)
2663 n = len(pool)
2664 indices = [0] * r
2665 yield tuple(pool[i] for i in indices)
2666 while 1:
2667 for i in reversed(range(r)):
2668 if indices[i] != n - 1:
2669 break
2670 else:
2671 return
2672 indices[i:] = [indices[i] + 1] * (r - i)
2673 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002674
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002675 def combinations_with_replacement2(iterable, r):
2676 'Alternate version that filters from product()'
2677 pool = tuple(iterable)
2678 n = len(pool)
2679 for indices in product(range(n), repeat=r):
2680 if sorted(indices) == list(indices):
2681 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002682*/
2683typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002684 PyObject_HEAD
2685 PyObject *pool; /* input converted to a tuple */
2686 Py_ssize_t *indices; /* one index per result element */
2687 PyObject *result; /* most recently returned result tuple */
2688 Py_ssize_t r; /* size of result tuple */
2689 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002690} cwrobject;
2691
2692static PyTypeObject cwr_type;
2693
2694static PyObject *
2695cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2696{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002697 cwrobject *co;
2698 Py_ssize_t n;
2699 Py_ssize_t r;
2700 PyObject *pool = NULL;
2701 PyObject *iterable = NULL;
2702 Py_ssize_t *indices = NULL;
2703 Py_ssize_t i;
2704 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002705
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002706 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2707 &iterable, &r))
2708 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002709
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002710 pool = PySequence_Tuple(iterable);
2711 if (pool == NULL)
2712 goto error;
2713 n = PyTuple_GET_SIZE(pool);
2714 if (r < 0) {
2715 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2716 goto error;
2717 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002718
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002719 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002720 if (indices == NULL) {
2721 PyErr_NoMemory();
2722 goto error;
2723 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002724
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002725 for (i=0 ; i<r ; i++)
2726 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002727
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002728 /* create cwrobject structure */
2729 co = (cwrobject *)type->tp_alloc(type, 0);
2730 if (co == NULL)
2731 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002732
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002733 co->pool = pool;
2734 co->indices = indices;
2735 co->result = NULL;
2736 co->r = r;
2737 co->stopped = !n && r;
2738
2739 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002740
2741error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002742 if (indices != NULL)
2743 PyMem_Free(indices);
2744 Py_XDECREF(pool);
2745 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002746}
2747
2748static void
2749cwr_dealloc(cwrobject *co)
2750{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002751 PyObject_GC_UnTrack(co);
2752 Py_XDECREF(co->pool);
2753 Py_XDECREF(co->result);
2754 if (co->indices != NULL)
2755 PyMem_Free(co->indices);
2756 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002757}
2758
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002759static PyObject *
2760cwr_sizeof(cwrobject *co, void *unused)
2761{
2762 Py_ssize_t res;
2763
2764 res = sizeof(cwrobject);
2765 res += co->r * sizeof(Py_ssize_t);
2766 return PyLong_FromSsize_t(res);
2767}
2768
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002769static int
2770cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2771{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002772 Py_VISIT(co->pool);
2773 Py_VISIT(co->result);
2774 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002775}
2776
2777static PyObject *
2778cwr_next(cwrobject *co)
2779{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002780 PyObject *elem;
2781 PyObject *oldelem;
2782 PyObject *pool = co->pool;
2783 Py_ssize_t *indices = co->indices;
2784 PyObject *result = co->result;
2785 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2786 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002787 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002788
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002789 if (co->stopped)
2790 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002791
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002792 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002793 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002794 result = PyTuple_New(r);
2795 if (result == NULL)
2796 goto empty;
2797 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002798 if (n > 0) {
2799 elem = PyTuple_GET_ITEM(pool, 0);
2800 for (i=0; i<r ; i++) {
2801 assert(indices[i] == 0);
2802 Py_INCREF(elem);
2803 PyTuple_SET_ITEM(result, i, elem);
2804 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002805 }
2806 } else {
2807 /* Copy the previous result tuple or re-use it if available */
2808 if (Py_REFCNT(result) > 1) {
2809 PyObject *old_result = result;
2810 result = PyTuple_New(r);
2811 if (result == NULL)
2812 goto empty;
2813 co->result = result;
2814 for (i=0; i<r ; i++) {
2815 elem = PyTuple_GET_ITEM(old_result, i);
2816 Py_INCREF(elem);
2817 PyTuple_SET_ITEM(result, i, elem);
2818 }
2819 Py_DECREF(old_result);
2820 }
2821 /* Now, we've got the only copy so we can update it in-place CPython's
2822 empty tuple is a singleton and cached in PyTuple's freelist. */
2823 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002824
Tim Peters9edb1682013-09-03 11:49:31 -05002825 /* Scan indices right-to-left until finding one that is not
2826 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002827 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2828 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002829
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002830 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002831 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002832 if (i < 0)
2833 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002834
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002835 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002836 maximum. Then set all to the right to the same value. */
2837 index = indices[i] + 1;
2838 assert(index < n);
2839 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002840 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002841 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002842 Py_INCREF(elem);
2843 oldelem = PyTuple_GET_ITEM(result, i);
2844 PyTuple_SET_ITEM(result, i, elem);
2845 Py_DECREF(oldelem);
2846 }
2847 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002848
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002849 Py_INCREF(result);
2850 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002851
2852empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002853 co->stopped = 1;
2854 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002855}
2856
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002857static PyObject *
2858cwr_reduce(cwrobject *lz)
2859{
2860 if (lz->result == NULL) {
2861 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2862 } else if (lz->stopped) {
2863 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2864 } else {
2865 PyObject *indices;
2866 Py_ssize_t i;
2867
2868 /* we must pickle the indices and use them for setstate */
2869 indices = PyTuple_New(lz->r);
2870 if (!indices)
2871 return NULL;
2872 for (i=0; i<lz->r; i++)
2873 {
2874 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2875 if (!index) {
2876 Py_DECREF(indices);
2877 return NULL;
2878 }
2879 PyTuple_SET_ITEM(indices, i, index);
2880 }
2881
2882 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2883 }
2884}
2885
2886static PyObject *
2887cwr_setstate(cwrobject *lz, PyObject *state)
2888{
2889 PyObject *result;
2890 Py_ssize_t n, i;
2891
2892 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2893 {
2894 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2895 return NULL;
2896 }
2897
2898 n = PyTuple_GET_SIZE(lz->pool);
2899 for (i=0; i<lz->r; i++)
2900 {
2901 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2902 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2903 if (index < 0 && PyErr_Occurred())
2904 return NULL; /* not an integer */
2905 /* clamp the index */
2906 if (index < 0)
2907 index = 0;
2908 else if (index > n-1)
2909 index = n-1;
2910 lz->indices[i] = index;
2911 }
2912 result = PyTuple_New(lz->r);
2913 if (result == NULL)
2914 return NULL;
2915 for (i=0; i<lz->r; i++) {
2916 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2917 Py_INCREF(element);
2918 PyTuple_SET_ITEM(result, i, element);
2919 }
2920 Py_CLEAR(lz->result);
2921 lz->result = result;
2922 Py_RETURN_NONE;
2923}
2924
2925static PyMethodDef cwr_methods[] = {
2926 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2927 reduce_doc},
2928 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2929 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002930 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
2931 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002932 {NULL, NULL} /* sentinel */
2933};
2934
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002935PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002936"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002937\n\
2938Return successive r-length combinations of elements in the iterable\n\
2939allowing individual elements to have successive repeats.\n\
2940combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2941
2942static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002943 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002944 "itertools.combinations_with_replacement", /* tp_name */
2945 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002946 0, /* tp_itemsize */
2947 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002948 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002949 0, /* tp_print */
2950 0, /* tp_getattr */
2951 0, /* tp_setattr */
2952 0, /* tp_reserved */
2953 0, /* tp_repr */
2954 0, /* tp_as_number */
2955 0, /* tp_as_sequence */
2956 0, /* tp_as_mapping */
2957 0, /* tp_hash */
2958 0, /* tp_call */
2959 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002960 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002961 0, /* tp_setattro */
2962 0, /* tp_as_buffer */
2963 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002964 Py_TPFLAGS_BASETYPE, /* tp_flags */
2965 cwr_doc, /* tp_doc */
2966 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002967 0, /* tp_clear */
2968 0, /* tp_richcompare */
2969 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002970 PyObject_SelfIter, /* tp_iter */
2971 (iternextfunc)cwr_next, /* tp_iternext */
2972 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002973 0, /* tp_members */
2974 0, /* tp_getset */
2975 0, /* tp_base */
2976 0, /* tp_dict */
2977 0, /* tp_descr_get */
2978 0, /* tp_descr_set */
2979 0, /* tp_dictoffset */
2980 0, /* tp_init */
2981 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002982 cwr_new, /* tp_new */
2983 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002984};
2985
2986
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002987/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002988
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002989def permutations(iterable, r=None):
2990 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2991 pool = tuple(iterable)
2992 n = len(pool)
2993 r = n if r is None else r
2994 indices = range(n)
2995 cycles = range(n-r+1, n+1)[::-1]
2996 yield tuple(pool[i] for i in indices[:r])
2997 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03002998 for i in reversed(range(r)):
2999 cycles[i] -= 1
3000 if cycles[i] == 0:
3001 indices[i:] = indices[i+1:] + indices[i:i+1]
3002 cycles[i] = n - i
3003 else:
3004 j = cycles[i]
3005 indices[i], indices[-j] = indices[-j], indices[i]
3006 yield tuple(pool[i] for i in indices[:r])
3007 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003008 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003009 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003010*/
3011
3012typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003013 PyObject_HEAD
3014 PyObject *pool; /* input converted to a tuple */
3015 Py_ssize_t *indices; /* one index per element in the pool */
3016 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3017 PyObject *result; /* most recently returned result tuple */
3018 Py_ssize_t r; /* size of result tuple */
3019 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003020} permutationsobject;
3021
3022static PyTypeObject permutations_type;
3023
3024static PyObject *
3025permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3026{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003027 permutationsobject *po;
3028 Py_ssize_t n;
3029 Py_ssize_t r;
3030 PyObject *robj = Py_None;
3031 PyObject *pool = NULL;
3032 PyObject *iterable = NULL;
3033 Py_ssize_t *indices = NULL;
3034 Py_ssize_t *cycles = NULL;
3035 Py_ssize_t i;
3036 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003037
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003038 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
3039 &iterable, &robj))
3040 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003041
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003042 pool = PySequence_Tuple(iterable);
3043 if (pool == NULL)
3044 goto error;
3045 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003046
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003047 r = n;
3048 if (robj != Py_None) {
3049 if (!PyLong_Check(robj)) {
3050 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3051 goto error;
3052 }
3053 r = PyLong_AsSsize_t(robj);
3054 if (r == -1 && PyErr_Occurred())
3055 goto error;
3056 }
3057 if (r < 0) {
3058 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3059 goto error;
3060 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003061
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003062 indices = PyMem_New(Py_ssize_t, n);
3063 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003064 if (indices == NULL || cycles == NULL) {
3065 PyErr_NoMemory();
3066 goto error;
3067 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003068
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003069 for (i=0 ; i<n ; i++)
3070 indices[i] = i;
3071 for (i=0 ; i<r ; i++)
3072 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003073
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003074 /* create permutationsobject structure */
3075 po = (permutationsobject *)type->tp_alloc(type, 0);
3076 if (po == NULL)
3077 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003078
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003079 po->pool = pool;
3080 po->indices = indices;
3081 po->cycles = cycles;
3082 po->result = NULL;
3083 po->r = r;
3084 po->stopped = r > n ? 1 : 0;
3085
3086 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003087
3088error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003089 if (indices != NULL)
3090 PyMem_Free(indices);
3091 if (cycles != NULL)
3092 PyMem_Free(cycles);
3093 Py_XDECREF(pool);
3094 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003095}
3096
3097static void
3098permutations_dealloc(permutationsobject *po)
3099{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003100 PyObject_GC_UnTrack(po);
3101 Py_XDECREF(po->pool);
3102 Py_XDECREF(po->result);
3103 PyMem_Free(po->indices);
3104 PyMem_Free(po->cycles);
3105 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003106}
3107
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003108static PyObject *
3109permutations_sizeof(permutationsobject *po, void *unused)
3110{
3111 Py_ssize_t res;
3112
3113 res = sizeof(permutationsobject);
3114 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3115 res += po->r * sizeof(Py_ssize_t);
3116 return PyLong_FromSsize_t(res);
3117}
3118
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003119static int
3120permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3121{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003122 Py_VISIT(po->pool);
3123 Py_VISIT(po->result);
3124 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003125}
3126
3127static PyObject *
3128permutations_next(permutationsobject *po)
3129{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003130 PyObject *elem;
3131 PyObject *oldelem;
3132 PyObject *pool = po->pool;
3133 Py_ssize_t *indices = po->indices;
3134 Py_ssize_t *cycles = po->cycles;
3135 PyObject *result = po->result;
3136 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3137 Py_ssize_t r = po->r;
3138 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003139
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003140 if (po->stopped)
3141 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003142
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003143 if (result == NULL) {
3144 /* On the first pass, initialize result tuple using the indices */
3145 result = PyTuple_New(r);
3146 if (result == NULL)
3147 goto empty;
3148 po->result = result;
3149 for (i=0; i<r ; i++) {
3150 index = indices[i];
3151 elem = PyTuple_GET_ITEM(pool, index);
3152 Py_INCREF(elem);
3153 PyTuple_SET_ITEM(result, i, elem);
3154 }
3155 } else {
3156 if (n == 0)
3157 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003158
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003159 /* Copy the previous result tuple or re-use it if available */
3160 if (Py_REFCNT(result) > 1) {
3161 PyObject *old_result = result;
3162 result = PyTuple_New(r);
3163 if (result == NULL)
3164 goto empty;
3165 po->result = result;
3166 for (i=0; i<r ; i++) {
3167 elem = PyTuple_GET_ITEM(old_result, i);
3168 Py_INCREF(elem);
3169 PyTuple_SET_ITEM(result, i, elem);
3170 }
3171 Py_DECREF(old_result);
3172 }
3173 /* Now, we've got the only copy so we can update it in-place */
3174 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003175
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003176 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3177 for (i=r-1 ; i>=0 ; i--) {
3178 cycles[i] -= 1;
3179 if (cycles[i] == 0) {
3180 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3181 index = indices[i];
3182 for (j=i ; j<n-1 ; j++)
3183 indices[j] = indices[j+1];
3184 indices[n-1] = index;
3185 cycles[i] = n - i;
3186 } else {
3187 j = cycles[i];
3188 index = indices[i];
3189 indices[i] = indices[n-j];
3190 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003191
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003192 for (k=i; k<r ; k++) {
3193 /* start with i, the leftmost element that changed */
3194 /* yield tuple(pool[k] for k in indices[:r]) */
3195 index = indices[k];
3196 elem = PyTuple_GET_ITEM(pool, index);
3197 Py_INCREF(elem);
3198 oldelem = PyTuple_GET_ITEM(result, k);
3199 PyTuple_SET_ITEM(result, k, elem);
3200 Py_DECREF(oldelem);
3201 }
3202 break;
3203 }
3204 }
3205 /* If i is negative, then the cycles have all
3206 rolled-over and we're done. */
3207 if (i < 0)
3208 goto empty;
3209 }
3210 Py_INCREF(result);
3211 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003212
3213empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003214 po->stopped = 1;
3215 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003216}
3217
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003218static PyObject *
3219permutations_reduce(permutationsobject *po)
3220{
3221 if (po->result == NULL) {
3222 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3223 } else if (po->stopped) {
3224 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3225 } else {
3226 PyObject *indices=NULL, *cycles=NULL;
3227 Py_ssize_t n, i;
3228
3229 /* we must pickle the indices and cycles and use them for setstate */
3230 n = PyTuple_GET_SIZE(po->pool);
3231 indices = PyTuple_New(n);
3232 if (indices == NULL)
3233 goto err;
3234 for (i=0; i<n; i++){
3235 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3236 if (!index)
3237 goto err;
3238 PyTuple_SET_ITEM(indices, i, index);
3239 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003240
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003241 cycles = PyTuple_New(po->r);
3242 if (cycles == NULL)
3243 goto err;
3244 for (i=0; i<po->r; i++)
3245 {
3246 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3247 if (!index)
3248 goto err;
3249 PyTuple_SET_ITEM(cycles, i, index);
3250 }
3251 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3252 po->pool, po->r,
3253 indices, cycles);
3254 err:
3255 Py_XDECREF(indices);
3256 Py_XDECREF(cycles);
3257 return NULL;
3258 }
3259}
3260
3261static PyObject *
3262permutations_setstate(permutationsobject *po, PyObject *state)
3263{
3264 PyObject *indices, *cycles, *result;
3265 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003266
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003267 if (!PyArg_ParseTuple(state, "O!O!",
3268 &PyTuple_Type, &indices,
3269 &PyTuple_Type, &cycles))
3270 return NULL;
3271
3272 n = PyTuple_GET_SIZE(po->pool);
3273 if (PyTuple_GET_SIZE(indices) != n ||
3274 PyTuple_GET_SIZE(cycles) != po->r)
3275 {
3276 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3277 return NULL;
3278 }
3279
3280 for (i=0; i<n; i++)
3281 {
3282 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3283 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3284 if (index < 0 && PyErr_Occurred())
3285 return NULL; /* not an integer */
3286 /* clamp the index */
3287 if (index < 0)
3288 index = 0;
3289 else if (index > n-1)
3290 index = n-1;
3291 po->indices[i] = index;
3292 }
3293
3294 for (i=0; i<po->r; i++)
3295 {
3296 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3297 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3298 if (index < 0 && PyErr_Occurred())
3299 return NULL; /* not an integer */
3300 if (index < 1)
3301 index = 1;
3302 else if (index > n-i)
3303 index = n-i;
3304 po->cycles[i] = index;
3305 }
3306 result = PyTuple_New(po->r);
3307 if (result == NULL)
3308 return NULL;
3309 for (i=0; i<po->r; i++) {
3310 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3311 Py_INCREF(element);
3312 PyTuple_SET_ITEM(result, i, element);
3313 }
3314 Py_CLEAR(po->result);
3315 po->result = result;
3316 Py_RETURN_NONE;
3317}
3318
3319static PyMethodDef permuations_methods[] = {
3320 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3321 reduce_doc},
3322 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3323 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003324 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3325 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003326 {NULL, NULL} /* sentinel */
3327};
3328
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003329PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003330"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003331\n\
3332Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003333permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003334
3335static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003336 PyVarObject_HEAD_INIT(NULL, 0)
3337 "itertools.permutations", /* tp_name */
3338 sizeof(permutationsobject), /* tp_basicsize */
3339 0, /* tp_itemsize */
3340 /* methods */
3341 (destructor)permutations_dealloc, /* tp_dealloc */
3342 0, /* tp_print */
3343 0, /* tp_getattr */
3344 0, /* tp_setattr */
3345 0, /* tp_reserved */
3346 0, /* tp_repr */
3347 0, /* tp_as_number */
3348 0, /* tp_as_sequence */
3349 0, /* tp_as_mapping */
3350 0, /* tp_hash */
3351 0, /* tp_call */
3352 0, /* tp_str */
3353 PyObject_GenericGetAttr, /* tp_getattro */
3354 0, /* tp_setattro */
3355 0, /* tp_as_buffer */
3356 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3357 Py_TPFLAGS_BASETYPE, /* tp_flags */
3358 permutations_doc, /* tp_doc */
3359 (traverseproc)permutations_traverse, /* tp_traverse */
3360 0, /* tp_clear */
3361 0, /* tp_richcompare */
3362 0, /* tp_weaklistoffset */
3363 PyObject_SelfIter, /* tp_iter */
3364 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003365 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003366 0, /* tp_members */
3367 0, /* tp_getset */
3368 0, /* tp_base */
3369 0, /* tp_dict */
3370 0, /* tp_descr_get */
3371 0, /* tp_descr_set */
3372 0, /* tp_dictoffset */
3373 0, /* tp_init */
3374 0, /* tp_alloc */
3375 permutations_new, /* tp_new */
3376 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003377};
3378
Raymond Hettinger482ba772010-12-01 22:48:00 +00003379/* accumulate object ************************************************************/
3380
3381typedef struct {
3382 PyObject_HEAD
3383 PyObject *total;
3384 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003385 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003386} accumulateobject;
3387
3388static PyTypeObject accumulate_type;
3389
3390static PyObject *
3391accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3392{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003393 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003394 PyObject *iterable;
3395 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003396 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003397 accumulateobject *lz;
3398
Raymond Hettinger5d446132011-03-27 18:52:10 -07003399 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3400 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003401 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003402
3403 /* Get iterator. */
3404 it = PyObject_GetIter(iterable);
3405 if (it == NULL)
3406 return NULL;
3407
Raymond Hettinger482ba772010-12-01 22:48:00 +00003408 /* create accumulateobject structure */
3409 lz = (accumulateobject *)type->tp_alloc(type, 0);
3410 if (lz == NULL) {
3411 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003412 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003413 }
3414
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003415 if (binop != Py_None) {
3416 Py_XINCREF(binop);
3417 lz->binop = binop;
3418 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003419 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003420 lz->it = it;
3421 return (PyObject *)lz;
3422}
3423
3424static void
3425accumulate_dealloc(accumulateobject *lz)
3426{
3427 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003428 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003429 Py_XDECREF(lz->total);
3430 Py_XDECREF(lz->it);
3431 Py_TYPE(lz)->tp_free(lz);
3432}
3433
3434static int
3435accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3436{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003437 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003438 Py_VISIT(lz->it);
3439 Py_VISIT(lz->total);
3440 return 0;
3441}
3442
3443static PyObject *
3444accumulate_next(accumulateobject *lz)
3445{
3446 PyObject *val, *oldtotal, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003447
Raymond Hettinger482ba772010-12-01 22:48:00 +00003448 val = PyIter_Next(lz->it);
3449 if (val == NULL)
3450 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003451
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003452 if (lz->total == NULL) {
3453 Py_INCREF(val);
3454 lz->total = val;
3455 return lz->total;
3456 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003457
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003458 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003459 newtotal = PyNumber_Add(lz->total, val);
3460 else
3461 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003462 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003463 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003464 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003465
3466 oldtotal = lz->total;
3467 lz->total = newtotal;
3468 Py_DECREF(oldtotal);
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003469
Raymond Hettinger482ba772010-12-01 22:48:00 +00003470 Py_INCREF(newtotal);
3471 return newtotal;
3472}
3473
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003474static PyObject *
3475accumulate_reduce(accumulateobject *lz)
3476{
3477 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3478 lz->it, lz->binop?lz->binop:Py_None,
3479 lz->total?lz->total:Py_None);
3480 }
3481
3482static PyObject *
3483accumulate_setstate(accumulateobject *lz, PyObject *state)
3484{
3485 Py_CLEAR(lz->total);
3486 lz->total = state;
3487 Py_INCREF(lz->total);
3488 Py_RETURN_NONE;
3489}
3490
3491static PyMethodDef accumulate_methods[] = {
3492 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3493 reduce_doc},
3494 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3495 setstate_doc},
3496 {NULL, NULL} /* sentinel */
3497};
3498
Raymond Hettinger482ba772010-12-01 22:48:00 +00003499PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003500"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003501\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003502Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003503
3504static PyTypeObject accumulate_type = {
3505 PyVarObject_HEAD_INIT(NULL, 0)
3506 "itertools.accumulate", /* tp_name */
3507 sizeof(accumulateobject), /* tp_basicsize */
3508 0, /* tp_itemsize */
3509 /* methods */
3510 (destructor)accumulate_dealloc, /* tp_dealloc */
3511 0, /* tp_print */
3512 0, /* tp_getattr */
3513 0, /* tp_setattr */
3514 0, /* tp_reserved */
3515 0, /* tp_repr */
3516 0, /* tp_as_number */
3517 0, /* tp_as_sequence */
3518 0, /* tp_as_mapping */
3519 0, /* tp_hash */
3520 0, /* tp_call */
3521 0, /* tp_str */
3522 PyObject_GenericGetAttr, /* tp_getattro */
3523 0, /* tp_setattro */
3524 0, /* tp_as_buffer */
3525 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3526 Py_TPFLAGS_BASETYPE, /* tp_flags */
3527 accumulate_doc, /* tp_doc */
3528 (traverseproc)accumulate_traverse, /* tp_traverse */
3529 0, /* tp_clear */
3530 0, /* tp_richcompare */
3531 0, /* tp_weaklistoffset */
3532 PyObject_SelfIter, /* tp_iter */
3533 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003534 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003535 0, /* tp_members */
3536 0, /* tp_getset */
3537 0, /* tp_base */
3538 0, /* tp_dict */
3539 0, /* tp_descr_get */
3540 0, /* tp_descr_set */
3541 0, /* tp_dictoffset */
3542 0, /* tp_init */
3543 0, /* tp_alloc */
3544 accumulate_new, /* tp_new */
3545 PyObject_GC_Del, /* tp_free */
3546};
3547
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003548
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003549/* compress object ************************************************************/
3550
3551/* Equivalent to:
3552
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003553 def compress(data, selectors):
3554 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3555 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003556*/
3557
3558typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003559 PyObject_HEAD
3560 PyObject *data;
3561 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003562} compressobject;
3563
3564static PyTypeObject compress_type;
3565
3566static PyObject *
3567compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3568{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003569 PyObject *seq1, *seq2;
3570 PyObject *data=NULL, *selectors=NULL;
3571 compressobject *lz;
3572 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003573
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003574 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3575 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003576
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003577 data = PyObject_GetIter(seq1);
3578 if (data == NULL)
3579 goto fail;
3580 selectors = PyObject_GetIter(seq2);
3581 if (selectors == NULL)
3582 goto fail;
3583
3584 /* create compressobject structure */
3585 lz = (compressobject *)type->tp_alloc(type, 0);
3586 if (lz == NULL)
3587 goto fail;
3588 lz->data = data;
3589 lz->selectors = selectors;
3590 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003591
3592fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003593 Py_XDECREF(data);
3594 Py_XDECREF(selectors);
3595 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003596}
3597
3598static void
3599compress_dealloc(compressobject *lz)
3600{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003601 PyObject_GC_UnTrack(lz);
3602 Py_XDECREF(lz->data);
3603 Py_XDECREF(lz->selectors);
3604 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003605}
3606
3607static int
3608compress_traverse(compressobject *lz, visitproc visit, void *arg)
3609{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003610 Py_VISIT(lz->data);
3611 Py_VISIT(lz->selectors);
3612 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003613}
3614
3615static PyObject *
3616compress_next(compressobject *lz)
3617{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003618 PyObject *data = lz->data, *selectors = lz->selectors;
3619 PyObject *datum, *selector;
3620 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3621 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3622 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003623
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003624 while (1) {
3625 /* Steps: get datum, get selector, evaluate selector.
3626 Order is important (to match the pure python version
3627 in terms of which input gets a chance to raise an
3628 exception first).
3629 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003630
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003631 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003632 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003633 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003634
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003635 selector = selectornext(selectors);
3636 if (selector == NULL) {
3637 Py_DECREF(datum);
3638 return NULL;
3639 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003640
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003641 ok = PyObject_IsTrue(selector);
3642 Py_DECREF(selector);
3643 if (ok == 1)
3644 return datum;
3645 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003646 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003647 return NULL;
3648 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003649}
3650
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003651static PyObject *
3652compress_reduce(compressobject *lz)
3653{
3654 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3655 lz->data, lz->selectors);
3656 }
3657
3658static PyMethodDef compress_methods[] = {
3659 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3660 reduce_doc},
3661 {NULL, NULL} /* sentinel */
3662};
3663
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003664PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003665"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003666\n\
3667Return data elements corresponding to true selector elements.\n\
3668Forms a shorter iterator from selected data elements using the\n\
3669selectors to choose the data elements.");
3670
3671static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003672 PyVarObject_HEAD_INIT(NULL, 0)
3673 "itertools.compress", /* tp_name */
3674 sizeof(compressobject), /* tp_basicsize */
3675 0, /* tp_itemsize */
3676 /* methods */
3677 (destructor)compress_dealloc, /* tp_dealloc */
3678 0, /* tp_print */
3679 0, /* tp_getattr */
3680 0, /* tp_setattr */
3681 0, /* tp_reserved */
3682 0, /* tp_repr */
3683 0, /* tp_as_number */
3684 0, /* tp_as_sequence */
3685 0, /* tp_as_mapping */
3686 0, /* tp_hash */
3687 0, /* tp_call */
3688 0, /* tp_str */
3689 PyObject_GenericGetAttr, /* tp_getattro */
3690 0, /* tp_setattro */
3691 0, /* tp_as_buffer */
3692 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3693 Py_TPFLAGS_BASETYPE, /* tp_flags */
3694 compress_doc, /* tp_doc */
3695 (traverseproc)compress_traverse, /* tp_traverse */
3696 0, /* tp_clear */
3697 0, /* tp_richcompare */
3698 0, /* tp_weaklistoffset */
3699 PyObject_SelfIter, /* tp_iter */
3700 (iternextfunc)compress_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003701 compress_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003702 0, /* tp_members */
3703 0, /* tp_getset */
3704 0, /* tp_base */
3705 0, /* tp_dict */
3706 0, /* tp_descr_get */
3707 0, /* tp_descr_set */
3708 0, /* tp_dictoffset */
3709 0, /* tp_init */
3710 0, /* tp_alloc */
3711 compress_new, /* tp_new */
3712 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003713};
3714
3715
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003716/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003717
3718typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003719 PyObject_HEAD
3720 PyObject *func;
3721 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003722} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003723
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003724static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003725
3726static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003727filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003728{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003729 PyObject *func, *seq;
3730 PyObject *it;
3731 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003732
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003733 if (type == &filterfalse_type &&
3734 !_PyArg_NoKeywords("filterfalse()", kwds))
3735 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003736
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003737 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3738 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003739
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003740 /* Get iterator. */
3741 it = PyObject_GetIter(seq);
3742 if (it == NULL)
3743 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003744
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003745 /* create filterfalseobject structure */
3746 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3747 if (lz == NULL) {
3748 Py_DECREF(it);
3749 return NULL;
3750 }
3751 Py_INCREF(func);
3752 lz->func = func;
3753 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003754
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003755 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003756}
3757
3758static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003759filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003760{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003761 PyObject_GC_UnTrack(lz);
3762 Py_XDECREF(lz->func);
3763 Py_XDECREF(lz->it);
3764 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003765}
3766
3767static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003768filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003769{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003770 Py_VISIT(lz->it);
3771 Py_VISIT(lz->func);
3772 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003773}
3774
3775static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003776filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003777{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003778 PyObject *item;
3779 PyObject *it = lz->it;
3780 long ok;
3781 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003782
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003783 iternext = *Py_TYPE(it)->tp_iternext;
3784 for (;;) {
3785 item = iternext(it);
3786 if (item == NULL)
3787 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003788
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003789 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3790 ok = PyObject_IsTrue(item);
3791 } else {
3792 PyObject *good;
3793 good = PyObject_CallFunctionObjArgs(lz->func,
3794 item, NULL);
3795 if (good == NULL) {
3796 Py_DECREF(item);
3797 return NULL;
3798 }
3799 ok = PyObject_IsTrue(good);
3800 Py_DECREF(good);
3801 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003802 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003803 return item;
3804 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003805 if (ok < 0)
3806 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003807 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003808}
3809
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003810static PyObject *
3811filterfalse_reduce(filterfalseobject *lz)
3812{
3813 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3814 lz->func, lz->it);
3815 }
3816
3817static PyMethodDef filterfalse_methods[] = {
3818 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3819 reduce_doc},
3820 {NULL, NULL} /* sentinel */
3821};
3822
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003823PyDoc_STRVAR(filterfalse_doc,
3824"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003825\n\
3826Return those items of sequence for which function(item) is false.\n\
3827If function is None, return the items that are false.");
3828
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003829static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003830 PyVarObject_HEAD_INIT(NULL, 0)
3831 "itertools.filterfalse", /* tp_name */
3832 sizeof(filterfalseobject), /* tp_basicsize */
3833 0, /* tp_itemsize */
3834 /* methods */
3835 (destructor)filterfalse_dealloc, /* tp_dealloc */
3836 0, /* tp_print */
3837 0, /* tp_getattr */
3838 0, /* tp_setattr */
3839 0, /* tp_reserved */
3840 0, /* tp_repr */
3841 0, /* tp_as_number */
3842 0, /* tp_as_sequence */
3843 0, /* tp_as_mapping */
3844 0, /* tp_hash */
3845 0, /* tp_call */
3846 0, /* tp_str */
3847 PyObject_GenericGetAttr, /* tp_getattro */
3848 0, /* tp_setattro */
3849 0, /* tp_as_buffer */
3850 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3851 Py_TPFLAGS_BASETYPE, /* tp_flags */
3852 filterfalse_doc, /* tp_doc */
3853 (traverseproc)filterfalse_traverse, /* tp_traverse */
3854 0, /* tp_clear */
3855 0, /* tp_richcompare */
3856 0, /* tp_weaklistoffset */
3857 PyObject_SelfIter, /* tp_iter */
3858 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003859 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003860 0, /* tp_members */
3861 0, /* tp_getset */
3862 0, /* tp_base */
3863 0, /* tp_dict */
3864 0, /* tp_descr_get */
3865 0, /* tp_descr_set */
3866 0, /* tp_dictoffset */
3867 0, /* tp_init */
3868 0, /* tp_alloc */
3869 filterfalse_new, /* tp_new */
3870 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003871};
3872
3873
3874/* count object ************************************************************/
3875
3876typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003877 PyObject_HEAD
3878 Py_ssize_t cnt;
3879 PyObject *long_cnt;
3880 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003881} countobject;
3882
Raymond Hettinger30729212009-02-12 06:28:27 +00003883/* Counting logic and invariants:
3884
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003885fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003886
Serhiy Storchaka483405b2015-02-17 10:14:30 +02003887 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003888 Advances with: cnt += 1
3889 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003890
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003891slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003892
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003893 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3894 All counting is done with python objects (no overflows or underflows).
3895 Advances with: long_cnt += long_step
3896 Step may be zero -- effectively a slow version of repeat(cnt).
3897 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003898*/
3899
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003900static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003901
3902static PyObject *
3903count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3904{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003905 countobject *lz;
3906 int slow_mode = 0;
3907 Py_ssize_t cnt = 0;
3908 PyObject *long_cnt = NULL;
3909 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003910 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003911 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003912
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003913 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3914 kwlist, &long_cnt, &long_step))
3915 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003916
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003917 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3918 (long_step != NULL && !PyNumber_Check(long_step))) {
3919 PyErr_SetString(PyExc_TypeError, "a number is required");
3920 return NULL;
3921 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003922
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003923 if (long_cnt != NULL) {
3924 cnt = PyLong_AsSsize_t(long_cnt);
3925 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3926 PyErr_Clear();
3927 slow_mode = 1;
3928 }
3929 Py_INCREF(long_cnt);
3930 } else {
3931 cnt = 0;
3932 long_cnt = PyLong_FromLong(0);
3933 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003934
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003935 /* If not specified, step defaults to 1 */
3936 if (long_step == NULL) {
3937 long_step = PyLong_FromLong(1);
3938 if (long_step == NULL) {
3939 Py_DECREF(long_cnt);
3940 return NULL;
3941 }
3942 } else
3943 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003944
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003945 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003946
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003947 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003948 step = PyLong_AsLong(long_step);
3949 if (step != 1) {
3950 slow_mode = 1;
3951 if (step == -1 && PyErr_Occurred())
3952 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003953 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003954
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003955 if (slow_mode)
3956 cnt = PY_SSIZE_T_MAX;
3957 else
3958 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003959
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003960 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3961 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3962 assert(slow_mode ||
3963 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003964
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003965 /* create countobject structure */
3966 lz = (countobject *)type->tp_alloc(type, 0);
3967 if (lz == NULL) {
3968 Py_XDECREF(long_cnt);
3969 return NULL;
3970 }
3971 lz->cnt = cnt;
3972 lz->long_cnt = long_cnt;
3973 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003974
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003975 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003976}
3977
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003978static void
3979count_dealloc(countobject *lz)
3980{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003981 PyObject_GC_UnTrack(lz);
3982 Py_XDECREF(lz->long_cnt);
3983 Py_XDECREF(lz->long_step);
3984 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003985}
3986
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003987static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003988count_traverse(countobject *lz, visitproc visit, void *arg)
3989{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003990 Py_VISIT(lz->long_cnt);
3991 Py_VISIT(lz->long_step);
3992 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003993}
3994
3995static PyObject *
3996count_nextlong(countobject *lz)
3997{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003998 PyObject *long_cnt;
3999 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004000
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004001 long_cnt = lz->long_cnt;
4002 if (long_cnt == NULL) {
4003 /* Switch to slow_mode */
4004 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4005 if (long_cnt == NULL)
4006 return NULL;
4007 }
4008 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004009
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004010 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4011 if (stepped_up == NULL)
4012 return NULL;
4013 lz->long_cnt = stepped_up;
4014 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004015}
4016
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004017static PyObject *
4018count_next(countobject *lz)
4019{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004020 if (lz->cnt == PY_SSIZE_T_MAX)
4021 return count_nextlong(lz);
4022 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004023}
4024
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004025static PyObject *
4026count_repr(countobject *lz)
4027{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004028 if (lz->cnt != PY_SSIZE_T_MAX)
4029 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004030
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004031 if (PyLong_Check(lz->long_step)) {
4032 long step = PyLong_AsLong(lz->long_step);
4033 if (step == -1 && PyErr_Occurred()) {
4034 PyErr_Clear();
4035 }
4036 if (step == 1) {
4037 /* Don't display step when it is an integer equal to 1 */
4038 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
4039 }
4040 }
4041 return PyUnicode_FromFormat("count(%R, %R)",
4042 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004043}
4044
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004045static PyObject *
4046count_reduce(countobject *lz)
4047{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004048 if (lz->cnt == PY_SSIZE_T_MAX)
4049 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4050 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004051}
4052
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004053static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004054 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004055 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004056 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004057};
4058
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004059PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00004060 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004061\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004062Return a count object whose .__next__() method returns consecutive values.\n\
4063Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00004064 def count(firstval=0, step=1):\n\
Eli Benderskyc554f722013-09-03 06:37:19 -07004065 x = firstval\n\
4066 while 1:\n\
4067 yield x\n\
4068 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004069
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004070static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004071 PyVarObject_HEAD_INIT(NULL, 0)
4072 "itertools.count", /* tp_name */
4073 sizeof(countobject), /* tp_basicsize */
4074 0, /* tp_itemsize */
4075 /* methods */
4076 (destructor)count_dealloc, /* tp_dealloc */
4077 0, /* tp_print */
4078 0, /* tp_getattr */
4079 0, /* tp_setattr */
4080 0, /* tp_reserved */
4081 (reprfunc)count_repr, /* tp_repr */
4082 0, /* tp_as_number */
4083 0, /* tp_as_sequence */
4084 0, /* tp_as_mapping */
4085 0, /* tp_hash */
4086 0, /* tp_call */
4087 0, /* tp_str */
4088 PyObject_GenericGetAttr, /* tp_getattro */
4089 0, /* tp_setattro */
4090 0, /* tp_as_buffer */
4091 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4092 Py_TPFLAGS_BASETYPE, /* tp_flags */
4093 count_doc, /* tp_doc */
4094 (traverseproc)count_traverse, /* tp_traverse */
4095 0, /* tp_clear */
4096 0, /* tp_richcompare */
4097 0, /* tp_weaklistoffset */
4098 PyObject_SelfIter, /* tp_iter */
4099 (iternextfunc)count_next, /* tp_iternext */
4100 count_methods, /* tp_methods */
4101 0, /* tp_members */
4102 0, /* tp_getset */
4103 0, /* tp_base */
4104 0, /* tp_dict */
4105 0, /* tp_descr_get */
4106 0, /* tp_descr_set */
4107 0, /* tp_dictoffset */
4108 0, /* tp_init */
4109 0, /* tp_alloc */
4110 count_new, /* tp_new */
4111 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004112};
4113
4114
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004115/* repeat object ************************************************************/
4116
4117typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004118 PyObject_HEAD
4119 PyObject *element;
4120 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004121} repeatobject;
4122
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004123static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004124
4125static PyObject *
4126repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4127{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004128 repeatobject *ro;
4129 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004130 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004131 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004132
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004133 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4134 &element, &cnt))
4135 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004136
Raymond Hettinger97d35552014-06-24 21:36:58 -07004137 if (kwds != NULL)
4138 n_kwds = PyDict_Size(kwds);
4139 /* Does user supply times argument? */
4140 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004141 cnt = 0;
4142
4143 ro = (repeatobject *)type->tp_alloc(type, 0);
4144 if (ro == NULL)
4145 return NULL;
4146 Py_INCREF(element);
4147 ro->element = element;
4148 ro->cnt = cnt;
4149 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004150}
4151
4152static void
4153repeat_dealloc(repeatobject *ro)
4154{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004155 PyObject_GC_UnTrack(ro);
4156 Py_XDECREF(ro->element);
4157 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004158}
4159
4160static int
4161repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4162{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004163 Py_VISIT(ro->element);
4164 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004165}
4166
4167static PyObject *
4168repeat_next(repeatobject *ro)
4169{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004170 if (ro->cnt == 0)
4171 return NULL;
4172 if (ro->cnt > 0)
4173 ro->cnt--;
4174 Py_INCREF(ro->element);
4175 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004176}
4177
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004178static PyObject *
4179repeat_repr(repeatobject *ro)
4180{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004181 if (ro->cnt == -1)
4182 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4183 else
4184 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4185}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004186
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004187static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004188repeat_len(repeatobject *ro)
4189{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004190 if (ro->cnt == -1) {
4191 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4192 return NULL;
4193 }
4194 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004195}
4196
Armin Rigof5b3e362006-02-11 21:32:43 +00004197PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004198
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004199static PyObject *
4200repeat_reduce(repeatobject *ro)
4201{
4202 /* unpickle this so that a new repeat iterator is constructed with an
4203 * object, then call __setstate__ on it to set cnt
4204 */
4205 if (ro->cnt >= 0)
4206 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4207 else
4208 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4209}
4210
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004211static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004212 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004213 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004214 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004215};
4216
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004217PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004218"repeat(object [,times]) -> create an iterator which returns the object\n\
4219for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004220endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004221
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004222static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004223 PyVarObject_HEAD_INIT(NULL, 0)
4224 "itertools.repeat", /* tp_name */
4225 sizeof(repeatobject), /* tp_basicsize */
4226 0, /* tp_itemsize */
4227 /* methods */
4228 (destructor)repeat_dealloc, /* tp_dealloc */
4229 0, /* tp_print */
4230 0, /* tp_getattr */
4231 0, /* tp_setattr */
4232 0, /* tp_reserved */
4233 (reprfunc)repeat_repr, /* tp_repr */
4234 0, /* tp_as_number */
4235 0, /* tp_as_sequence */
4236 0, /* tp_as_mapping */
4237 0, /* tp_hash */
4238 0, /* tp_call */
4239 0, /* tp_str */
4240 PyObject_GenericGetAttr, /* tp_getattro */
4241 0, /* tp_setattro */
4242 0, /* tp_as_buffer */
4243 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4244 Py_TPFLAGS_BASETYPE, /* tp_flags */
4245 repeat_doc, /* tp_doc */
4246 (traverseproc)repeat_traverse, /* tp_traverse */
4247 0, /* tp_clear */
4248 0, /* tp_richcompare */
4249 0, /* tp_weaklistoffset */
4250 PyObject_SelfIter, /* tp_iter */
4251 (iternextfunc)repeat_next, /* tp_iternext */
4252 repeat_methods, /* tp_methods */
4253 0, /* tp_members */
4254 0, /* tp_getset */
4255 0, /* tp_base */
4256 0, /* tp_dict */
4257 0, /* tp_descr_get */
4258 0, /* tp_descr_set */
4259 0, /* tp_dictoffset */
4260 0, /* tp_init */
4261 0, /* tp_alloc */
4262 repeat_new, /* tp_new */
4263 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004264};
4265
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004266/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004267
4268#include "Python.h"
4269
4270typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004271 PyObject_HEAD
4272 Py_ssize_t tuplesize;
4273 Py_ssize_t numactive;
4274 PyObject *ittuple; /* tuple of iterators */
4275 PyObject *result;
4276 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004277} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004278
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004279static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004280
4281static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004282zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004283{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004284 ziplongestobject *lz;
4285 Py_ssize_t i;
4286 PyObject *ittuple; /* tuple of iterators */
4287 PyObject *result;
4288 PyObject *fillvalue = Py_None;
4289 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004290
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004291 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
4292 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4293 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
4294 PyErr_SetString(PyExc_TypeError,
4295 "zip_longest() got an unexpected keyword argument");
4296 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004297 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004298 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004299
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004300 /* args must be a tuple */
4301 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00004302
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004303 /* obtain iterators */
4304 ittuple = PyTuple_New(tuplesize);
4305 if (ittuple == NULL)
4306 return NULL;
4307 for (i=0; i < tuplesize; ++i) {
4308 PyObject *item = PyTuple_GET_ITEM(args, i);
4309 PyObject *it = PyObject_GetIter(item);
4310 if (it == NULL) {
4311 if (PyErr_ExceptionMatches(PyExc_TypeError))
4312 PyErr_Format(PyExc_TypeError,
4313 "zip_longest argument #%zd must support iteration",
4314 i+1);
4315 Py_DECREF(ittuple);
4316 return NULL;
4317 }
4318 PyTuple_SET_ITEM(ittuple, i, it);
4319 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004320
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004321 /* create a result holder */
4322 result = PyTuple_New(tuplesize);
4323 if (result == NULL) {
4324 Py_DECREF(ittuple);
4325 return NULL;
4326 }
4327 for (i=0 ; i < tuplesize ; i++) {
4328 Py_INCREF(Py_None);
4329 PyTuple_SET_ITEM(result, i, Py_None);
4330 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004331
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004332 /* create ziplongestobject structure */
4333 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4334 if (lz == NULL) {
4335 Py_DECREF(ittuple);
4336 Py_DECREF(result);
4337 return NULL;
4338 }
4339 lz->ittuple = ittuple;
4340 lz->tuplesize = tuplesize;
4341 lz->numactive = tuplesize;
4342 lz->result = result;
4343 Py_INCREF(fillvalue);
4344 lz->fillvalue = fillvalue;
4345 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004346}
4347
4348static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004349zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004350{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004351 PyObject_GC_UnTrack(lz);
4352 Py_XDECREF(lz->ittuple);
4353 Py_XDECREF(lz->result);
4354 Py_XDECREF(lz->fillvalue);
4355 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004356}
4357
4358static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004359zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004360{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004361 Py_VISIT(lz->ittuple);
4362 Py_VISIT(lz->result);
4363 Py_VISIT(lz->fillvalue);
4364 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004365}
4366
4367static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004368zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004369{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004370 Py_ssize_t i;
4371 Py_ssize_t tuplesize = lz->tuplesize;
4372 PyObject *result = lz->result;
4373 PyObject *it;
4374 PyObject *item;
4375 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004376
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004377 if (tuplesize == 0)
4378 return NULL;
4379 if (lz->numactive == 0)
4380 return NULL;
4381 if (Py_REFCNT(result) == 1) {
4382 Py_INCREF(result);
4383 for (i=0 ; i < tuplesize ; i++) {
4384 it = PyTuple_GET_ITEM(lz->ittuple, i);
4385 if (it == NULL) {
4386 Py_INCREF(lz->fillvalue);
4387 item = lz->fillvalue;
4388 } else {
4389 item = PyIter_Next(it);
4390 if (item == NULL) {
4391 lz->numactive -= 1;
4392 if (lz->numactive == 0 || PyErr_Occurred()) {
4393 lz->numactive = 0;
4394 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004395 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004396 } else {
4397 Py_INCREF(lz->fillvalue);
4398 item = lz->fillvalue;
4399 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4400 Py_DECREF(it);
4401 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004402 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004403 }
4404 olditem = PyTuple_GET_ITEM(result, i);
4405 PyTuple_SET_ITEM(result, i, item);
4406 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004407 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004408 } else {
4409 result = PyTuple_New(tuplesize);
4410 if (result == NULL)
4411 return NULL;
4412 for (i=0 ; i < tuplesize ; i++) {
4413 it = PyTuple_GET_ITEM(lz->ittuple, i);
4414 if (it == NULL) {
4415 Py_INCREF(lz->fillvalue);
4416 item = lz->fillvalue;
4417 } else {
4418 item = PyIter_Next(it);
4419 if (item == NULL) {
4420 lz->numactive -= 1;
4421 if (lz->numactive == 0 || PyErr_Occurred()) {
4422 lz->numactive = 0;
4423 Py_DECREF(result);
4424 return NULL;
4425 } else {
4426 Py_INCREF(lz->fillvalue);
4427 item = lz->fillvalue;
4428 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4429 Py_DECREF(it);
4430 }
4431 }
4432 }
4433 PyTuple_SET_ITEM(result, i, item);
4434 }
4435 }
4436 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004437}
4438
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004439static PyObject *
4440zip_longest_reduce(ziplongestobject *lz)
4441{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004442
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004443 /* Create a new tuple with empty sequences where appropriate to pickle.
4444 * Then use setstate to set the fillvalue
4445 */
4446 int i;
4447 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
4448 if (args == NULL)
4449 return NULL;
4450 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4451 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4452 if (elem == NULL) {
4453 elem = PyTuple_New(0);
4454 if (elem == NULL) {
4455 Py_DECREF(args);
4456 return NULL;
4457 }
4458 } else
4459 Py_INCREF(elem);
4460 PyTuple_SET_ITEM(args, i, elem);
4461 }
4462 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4463}
4464
4465static PyObject *
4466zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4467{
4468 Py_CLEAR(lz->fillvalue);
4469 lz->fillvalue = state;
4470 Py_INCREF(lz->fillvalue);
4471 Py_RETURN_NONE;
4472}
4473
4474static PyMethodDef zip_longest_methods[] = {
4475 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4476 reduce_doc},
4477 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4478 setstate_doc},
4479 {NULL, NULL} /* sentinel */
4480};
4481
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004482PyDoc_STRVAR(zip_longest_doc,
4483"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004484\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004485Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004486the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004487method continues until the longest iterable in the argument sequence\n\
4488is exhausted and then it raises StopIteration. When the shorter iterables\n\
4489are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4490defaults to None or can be specified by a keyword argument.\n\
4491");
4492
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004493static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004494 PyVarObject_HEAD_INIT(NULL, 0)
4495 "itertools.zip_longest", /* tp_name */
4496 sizeof(ziplongestobject), /* tp_basicsize */
4497 0, /* tp_itemsize */
4498 /* methods */
4499 (destructor)zip_longest_dealloc, /* tp_dealloc */
4500 0, /* tp_print */
4501 0, /* tp_getattr */
4502 0, /* tp_setattr */
4503 0, /* tp_reserved */
4504 0, /* tp_repr */
4505 0, /* tp_as_number */
4506 0, /* tp_as_sequence */
4507 0, /* tp_as_mapping */
4508 0, /* tp_hash */
4509 0, /* tp_call */
4510 0, /* tp_str */
4511 PyObject_GenericGetAttr, /* tp_getattro */
4512 0, /* tp_setattro */
4513 0, /* tp_as_buffer */
4514 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4515 Py_TPFLAGS_BASETYPE, /* tp_flags */
4516 zip_longest_doc, /* tp_doc */
4517 (traverseproc)zip_longest_traverse, /* tp_traverse */
4518 0, /* tp_clear */
4519 0, /* tp_richcompare */
4520 0, /* tp_weaklistoffset */
4521 PyObject_SelfIter, /* tp_iter */
4522 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004523 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004524 0, /* tp_members */
4525 0, /* tp_getset */
4526 0, /* tp_base */
4527 0, /* tp_dict */
4528 0, /* tp_descr_get */
4529 0, /* tp_descr_set */
4530 0, /* tp_dictoffset */
4531 0, /* tp_init */
4532 0, /* tp_alloc */
4533 zip_longest_new, /* tp_new */
4534 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004535};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004536
4537/* module level code ********************************************************/
4538
4539PyDoc_STRVAR(module_doc,
4540"Functional tools for creating and using iterators.\n\
4541\n\
4542Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004543count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004544cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004545repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004546\n\
4547Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004548accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004549chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger8df58f72013-09-09 01:29:40 -05004550chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004551compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4552dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4553groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004554filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004555islice(seq, [start,] stop [, step]) --> elements from\n\
4556 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004557starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004558tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004559takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004560zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004561\n\
4562Combinatoric generators:\n\
4563product(p, q, ... [repeat=1]) --> cartesian product\n\
4564permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004565combinations(p, r)\n\
4566combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004567");
4568
4569
Raymond Hettingerad983e72003-11-12 14:32:26 +00004570static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004571 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4572 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004573};
4574
Martin v. Löwis1a214512008-06-11 05:26:20 +00004575
4576static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004577 PyModuleDef_HEAD_INIT,
4578 "itertools",
4579 module_doc,
4580 -1,
4581 module_methods,
4582 NULL,
4583 NULL,
4584 NULL,
4585 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004586};
4587
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004588PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004589PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004590{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004591 int i;
4592 PyObject *m;
4593 char *name;
4594 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004595 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004596 &combinations_type,
4597 &cwr_type,
4598 &cycle_type,
4599 &dropwhile_type,
4600 &takewhile_type,
4601 &islice_type,
4602 &starmap_type,
4603 &chain_type,
4604 &compress_type,
4605 &filterfalse_type,
4606 &count_type,
4607 &ziplongest_type,
4608 &permutations_type,
4609 &product_type,
4610 &repeat_type,
4611 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004612 &_grouper_type,
4613 &tee_type,
4614 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004615 NULL
4616 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004617
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004618 Py_TYPE(&teedataobject_type) = &PyType_Type;
4619 m = PyModule_Create(&itertoolsmodule);
4620 if (m == NULL)
4621 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004622
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004623 for (i=0 ; typelist[i] != NULL ; i++) {
4624 if (PyType_Ready(typelist[i]) < 0)
4625 return NULL;
4626 name = strchr(typelist[i]->tp_name, '.');
4627 assert (name != NULL);
4628 Py_INCREF(typelist[i]);
4629 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4630 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004631
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004632 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004633}