blob: 01231816123bdb29ca6c374cee46d858789636bb [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 Hettingerfb92f392013-09-09 02:01:35 -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;
404 int numread;
405 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;
412 int index;
413 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 Pitrouf95a1b32010-05-09 15:52:27 +00001495 iternext = *Py_TYPE(it)->tp_iternext;
1496 while (lz->cnt < lz->next) {
1497 item = iternext(it);
1498 if (item == NULL)
1499 return NULL;
1500 Py_DECREF(item);
1501 lz->cnt++;
1502 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001503 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001504 return NULL;
1505 item = iternext(it);
1506 if (item == NULL)
1507 return NULL;
1508 lz->cnt++;
1509 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001510 /* The (size_t) cast below avoids the danger of undefined
1511 behaviour from signed integer overflow. */
1512 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001513 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1514 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001515 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001516}
1517
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001518static PyObject *
1519islice_reduce(isliceobject *lz)
1520{
1521 /* When unpickled, generate a new object with the same bounds,
1522 * then 'setstate' with the next and count
1523 */
1524 PyObject *stop;
1525 if (lz->stop == -1) {
1526 stop = Py_None;
1527 Py_INCREF(stop);
1528 } else {
1529 stop = PyLong_FromSsize_t(lz->stop);
1530 if (stop == NULL)
1531 return NULL;
1532 }
1533 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1534 lz->it, lz->next, stop, lz->step,
1535 lz->cnt);
1536}
1537
1538static PyObject *
1539islice_setstate(isliceobject *lz, PyObject *state)
1540{
1541 Py_ssize_t cnt = PyLong_AsSsize_t(state);
1542 if (cnt == -1 && PyErr_Occurred())
1543 return NULL;
1544 lz->cnt = cnt;
1545 Py_RETURN_NONE;
1546}
1547
1548static PyMethodDef islice_methods[] = {
1549 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1550 reduce_doc},
1551 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1552 setstate_doc},
1553 {NULL, NULL} /* sentinel */
1554};
1555
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001556PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001557"islice(iterable, stop) --> islice object\n\
1558islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001559\n\
1560Return an iterator whose next() method returns selected values from an\n\
1561iterable. If start is specified, will skip all preceding elements;\n\
1562otherwise, start defaults to zero. Step defaults to one. If\n\
1563specified as another value, step determines how many values are \n\
1564skipped between successive calls. Works like a slice() on a list\n\
1565but returns an iterator.");
1566
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001567static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001568 PyVarObject_HEAD_INIT(NULL, 0)
1569 "itertools.islice", /* tp_name */
1570 sizeof(isliceobject), /* tp_basicsize */
1571 0, /* tp_itemsize */
1572 /* methods */
1573 (destructor)islice_dealloc, /* tp_dealloc */
1574 0, /* tp_print */
1575 0, /* tp_getattr */
1576 0, /* tp_setattr */
1577 0, /* tp_reserved */
1578 0, /* tp_repr */
1579 0, /* tp_as_number */
1580 0, /* tp_as_sequence */
1581 0, /* tp_as_mapping */
1582 0, /* tp_hash */
1583 0, /* tp_call */
1584 0, /* tp_str */
1585 PyObject_GenericGetAttr, /* tp_getattro */
1586 0, /* tp_setattro */
1587 0, /* tp_as_buffer */
1588 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1589 Py_TPFLAGS_BASETYPE, /* tp_flags */
1590 islice_doc, /* tp_doc */
1591 (traverseproc)islice_traverse, /* tp_traverse */
1592 0, /* tp_clear */
1593 0, /* tp_richcompare */
1594 0, /* tp_weaklistoffset */
1595 PyObject_SelfIter, /* tp_iter */
1596 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001597 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001598 0, /* tp_members */
1599 0, /* tp_getset */
1600 0, /* tp_base */
1601 0, /* tp_dict */
1602 0, /* tp_descr_get */
1603 0, /* tp_descr_set */
1604 0, /* tp_dictoffset */
1605 0, /* tp_init */
1606 0, /* tp_alloc */
1607 islice_new, /* tp_new */
1608 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001609};
1610
1611
1612/* starmap object ************************************************************/
1613
1614typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001615 PyObject_HEAD
1616 PyObject *func;
1617 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001618} starmapobject;
1619
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001620static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001621
1622static PyObject *
1623starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1624{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001625 PyObject *func, *seq;
1626 PyObject *it;
1627 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001628
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001629 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1630 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001631
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001632 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1633 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001634
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001635 /* Get iterator. */
1636 it = PyObject_GetIter(seq);
1637 if (it == NULL)
1638 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001639
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001640 /* create starmapobject structure */
1641 lz = (starmapobject *)type->tp_alloc(type, 0);
1642 if (lz == NULL) {
1643 Py_DECREF(it);
1644 return NULL;
1645 }
1646 Py_INCREF(func);
1647 lz->func = func;
1648 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001649
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001650 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001651}
1652
1653static void
1654starmap_dealloc(starmapobject *lz)
1655{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001656 PyObject_GC_UnTrack(lz);
1657 Py_XDECREF(lz->func);
1658 Py_XDECREF(lz->it);
1659 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001660}
1661
1662static int
1663starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1664{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001665 Py_VISIT(lz->it);
1666 Py_VISIT(lz->func);
1667 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001668}
1669
1670static PyObject *
1671starmap_next(starmapobject *lz)
1672{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001673 PyObject *args;
1674 PyObject *result;
1675 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001676
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001677 args = (*Py_TYPE(it)->tp_iternext)(it);
1678 if (args == NULL)
1679 return NULL;
1680 if (!PyTuple_CheckExact(args)) {
1681 PyObject *newargs = PySequence_Tuple(args);
1682 Py_DECREF(args);
1683 if (newargs == NULL)
1684 return NULL;
1685 args = newargs;
1686 }
1687 result = PyObject_Call(lz->func, args, NULL);
1688 Py_DECREF(args);
1689 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001690}
1691
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001692static PyObject *
1693starmap_reduce(starmapobject *lz)
1694{
1695 /* Just pickle the iterator */
1696 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1697}
1698
1699static PyMethodDef starmap_methods[] = {
1700 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1701 reduce_doc},
1702 {NULL, NULL} /* sentinel */
1703};
1704
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001705PyDoc_STRVAR(starmap_doc,
1706"starmap(function, sequence) --> starmap object\n\
1707\n\
1708Return an iterator whose values are returned from the function evaluated\n\
1709with a argument tuple taken from the given sequence.");
1710
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001711static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001712 PyVarObject_HEAD_INIT(NULL, 0)
1713 "itertools.starmap", /* tp_name */
1714 sizeof(starmapobject), /* tp_basicsize */
1715 0, /* tp_itemsize */
1716 /* methods */
1717 (destructor)starmap_dealloc, /* tp_dealloc */
1718 0, /* tp_print */
1719 0, /* tp_getattr */
1720 0, /* tp_setattr */
1721 0, /* tp_reserved */
1722 0, /* tp_repr */
1723 0, /* tp_as_number */
1724 0, /* tp_as_sequence */
1725 0, /* tp_as_mapping */
1726 0, /* tp_hash */
1727 0, /* tp_call */
1728 0, /* tp_str */
1729 PyObject_GenericGetAttr, /* tp_getattro */
1730 0, /* tp_setattro */
1731 0, /* tp_as_buffer */
1732 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1733 Py_TPFLAGS_BASETYPE, /* tp_flags */
1734 starmap_doc, /* tp_doc */
1735 (traverseproc)starmap_traverse, /* tp_traverse */
1736 0, /* tp_clear */
1737 0, /* tp_richcompare */
1738 0, /* tp_weaklistoffset */
1739 PyObject_SelfIter, /* tp_iter */
1740 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001741 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001742 0, /* tp_members */
1743 0, /* tp_getset */
1744 0, /* tp_base */
1745 0, /* tp_dict */
1746 0, /* tp_descr_get */
1747 0, /* tp_descr_set */
1748 0, /* tp_dictoffset */
1749 0, /* tp_init */
1750 0, /* tp_alloc */
1751 starmap_new, /* tp_new */
1752 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001753};
1754
1755
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001756/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001757
1758typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001759 PyObject_HEAD
1760 PyObject *source; /* Iterator over input iterables */
1761 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001762} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001763
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001764static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001765
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001766static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001767chain_new_internal(PyTypeObject *type, PyObject *source)
1768{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001769 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001770
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001771 lz = (chainobject *)type->tp_alloc(type, 0);
1772 if (lz == NULL) {
1773 Py_DECREF(source);
1774 return NULL;
1775 }
1776
1777 lz->source = source;
1778 lz->active = NULL;
1779 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001780}
1781
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001782static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001783chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001784{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001785 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001786
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001787 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1788 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001789
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001790 source = PyObject_GetIter(args);
1791 if (source == NULL)
1792 return NULL;
1793
1794 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001795}
1796
1797static PyObject *
1798chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1799{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001800 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001801
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001802 source = PyObject_GetIter(arg);
1803 if (source == NULL)
1804 return NULL;
1805
1806 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001807}
1808
1809static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001810chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001811{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001812 PyObject_GC_UnTrack(lz);
1813 Py_XDECREF(lz->active);
1814 Py_XDECREF(lz->source);
1815 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001816}
1817
Raymond Hettinger2012f172003-02-07 05:32:58 +00001818static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001819chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001820{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001821 Py_VISIT(lz->source);
1822 Py_VISIT(lz->active);
1823 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001824}
1825
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001826static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001827chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001828{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001829 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001830
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001831 if (lz->source == NULL)
1832 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001833
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001834 if (lz->active == NULL) {
1835 PyObject *iterable = PyIter_Next(lz->source);
1836 if (iterable == NULL) {
1837 Py_CLEAR(lz->source);
1838 return NULL; /* no more input sources */
1839 }
1840 lz->active = PyObject_GetIter(iterable);
1841 Py_DECREF(iterable);
1842 if (lz->active == NULL) {
1843 Py_CLEAR(lz->source);
1844 return NULL; /* input not iterable */
1845 }
1846 }
1847 item = PyIter_Next(lz->active);
1848 if (item != NULL)
1849 return item;
1850 if (PyErr_Occurred()) {
1851 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1852 PyErr_Clear();
1853 else
1854 return NULL; /* input raised an exception */
1855 }
1856 Py_CLEAR(lz->active);
1857 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001858}
1859
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001860static PyObject *
1861chain_reduce(chainobject *lz)
1862{
1863 if (lz->source) {
1864 /* we can't pickle function objects (itertools.from_iterable) so
1865 * we must use setstate to replace the iterable. One day we
1866 * will fix pickling of functions
1867 */
1868 if (lz->active) {
1869 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1870 } else {
1871 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1872 }
1873 } else {
1874 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1875 }
1876 return NULL;
1877}
1878
1879static PyObject *
1880chain_setstate(chainobject *lz, PyObject *state)
1881{
1882 PyObject *source, *active=NULL;
1883 if (! PyArg_ParseTuple(state, "O|O", &source, &active))
1884 return NULL;
1885
1886 Py_CLEAR(lz->source);
1887 lz->source = source;
1888 Py_INCREF(lz->source);
1889 Py_CLEAR(lz->active);
1890 lz->active = active;
1891 Py_XINCREF(lz->active);
1892 Py_RETURN_NONE;
1893}
1894
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001895PyDoc_STRVAR(chain_doc,
1896"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001897\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001898Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001899first iterable until it is exhausted, then elements from the next\n\
1900iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001901
Christian Heimesf16baeb2008-02-29 14:57:44 +00001902PyDoc_STRVAR(chain_from_iterable_doc,
1903"chain.from_iterable(iterable) --> chain object\n\
1904\n\
1905Alternate chain() contructor taking a single iterable argument\n\
1906that evaluates lazily.");
1907
1908static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001909 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1910 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001911 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1912 reduce_doc},
1913 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1914 setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001915 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001916};
1917
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001918static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001919 PyVarObject_HEAD_INIT(NULL, 0)
1920 "itertools.chain", /* tp_name */
1921 sizeof(chainobject), /* tp_basicsize */
1922 0, /* tp_itemsize */
1923 /* methods */
1924 (destructor)chain_dealloc, /* tp_dealloc */
1925 0, /* tp_print */
1926 0, /* tp_getattr */
1927 0, /* tp_setattr */
1928 0, /* tp_reserved */
1929 0, /* tp_repr */
1930 0, /* tp_as_number */
1931 0, /* tp_as_sequence */
1932 0, /* tp_as_mapping */
1933 0, /* tp_hash */
1934 0, /* tp_call */
1935 0, /* tp_str */
1936 PyObject_GenericGetAttr, /* tp_getattro */
1937 0, /* tp_setattro */
1938 0, /* tp_as_buffer */
1939 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1940 Py_TPFLAGS_BASETYPE, /* tp_flags */
1941 chain_doc, /* tp_doc */
1942 (traverseproc)chain_traverse, /* tp_traverse */
1943 0, /* tp_clear */
1944 0, /* tp_richcompare */
1945 0, /* tp_weaklistoffset */
1946 PyObject_SelfIter, /* tp_iter */
1947 (iternextfunc)chain_next, /* tp_iternext */
1948 chain_methods, /* tp_methods */
1949 0, /* tp_members */
1950 0, /* tp_getset */
1951 0, /* tp_base */
1952 0, /* tp_dict */
1953 0, /* tp_descr_get */
1954 0, /* tp_descr_set */
1955 0, /* tp_dictoffset */
1956 0, /* tp_init */
1957 0, /* tp_alloc */
1958 chain_new, /* tp_new */
1959 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001960};
1961
1962
Christian Heimesc3f30c42008-02-22 16:37:40 +00001963/* product object ************************************************************/
1964
1965typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001966 PyObject_HEAD
1967 PyObject *pools; /* tuple of pool tuples */
1968 Py_ssize_t *indices; /* one index per pool */
1969 PyObject *result; /* most recently returned result tuple */
1970 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001971} productobject;
1972
1973static PyTypeObject product_type;
1974
1975static PyObject *
1976product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1977{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001978 productobject *lz;
1979 Py_ssize_t nargs, npools, repeat=1;
1980 PyObject *pools = NULL;
1981 Py_ssize_t *indices = NULL;
1982 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001983
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001984 if (kwds != NULL) {
1985 char *kwlist[] = {"repeat", 0};
1986 PyObject *tmpargs = PyTuple_New(0);
1987 if (tmpargs == NULL)
1988 return NULL;
1989 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1990 Py_DECREF(tmpargs);
1991 return NULL;
1992 }
1993 Py_DECREF(tmpargs);
1994 if (repeat < 0) {
1995 PyErr_SetString(PyExc_ValueError,
1996 "repeat argument cannot be negative");
1997 return NULL;
1998 }
1999 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002000
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002001 assert(PyTuple_Check(args));
2002 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
2003 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002005 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
2006 if (indices == NULL) {
2007 PyErr_NoMemory();
2008 goto error;
2009 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002010
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002011 pools = PyTuple_New(npools);
2012 if (pools == NULL)
2013 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002014
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002015 for (i=0; i < nargs ; ++i) {
2016 PyObject *item = PyTuple_GET_ITEM(args, i);
2017 PyObject *pool = PySequence_Tuple(item);
2018 if (pool == NULL)
2019 goto error;
2020 PyTuple_SET_ITEM(pools, i, pool);
2021 indices[i] = 0;
2022 }
2023 for ( ; i < npools; ++i) {
2024 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2025 Py_INCREF(pool);
2026 PyTuple_SET_ITEM(pools, i, pool);
2027 indices[i] = 0;
2028 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002029
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002030 /* create productobject structure */
2031 lz = (productobject *)type->tp_alloc(type, 0);
2032 if (lz == NULL)
2033 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002034
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002035 lz->pools = pools;
2036 lz->indices = indices;
2037 lz->result = NULL;
2038 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002039
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002040 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002041
2042error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002043 if (indices != NULL)
2044 PyMem_Free(indices);
2045 Py_XDECREF(pools);
2046 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002047}
2048
2049static void
2050product_dealloc(productobject *lz)
2051{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002052 PyObject_GC_UnTrack(lz);
2053 Py_XDECREF(lz->pools);
2054 Py_XDECREF(lz->result);
2055 if (lz->indices != NULL)
2056 PyMem_Free(lz->indices);
2057 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002058}
2059
2060static int
2061product_traverse(productobject *lz, visitproc visit, void *arg)
2062{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002063 Py_VISIT(lz->pools);
2064 Py_VISIT(lz->result);
2065 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002066}
2067
2068static PyObject *
2069product_next(productobject *lz)
2070{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002071 PyObject *pool;
2072 PyObject *elem;
2073 PyObject *oldelem;
2074 PyObject *pools = lz->pools;
2075 PyObject *result = lz->result;
2076 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2077 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002078
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002079 if (lz->stopped)
2080 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002081
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002082 if (result == NULL) {
2083 /* On the first pass, return an initial tuple filled with the
2084 first element from each pool. */
2085 result = PyTuple_New(npools);
2086 if (result == NULL)
2087 goto empty;
2088 lz->result = result;
2089 for (i=0; i < npools; i++) {
2090 pool = PyTuple_GET_ITEM(pools, i);
2091 if (PyTuple_GET_SIZE(pool) == 0)
2092 goto empty;
2093 elem = PyTuple_GET_ITEM(pool, 0);
2094 Py_INCREF(elem);
2095 PyTuple_SET_ITEM(result, i, elem);
2096 }
2097 } else {
2098 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002099
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002100 /* Copy the previous result tuple or re-use it if available */
2101 if (Py_REFCNT(result) > 1) {
2102 PyObject *old_result = result;
2103 result = PyTuple_New(npools);
2104 if (result == NULL)
2105 goto empty;
2106 lz->result = result;
2107 for (i=0; i < npools; i++) {
2108 elem = PyTuple_GET_ITEM(old_result, i);
2109 Py_INCREF(elem);
2110 PyTuple_SET_ITEM(result, i, elem);
2111 }
2112 Py_DECREF(old_result);
2113 }
2114 /* Now, we've got the only copy so we can update it in-place */
2115 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002116
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002117 /* Update the pool indices right-to-left. Only advance to the
2118 next pool when the previous one rolls-over */
2119 for (i=npools-1 ; i >= 0 ; i--) {
2120 pool = PyTuple_GET_ITEM(pools, i);
2121 indices[i]++;
2122 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2123 /* Roll-over and advance to next pool */
2124 indices[i] = 0;
2125 elem = PyTuple_GET_ITEM(pool, 0);
2126 Py_INCREF(elem);
2127 oldelem = PyTuple_GET_ITEM(result, i);
2128 PyTuple_SET_ITEM(result, i, elem);
2129 Py_DECREF(oldelem);
2130 } else {
2131 /* No rollover. Just increment and stop here. */
2132 elem = PyTuple_GET_ITEM(pool, indices[i]);
2133 Py_INCREF(elem);
2134 oldelem = PyTuple_GET_ITEM(result, i);
2135 PyTuple_SET_ITEM(result, i, elem);
2136 Py_DECREF(oldelem);
2137 break;
2138 }
2139 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002140
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002141 /* If i is negative, then the indices have all rolled-over
2142 and we're done. */
2143 if (i < 0)
2144 goto empty;
2145 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002146
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002147 Py_INCREF(result);
2148 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002149
2150empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002151 lz->stopped = 1;
2152 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002153}
2154
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002155static PyObject *
2156product_reduce(productobject *lz)
2157{
2158 if (lz->stopped) {
2159 return Py_BuildValue("O(())", Py_TYPE(lz));
2160 } else if (lz->result == NULL) {
2161 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2162 } else {
2163 PyObject *indices;
2164 Py_ssize_t n, i;
2165
2166 /* we must pickle the indices use them for setstate, and
2167 * additionally indicate that the iterator has started
2168 */
2169 n = PyTuple_GET_SIZE(lz->pools);
2170 indices = PyTuple_New(n);
2171 if (indices == NULL)
2172 return NULL;
2173 for (i=0; i<n; i++){
2174 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2175 if (!index) {
2176 Py_DECREF(indices);
2177 return NULL;
2178 }
2179 PyTuple_SET_ITEM(indices, i, index);
2180 }
2181 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2182 }
2183}
2184
2185static PyObject *
2186product_setstate(productobject *lz, PyObject *state)
2187{
2188 PyObject *result;
2189 Py_ssize_t n, i;
2190
2191 n = PyTuple_GET_SIZE(lz->pools);
2192 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2193 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2194 return NULL;
2195 }
2196 for (i=0; i<n; i++)
2197 {
2198 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2199 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2200 if (index < 0 && PyErr_Occurred())
2201 return NULL; /* not an integer */
2202 /* clamp the index */
2203 if (index < 0)
2204 index = 0;
2205 else if (index > n-1)
2206 index = n-1;
2207 lz->indices[i] = index;
2208 }
2209
2210 result = PyTuple_New(n);
2211 if (!result)
2212 return NULL;
2213 for (i=0; i<n; i++) {
2214 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2215 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2216 Py_INCREF(element);
2217 PyTuple_SET_ITEM(result, i, element);
2218 }
2219 Py_CLEAR(lz->result);
2220 lz->result = result;
2221 Py_RETURN_NONE;
2222}
2223
2224static PyMethodDef product_methods[] = {
2225 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2226 reduce_doc},
2227 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2228 setstate_doc},
2229 {NULL, NULL} /* sentinel */
2230};
2231
Christian Heimesc3f30c42008-02-22 16:37:40 +00002232PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002233"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002234\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002235Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002236For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2237The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2238cycle in a manner similar to an odometer (with the rightmost element changing\n\
2239on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002240To compute the product of an iterable with itself, specify the number\n\
2241of repetitions with the optional repeat keyword argument. For example,\n\
2242product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002243product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2244product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2245
2246static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002247 PyVarObject_HEAD_INIT(NULL, 0)
2248 "itertools.product", /* tp_name */
2249 sizeof(productobject), /* tp_basicsize */
2250 0, /* tp_itemsize */
2251 /* methods */
2252 (destructor)product_dealloc, /* tp_dealloc */
2253 0, /* tp_print */
2254 0, /* tp_getattr */
2255 0, /* tp_setattr */
2256 0, /* tp_reserved */
2257 0, /* tp_repr */
2258 0, /* tp_as_number */
2259 0, /* tp_as_sequence */
2260 0, /* tp_as_mapping */
2261 0, /* tp_hash */
2262 0, /* tp_call */
2263 0, /* tp_str */
2264 PyObject_GenericGetAttr, /* tp_getattro */
2265 0, /* tp_setattro */
2266 0, /* tp_as_buffer */
2267 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2268 Py_TPFLAGS_BASETYPE, /* tp_flags */
2269 product_doc, /* tp_doc */
2270 (traverseproc)product_traverse, /* tp_traverse */
2271 0, /* tp_clear */
2272 0, /* tp_richcompare */
2273 0, /* tp_weaklistoffset */
2274 PyObject_SelfIter, /* tp_iter */
2275 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002276 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002277 0, /* tp_members */
2278 0, /* tp_getset */
2279 0, /* tp_base */
2280 0, /* tp_dict */
2281 0, /* tp_descr_get */
2282 0, /* tp_descr_set */
2283 0, /* tp_dictoffset */
2284 0, /* tp_init */
2285 0, /* tp_alloc */
2286 product_new, /* tp_new */
2287 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002288};
2289
2290
Christian Heimes380f7f22008-02-28 11:19:05 +00002291/* combinations object ************************************************************/
2292
2293typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002294 PyObject_HEAD
2295 PyObject *pool; /* input converted to a tuple */
2296 Py_ssize_t *indices; /* one index per result element */
2297 PyObject *result; /* most recently returned result tuple */
2298 Py_ssize_t r; /* size of result tuple */
2299 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002300} combinationsobject;
2301
2302static PyTypeObject combinations_type;
2303
2304static PyObject *
2305combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2306{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002307 combinationsobject *co;
2308 Py_ssize_t n;
2309 Py_ssize_t r;
2310 PyObject *pool = NULL;
2311 PyObject *iterable = NULL;
2312 Py_ssize_t *indices = NULL;
2313 Py_ssize_t i;
2314 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002315
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002316 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2317 &iterable, &r))
2318 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002319
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002320 pool = PySequence_Tuple(iterable);
2321 if (pool == NULL)
2322 goto error;
2323 n = PyTuple_GET_SIZE(pool);
2324 if (r < 0) {
2325 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2326 goto error;
2327 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002328
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002329 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2330 if (indices == NULL) {
2331 PyErr_NoMemory();
2332 goto error;
2333 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002334
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002335 for (i=0 ; i<r ; i++)
2336 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002337
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002338 /* create combinationsobject structure */
2339 co = (combinationsobject *)type->tp_alloc(type, 0);
2340 if (co == NULL)
2341 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002342
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002343 co->pool = pool;
2344 co->indices = indices;
2345 co->result = NULL;
2346 co->r = r;
2347 co->stopped = r > n ? 1 : 0;
2348
2349 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002350
2351error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002352 if (indices != NULL)
2353 PyMem_Free(indices);
2354 Py_XDECREF(pool);
2355 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002356}
2357
2358static void
2359combinations_dealloc(combinationsobject *co)
2360{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002361 PyObject_GC_UnTrack(co);
2362 Py_XDECREF(co->pool);
2363 Py_XDECREF(co->result);
2364 if (co->indices != NULL)
2365 PyMem_Free(co->indices);
2366 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002367}
2368
2369static int
2370combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2371{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002372 Py_VISIT(co->pool);
2373 Py_VISIT(co->result);
2374 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002375}
2376
2377static PyObject *
2378combinations_next(combinationsobject *co)
2379{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002380 PyObject *elem;
2381 PyObject *oldelem;
2382 PyObject *pool = co->pool;
2383 Py_ssize_t *indices = co->indices;
2384 PyObject *result = co->result;
2385 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2386 Py_ssize_t r = co->r;
2387 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002388
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002389 if (co->stopped)
2390 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002391
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002392 if (result == NULL) {
2393 /* On the first pass, initialize result tuple using the indices */
2394 result = PyTuple_New(r);
2395 if (result == NULL)
2396 goto empty;
2397 co->result = result;
2398 for (i=0; i<r ; i++) {
2399 index = indices[i];
2400 elem = PyTuple_GET_ITEM(pool, index);
2401 Py_INCREF(elem);
2402 PyTuple_SET_ITEM(result, i, elem);
2403 }
2404 } else {
2405 /* Copy the previous result tuple or re-use it if available */
2406 if (Py_REFCNT(result) > 1) {
2407 PyObject *old_result = result;
2408 result = PyTuple_New(r);
2409 if (result == NULL)
2410 goto empty;
2411 co->result = result;
2412 for (i=0; i<r ; i++) {
2413 elem = PyTuple_GET_ITEM(old_result, i);
2414 Py_INCREF(elem);
2415 PyTuple_SET_ITEM(result, i, elem);
2416 }
2417 Py_DECREF(old_result);
2418 }
2419 /* Now, we've got the only copy so we can update it in-place
2420 * CPython's empty tuple is a singleton and cached in
2421 * PyTuple's freelist.
2422 */
2423 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002424
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002425 /* Scan indices right-to-left until finding one that is not
2426 at its maximum (i + n - r). */
2427 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2428 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002429
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002430 /* If i is negative, then the indices are all at
2431 their maximum value and we're done. */
2432 if (i < 0)
2433 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002434
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002435 /* Increment the current index which we know is not at its
2436 maximum. Then move back to the right setting each index
2437 to its lowest possible value (one higher than the index
2438 to its left -- this maintains the sort order invariant). */
2439 indices[i]++;
2440 for (j=i+1 ; j<r ; j++)
2441 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002442
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002443 /* Update the result tuple for the new indices
2444 starting with i, the leftmost index that changed */
2445 for ( ; i<r ; i++) {
2446 index = indices[i];
2447 elem = PyTuple_GET_ITEM(pool, index);
2448 Py_INCREF(elem);
2449 oldelem = PyTuple_GET_ITEM(result, i);
2450 PyTuple_SET_ITEM(result, i, elem);
2451 Py_DECREF(oldelem);
2452 }
2453 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002454
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002455 Py_INCREF(result);
2456 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002457
2458empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002459 co->stopped = 1;
2460 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002461}
2462
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002463static PyObject *
2464combinations_reduce(combinationsobject *lz)
2465{
2466 if (lz->result == NULL) {
2467 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2468 } else if (lz->stopped) {
2469 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2470 } else {
2471 PyObject *indices;
2472 Py_ssize_t i;
2473
2474 /* we must pickle the indices and use them for setstate */
2475 indices = PyTuple_New(lz->r);
2476 if (!indices)
2477 return NULL;
2478 for (i=0; i<lz->r; i++)
2479 {
2480 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2481 if (!index) {
2482 Py_DECREF(indices);
2483 return NULL;
2484 }
2485 PyTuple_SET_ITEM(indices, i, index);
2486 }
2487
2488 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2489 }
2490}
2491
2492static PyObject *
2493combinations_setstate(combinationsobject *lz, PyObject *state)
2494{
2495 PyObject *result;
2496 Py_ssize_t i;
2497 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2498
2499 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2500 {
2501 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2502 return NULL;
2503 }
2504
2505 for (i=0; i<lz->r; i++)
2506 {
2507 Py_ssize_t max;
2508 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2509 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2510 if (index == -1 && PyErr_Occurred())
2511 return NULL; /* not an integer */
2512 max = i + n - lz->r;
2513 /* clamp the index (beware of negative max) */
2514 if (index > max)
2515 index = max;
2516 if (index < 0)
2517 index = 0;
2518 lz->indices[i] = index;
2519 }
2520
2521 result = PyTuple_New(lz->r);
2522 if (result == NULL)
2523 return NULL;
2524 for (i=0; i<lz->r; i++) {
2525 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2526 Py_INCREF(element);
2527 PyTuple_SET_ITEM(result, i, element);
2528 }
2529
2530 Py_CLEAR(lz->result);
2531 lz->result = result;
2532 Py_RETURN_NONE;
2533}
2534
2535static PyMethodDef combinations_methods[] = {
2536 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2537 reduce_doc},
2538 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2539 setstate_doc},
2540 {NULL, NULL} /* sentinel */
2541};
2542
Christian Heimes380f7f22008-02-28 11:19:05 +00002543PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002544"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002545\n\
2546Return successive r-length combinations of elements in the iterable.\n\n\
2547combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2548
2549static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002550 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002551 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002552 sizeof(combinationsobject), /* tp_basicsize */
2553 0, /* tp_itemsize */
2554 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002555 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002556 0, /* tp_print */
2557 0, /* tp_getattr */
2558 0, /* tp_setattr */
2559 0, /* tp_reserved */
2560 0, /* tp_repr */
2561 0, /* tp_as_number */
2562 0, /* tp_as_sequence */
2563 0, /* tp_as_mapping */
2564 0, /* tp_hash */
2565 0, /* tp_call */
2566 0, /* tp_str */
2567 PyObject_GenericGetAttr, /* tp_getattro */
2568 0, /* tp_setattro */
2569 0, /* tp_as_buffer */
2570 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2571 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002572 combinations_doc, /* tp_doc */
2573 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002574 0, /* tp_clear */
2575 0, /* tp_richcompare */
2576 0, /* tp_weaklistoffset */
2577 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002578 (iternextfunc)combinations_next, /* tp_iternext */
2579 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002580 0, /* tp_members */
2581 0, /* tp_getset */
2582 0, /* tp_base */
2583 0, /* tp_dict */
2584 0, /* tp_descr_get */
2585 0, /* tp_descr_set */
2586 0, /* tp_dictoffset */
2587 0, /* tp_init */
2588 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002589 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002590 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002591};
2592
2593
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002594/* combinations with replacement object *******************************************/
2595
2596/* Equivalent to:
2597
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002598 def combinations_with_replacement(iterable, r):
2599 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2600 # number items returned: (n+r-1)! / r! / (n-1)!
2601 pool = tuple(iterable)
2602 n = len(pool)
2603 indices = [0] * r
2604 yield tuple(pool[i] for i in indices)
2605 while 1:
2606 for i in reversed(range(r)):
2607 if indices[i] != n - 1:
2608 break
2609 else:
2610 return
2611 indices[i:] = [indices[i] + 1] * (r - i)
2612 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002613
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002614 def combinations_with_replacement2(iterable, r):
2615 'Alternate version that filters from product()'
2616 pool = tuple(iterable)
2617 n = len(pool)
2618 for indices in product(range(n), repeat=r):
2619 if sorted(indices) == list(indices):
2620 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002621*/
2622typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002623 PyObject_HEAD
2624 PyObject *pool; /* input converted to a tuple */
2625 Py_ssize_t *indices; /* one index per result element */
2626 PyObject *result; /* most recently returned result tuple */
2627 Py_ssize_t r; /* size of result tuple */
2628 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002629} cwrobject;
2630
2631static PyTypeObject cwr_type;
2632
2633static PyObject *
2634cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2635{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002636 cwrobject *co;
2637 Py_ssize_t n;
2638 Py_ssize_t r;
2639 PyObject *pool = NULL;
2640 PyObject *iterable = NULL;
2641 Py_ssize_t *indices = NULL;
2642 Py_ssize_t i;
2643 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002644
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002645 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2646 &iterable, &r))
2647 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002648
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002649 pool = PySequence_Tuple(iterable);
2650 if (pool == NULL)
2651 goto error;
2652 n = PyTuple_GET_SIZE(pool);
2653 if (r < 0) {
2654 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2655 goto error;
2656 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002657
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002658 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2659 if (indices == NULL) {
2660 PyErr_NoMemory();
2661 goto error;
2662 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002663
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002664 for (i=0 ; i<r ; i++)
2665 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002666
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002667 /* create cwrobject structure */
2668 co = (cwrobject *)type->tp_alloc(type, 0);
2669 if (co == NULL)
2670 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002671
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002672 co->pool = pool;
2673 co->indices = indices;
2674 co->result = NULL;
2675 co->r = r;
2676 co->stopped = !n && r;
2677
2678 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002679
2680error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002681 if (indices != NULL)
2682 PyMem_Free(indices);
2683 Py_XDECREF(pool);
2684 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002685}
2686
2687static void
2688cwr_dealloc(cwrobject *co)
2689{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002690 PyObject_GC_UnTrack(co);
2691 Py_XDECREF(co->pool);
2692 Py_XDECREF(co->result);
2693 if (co->indices != NULL)
2694 PyMem_Free(co->indices);
2695 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002696}
2697
2698static int
2699cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2700{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002701 Py_VISIT(co->pool);
2702 Py_VISIT(co->result);
2703 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002704}
2705
2706static PyObject *
2707cwr_next(cwrobject *co)
2708{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002709 PyObject *elem;
2710 PyObject *oldelem;
2711 PyObject *pool = co->pool;
2712 Py_ssize_t *indices = co->indices;
2713 PyObject *result = co->result;
2714 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2715 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002716 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002717
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002718 if (co->stopped)
2719 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002720
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002721 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002722 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002723 result = PyTuple_New(r);
2724 if (result == NULL)
2725 goto empty;
2726 co->result = result;
Tim Peters9edb1682013-09-03 11:49:31 -05002727 elem = PyTuple_GET_ITEM(pool, 0);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002728 for (i=0; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002729 assert(indices[i] == 0);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002730 Py_INCREF(elem);
2731 PyTuple_SET_ITEM(result, i, elem);
2732 }
2733 } else {
2734 /* Copy the previous result tuple or re-use it if available */
2735 if (Py_REFCNT(result) > 1) {
2736 PyObject *old_result = result;
2737 result = PyTuple_New(r);
2738 if (result == NULL)
2739 goto empty;
2740 co->result = result;
2741 for (i=0; i<r ; i++) {
2742 elem = PyTuple_GET_ITEM(old_result, i);
2743 Py_INCREF(elem);
2744 PyTuple_SET_ITEM(result, i, elem);
2745 }
2746 Py_DECREF(old_result);
2747 }
2748 /* Now, we've got the only copy so we can update it in-place CPython's
2749 empty tuple is a singleton and cached in PyTuple's freelist. */
2750 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002751
Tim Peters9edb1682013-09-03 11:49:31 -05002752 /* Scan indices right-to-left until finding one that is not
2753 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002754 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2755 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002756
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002757 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002758 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002759 if (i < 0)
2760 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002761
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002762 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002763 maximum. Then set all to the right to the same value. */
2764 index = indices[i] + 1;
2765 assert(index < n);
2766 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002767 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002768 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002769 Py_INCREF(elem);
2770 oldelem = PyTuple_GET_ITEM(result, i);
2771 PyTuple_SET_ITEM(result, i, elem);
2772 Py_DECREF(oldelem);
2773 }
2774 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002775
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002776 Py_INCREF(result);
2777 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002778
2779empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002780 co->stopped = 1;
2781 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002782}
2783
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002784static PyObject *
2785cwr_reduce(cwrobject *lz)
2786{
2787 if (lz->result == NULL) {
2788 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2789 } else if (lz->stopped) {
2790 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2791 } else {
2792 PyObject *indices;
2793 Py_ssize_t i;
2794
2795 /* we must pickle the indices and use them for setstate */
2796 indices = PyTuple_New(lz->r);
2797 if (!indices)
2798 return NULL;
2799 for (i=0; i<lz->r; i++)
2800 {
2801 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2802 if (!index) {
2803 Py_DECREF(indices);
2804 return NULL;
2805 }
2806 PyTuple_SET_ITEM(indices, i, index);
2807 }
2808
2809 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2810 }
2811}
2812
2813static PyObject *
2814cwr_setstate(cwrobject *lz, PyObject *state)
2815{
2816 PyObject *result;
2817 Py_ssize_t n, i;
2818
2819 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2820 {
2821 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2822 return NULL;
2823 }
2824
2825 n = PyTuple_GET_SIZE(lz->pool);
2826 for (i=0; i<lz->r; i++)
2827 {
2828 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2829 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2830 if (index < 0 && PyErr_Occurred())
2831 return NULL; /* not an integer */
2832 /* clamp the index */
2833 if (index < 0)
2834 index = 0;
2835 else if (index > n-1)
2836 index = n-1;
2837 lz->indices[i] = index;
2838 }
2839 result = PyTuple_New(lz->r);
2840 if (result == NULL)
2841 return NULL;
2842 for (i=0; i<lz->r; i++) {
2843 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2844 Py_INCREF(element);
2845 PyTuple_SET_ITEM(result, i, element);
2846 }
2847 Py_CLEAR(lz->result);
2848 lz->result = result;
2849 Py_RETURN_NONE;
2850}
2851
2852static PyMethodDef cwr_methods[] = {
2853 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2854 reduce_doc},
2855 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2856 setstate_doc},
2857 {NULL, NULL} /* sentinel */
2858};
2859
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002860PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002861"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002862\n\
2863Return successive r-length combinations of elements in the iterable\n\
2864allowing individual elements to have successive repeats.\n\
2865combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2866
2867static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002868 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002869 "itertools.combinations_with_replacement", /* tp_name */
2870 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002871 0, /* tp_itemsize */
2872 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002873 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002874 0, /* tp_print */
2875 0, /* tp_getattr */
2876 0, /* tp_setattr */
2877 0, /* tp_reserved */
2878 0, /* tp_repr */
2879 0, /* tp_as_number */
2880 0, /* tp_as_sequence */
2881 0, /* tp_as_mapping */
2882 0, /* tp_hash */
2883 0, /* tp_call */
2884 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002885 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002886 0, /* tp_setattro */
2887 0, /* tp_as_buffer */
2888 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002889 Py_TPFLAGS_BASETYPE, /* tp_flags */
2890 cwr_doc, /* tp_doc */
2891 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002892 0, /* tp_clear */
2893 0, /* tp_richcompare */
2894 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002895 PyObject_SelfIter, /* tp_iter */
2896 (iternextfunc)cwr_next, /* tp_iternext */
2897 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002898 0, /* tp_members */
2899 0, /* tp_getset */
2900 0, /* tp_base */
2901 0, /* tp_dict */
2902 0, /* tp_descr_get */
2903 0, /* tp_descr_set */
2904 0, /* tp_dictoffset */
2905 0, /* tp_init */
2906 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002907 cwr_new, /* tp_new */
2908 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002909};
2910
2911
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002912/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002913
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002914def permutations(iterable, r=None):
2915 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2916 pool = tuple(iterable)
2917 n = len(pool)
2918 r = n if r is None else r
2919 indices = range(n)
2920 cycles = range(n-r+1, n+1)[::-1]
2921 yield tuple(pool[i] for i in indices[:r])
2922 while n:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002923 for i in reversed(range(r)):
2924 cycles[i] -= 1
2925 if cycles[i] == 0:
2926 indices[i:] = indices[i+1:] + indices[i:i+1]
2927 cycles[i] = n - i
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002928 else:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002929 j = cycles[i]
2930 indices[i], indices[-j] = indices[-j], indices[i]
2931 yield tuple(pool[i] for i in indices[:r])
2932 break
2933 else:
2934 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002935*/
2936
2937typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002938 PyObject_HEAD
2939 PyObject *pool; /* input converted to a tuple */
2940 Py_ssize_t *indices; /* one index per element in the pool */
2941 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2942 PyObject *result; /* most recently returned result tuple */
2943 Py_ssize_t r; /* size of result tuple */
2944 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002945} permutationsobject;
2946
2947static PyTypeObject permutations_type;
2948
2949static PyObject *
2950permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2951{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002952 permutationsobject *po;
2953 Py_ssize_t n;
2954 Py_ssize_t r;
2955 PyObject *robj = Py_None;
2956 PyObject *pool = NULL;
2957 PyObject *iterable = NULL;
2958 Py_ssize_t *indices = NULL;
2959 Py_ssize_t *cycles = NULL;
2960 Py_ssize_t i;
2961 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002962
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002963 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2964 &iterable, &robj))
2965 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002966
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002967 pool = PySequence_Tuple(iterable);
2968 if (pool == NULL)
2969 goto error;
2970 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002971
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002972 r = n;
2973 if (robj != Py_None) {
2974 if (!PyLong_Check(robj)) {
2975 PyErr_SetString(PyExc_TypeError, "Expected int as r");
2976 goto error;
2977 }
2978 r = PyLong_AsSsize_t(robj);
2979 if (r == -1 && PyErr_Occurred())
2980 goto error;
2981 }
2982 if (r < 0) {
2983 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2984 goto error;
2985 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002986
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002987 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2988 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2989 if (indices == NULL || cycles == NULL) {
2990 PyErr_NoMemory();
2991 goto error;
2992 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002993
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002994 for (i=0 ; i<n ; i++)
2995 indices[i] = i;
2996 for (i=0 ; i<r ; i++)
2997 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002998
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002999 /* create permutationsobject structure */
3000 po = (permutationsobject *)type->tp_alloc(type, 0);
3001 if (po == NULL)
3002 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003003
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003004 po->pool = pool;
3005 po->indices = indices;
3006 po->cycles = cycles;
3007 po->result = NULL;
3008 po->r = r;
3009 po->stopped = r > n ? 1 : 0;
3010
3011 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003012
3013error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003014 if (indices != NULL)
3015 PyMem_Free(indices);
3016 if (cycles != NULL)
3017 PyMem_Free(cycles);
3018 Py_XDECREF(pool);
3019 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003020}
3021
3022static void
3023permutations_dealloc(permutationsobject *po)
3024{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003025 PyObject_GC_UnTrack(po);
3026 Py_XDECREF(po->pool);
3027 Py_XDECREF(po->result);
3028 PyMem_Free(po->indices);
3029 PyMem_Free(po->cycles);
3030 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003031}
3032
3033static int
3034permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3035{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003036 Py_VISIT(po->pool);
3037 Py_VISIT(po->result);
3038 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003039}
3040
3041static PyObject *
3042permutations_next(permutationsobject *po)
3043{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003044 PyObject *elem;
3045 PyObject *oldelem;
3046 PyObject *pool = po->pool;
3047 Py_ssize_t *indices = po->indices;
3048 Py_ssize_t *cycles = po->cycles;
3049 PyObject *result = po->result;
3050 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3051 Py_ssize_t r = po->r;
3052 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003053
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003054 if (po->stopped)
3055 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003056
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003057 if (result == NULL) {
3058 /* On the first pass, initialize result tuple using the indices */
3059 result = PyTuple_New(r);
3060 if (result == NULL)
3061 goto empty;
3062 po->result = result;
3063 for (i=0; i<r ; i++) {
3064 index = indices[i];
3065 elem = PyTuple_GET_ITEM(pool, index);
3066 Py_INCREF(elem);
3067 PyTuple_SET_ITEM(result, i, elem);
3068 }
3069 } else {
3070 if (n == 0)
3071 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003072
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003073 /* Copy the previous result tuple or re-use it if available */
3074 if (Py_REFCNT(result) > 1) {
3075 PyObject *old_result = result;
3076 result = PyTuple_New(r);
3077 if (result == NULL)
3078 goto empty;
3079 po->result = result;
3080 for (i=0; i<r ; i++) {
3081 elem = PyTuple_GET_ITEM(old_result, i);
3082 Py_INCREF(elem);
3083 PyTuple_SET_ITEM(result, i, elem);
3084 }
3085 Py_DECREF(old_result);
3086 }
3087 /* Now, we've got the only copy so we can update it in-place */
3088 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003089
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003090 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3091 for (i=r-1 ; i>=0 ; i--) {
3092 cycles[i] -= 1;
3093 if (cycles[i] == 0) {
3094 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3095 index = indices[i];
3096 for (j=i ; j<n-1 ; j++)
3097 indices[j] = indices[j+1];
3098 indices[n-1] = index;
3099 cycles[i] = n - i;
3100 } else {
3101 j = cycles[i];
3102 index = indices[i];
3103 indices[i] = indices[n-j];
3104 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003105
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003106 for (k=i; k<r ; k++) {
3107 /* start with i, the leftmost element that changed */
3108 /* yield tuple(pool[k] for k in indices[:r]) */
3109 index = indices[k];
3110 elem = PyTuple_GET_ITEM(pool, index);
3111 Py_INCREF(elem);
3112 oldelem = PyTuple_GET_ITEM(result, k);
3113 PyTuple_SET_ITEM(result, k, elem);
3114 Py_DECREF(oldelem);
3115 }
3116 break;
3117 }
3118 }
3119 /* If i is negative, then the cycles have all
3120 rolled-over and we're done. */
3121 if (i < 0)
3122 goto empty;
3123 }
3124 Py_INCREF(result);
3125 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003126
3127empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003128 po->stopped = 1;
3129 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003130}
3131
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003132static PyObject *
3133permutations_reduce(permutationsobject *po)
3134{
3135 if (po->result == NULL) {
3136 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3137 } else if (po->stopped) {
3138 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3139 } else {
3140 PyObject *indices=NULL, *cycles=NULL;
3141 Py_ssize_t n, i;
3142
3143 /* we must pickle the indices and cycles and use them for setstate */
3144 n = PyTuple_GET_SIZE(po->pool);
3145 indices = PyTuple_New(n);
3146 if (indices == NULL)
3147 goto err;
3148 for (i=0; i<n; i++){
3149 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3150 if (!index)
3151 goto err;
3152 PyTuple_SET_ITEM(indices, i, index);
3153 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003154
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003155 cycles = PyTuple_New(po->r);
3156 if (cycles == NULL)
3157 goto err;
3158 for (i=0; i<po->r; i++)
3159 {
3160 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3161 if (!index)
3162 goto err;
3163 PyTuple_SET_ITEM(cycles, i, index);
3164 }
3165 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3166 po->pool, po->r,
3167 indices, cycles);
3168 err:
3169 Py_XDECREF(indices);
3170 Py_XDECREF(cycles);
3171 return NULL;
3172 }
3173}
3174
3175static PyObject *
3176permutations_setstate(permutationsobject *po, PyObject *state)
3177{
3178 PyObject *indices, *cycles, *result;
3179 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003180
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003181 if (!PyArg_ParseTuple(state, "O!O!",
3182 &PyTuple_Type, &indices,
3183 &PyTuple_Type, &cycles))
3184 return NULL;
3185
3186 n = PyTuple_GET_SIZE(po->pool);
3187 if (PyTuple_GET_SIZE(indices) != n ||
3188 PyTuple_GET_SIZE(cycles) != po->r)
3189 {
3190 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3191 return NULL;
3192 }
3193
3194 for (i=0; i<n; i++)
3195 {
3196 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3197 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3198 if (index < 0 && PyErr_Occurred())
3199 return NULL; /* not an integer */
3200 /* clamp the index */
3201 if (index < 0)
3202 index = 0;
3203 else if (index > n-1)
3204 index = n-1;
3205 po->indices[i] = index;
3206 }
3207
3208 for (i=0; i<po->r; i++)
3209 {
3210 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3211 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3212 if (index < 0 && PyErr_Occurred())
3213 return NULL; /* not an integer */
3214 if (index < 1)
3215 index = 1;
3216 else if (index > n-i)
3217 index = n-i;
3218 po->cycles[i] = index;
3219 }
3220 result = PyTuple_New(po->r);
3221 if (result == NULL)
3222 return NULL;
3223 for (i=0; i<po->r; i++) {
3224 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3225 Py_INCREF(element);
3226 PyTuple_SET_ITEM(result, i, element);
3227 }
3228 Py_CLEAR(po->result);
3229 po->result = result;
3230 Py_RETURN_NONE;
3231}
3232
3233static PyMethodDef permuations_methods[] = {
3234 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3235 reduce_doc},
3236 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3237 setstate_doc},
3238 {NULL, NULL} /* sentinel */
3239};
3240
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003241PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003242"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003243\n\
3244Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003245permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003246
3247static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003248 PyVarObject_HEAD_INIT(NULL, 0)
3249 "itertools.permutations", /* tp_name */
3250 sizeof(permutationsobject), /* tp_basicsize */
3251 0, /* tp_itemsize */
3252 /* methods */
3253 (destructor)permutations_dealloc, /* tp_dealloc */
3254 0, /* tp_print */
3255 0, /* tp_getattr */
3256 0, /* tp_setattr */
3257 0, /* tp_reserved */
3258 0, /* tp_repr */
3259 0, /* tp_as_number */
3260 0, /* tp_as_sequence */
3261 0, /* tp_as_mapping */
3262 0, /* tp_hash */
3263 0, /* tp_call */
3264 0, /* tp_str */
3265 PyObject_GenericGetAttr, /* tp_getattro */
3266 0, /* tp_setattro */
3267 0, /* tp_as_buffer */
3268 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3269 Py_TPFLAGS_BASETYPE, /* tp_flags */
3270 permutations_doc, /* tp_doc */
3271 (traverseproc)permutations_traverse, /* tp_traverse */
3272 0, /* tp_clear */
3273 0, /* tp_richcompare */
3274 0, /* tp_weaklistoffset */
3275 PyObject_SelfIter, /* tp_iter */
3276 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003277 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003278 0, /* tp_members */
3279 0, /* tp_getset */
3280 0, /* tp_base */
3281 0, /* tp_dict */
3282 0, /* tp_descr_get */
3283 0, /* tp_descr_set */
3284 0, /* tp_dictoffset */
3285 0, /* tp_init */
3286 0, /* tp_alloc */
3287 permutations_new, /* tp_new */
3288 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003289};
3290
Raymond Hettinger482ba772010-12-01 22:48:00 +00003291/* accumulate object ************************************************************/
3292
3293typedef struct {
3294 PyObject_HEAD
3295 PyObject *total;
3296 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003297 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003298} accumulateobject;
3299
3300static PyTypeObject accumulate_type;
3301
3302static PyObject *
3303accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3304{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003305 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003306 PyObject *iterable;
3307 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003308 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003309 accumulateobject *lz;
3310
Raymond Hettinger5d446132011-03-27 18:52:10 -07003311 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3312 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003313 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003314
3315 /* Get iterator. */
3316 it = PyObject_GetIter(iterable);
3317 if (it == NULL)
3318 return NULL;
3319
Raymond Hettinger482ba772010-12-01 22:48:00 +00003320 /* create accumulateobject structure */
3321 lz = (accumulateobject *)type->tp_alloc(type, 0);
3322 if (lz == NULL) {
3323 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003324 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003325 }
3326
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003327 if (binop != Py_None) {
3328 Py_XINCREF(binop);
3329 lz->binop = binop;
3330 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003331 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003332 lz->it = it;
3333 return (PyObject *)lz;
3334}
3335
3336static void
3337accumulate_dealloc(accumulateobject *lz)
3338{
3339 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003340 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003341 Py_XDECREF(lz->total);
3342 Py_XDECREF(lz->it);
3343 Py_TYPE(lz)->tp_free(lz);
3344}
3345
3346static int
3347accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3348{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003349 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003350 Py_VISIT(lz->it);
3351 Py_VISIT(lz->total);
3352 return 0;
3353}
3354
3355static PyObject *
3356accumulate_next(accumulateobject *lz)
3357{
3358 PyObject *val, *oldtotal, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003359
Raymond Hettinger482ba772010-12-01 22:48:00 +00003360 val = PyIter_Next(lz->it);
3361 if (val == NULL)
3362 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003363
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003364 if (lz->total == NULL) {
3365 Py_INCREF(val);
3366 lz->total = val;
3367 return lz->total;
3368 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003369
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003370 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003371 newtotal = PyNumber_Add(lz->total, val);
3372 else
3373 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003374 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003375 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003376 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003377
3378 oldtotal = lz->total;
3379 lz->total = newtotal;
3380 Py_DECREF(oldtotal);
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003381
Raymond Hettinger482ba772010-12-01 22:48:00 +00003382 Py_INCREF(newtotal);
3383 return newtotal;
3384}
3385
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003386static PyObject *
3387accumulate_reduce(accumulateobject *lz)
3388{
3389 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3390 lz->it, lz->binop?lz->binop:Py_None,
3391 lz->total?lz->total:Py_None);
3392 }
3393
3394static PyObject *
3395accumulate_setstate(accumulateobject *lz, PyObject *state)
3396{
3397 Py_CLEAR(lz->total);
3398 lz->total = state;
3399 Py_INCREF(lz->total);
3400 Py_RETURN_NONE;
3401}
3402
3403static PyMethodDef accumulate_methods[] = {
3404 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3405 reduce_doc},
3406 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3407 setstate_doc},
3408 {NULL, NULL} /* sentinel */
3409};
3410
Raymond Hettinger482ba772010-12-01 22:48:00 +00003411PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003412"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003413\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003414Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003415
3416static PyTypeObject accumulate_type = {
3417 PyVarObject_HEAD_INIT(NULL, 0)
3418 "itertools.accumulate", /* tp_name */
3419 sizeof(accumulateobject), /* tp_basicsize */
3420 0, /* tp_itemsize */
3421 /* methods */
3422 (destructor)accumulate_dealloc, /* tp_dealloc */
3423 0, /* tp_print */
3424 0, /* tp_getattr */
3425 0, /* tp_setattr */
3426 0, /* tp_reserved */
3427 0, /* tp_repr */
3428 0, /* tp_as_number */
3429 0, /* tp_as_sequence */
3430 0, /* tp_as_mapping */
3431 0, /* tp_hash */
3432 0, /* tp_call */
3433 0, /* tp_str */
3434 PyObject_GenericGetAttr, /* tp_getattro */
3435 0, /* tp_setattro */
3436 0, /* tp_as_buffer */
3437 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3438 Py_TPFLAGS_BASETYPE, /* tp_flags */
3439 accumulate_doc, /* tp_doc */
3440 (traverseproc)accumulate_traverse, /* tp_traverse */
3441 0, /* tp_clear */
3442 0, /* tp_richcompare */
3443 0, /* tp_weaklistoffset */
3444 PyObject_SelfIter, /* tp_iter */
3445 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003446 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003447 0, /* tp_members */
3448 0, /* tp_getset */
3449 0, /* tp_base */
3450 0, /* tp_dict */
3451 0, /* tp_descr_get */
3452 0, /* tp_descr_set */
3453 0, /* tp_dictoffset */
3454 0, /* tp_init */
3455 0, /* tp_alloc */
3456 accumulate_new, /* tp_new */
3457 PyObject_GC_Del, /* tp_free */
3458};
3459
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003460
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003461/* compress object ************************************************************/
3462
3463/* Equivalent to:
3464
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003465 def compress(data, selectors):
3466 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3467 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003468*/
3469
3470typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003471 PyObject_HEAD
3472 PyObject *data;
3473 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003474} compressobject;
3475
3476static PyTypeObject compress_type;
3477
3478static PyObject *
3479compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3480{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003481 PyObject *seq1, *seq2;
3482 PyObject *data=NULL, *selectors=NULL;
3483 compressobject *lz;
3484 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003485
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003486 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3487 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003488
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003489 data = PyObject_GetIter(seq1);
3490 if (data == NULL)
3491 goto fail;
3492 selectors = PyObject_GetIter(seq2);
3493 if (selectors == NULL)
3494 goto fail;
3495
3496 /* create compressobject structure */
3497 lz = (compressobject *)type->tp_alloc(type, 0);
3498 if (lz == NULL)
3499 goto fail;
3500 lz->data = data;
3501 lz->selectors = selectors;
3502 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003503
3504fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003505 Py_XDECREF(data);
3506 Py_XDECREF(selectors);
3507 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003508}
3509
3510static void
3511compress_dealloc(compressobject *lz)
3512{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003513 PyObject_GC_UnTrack(lz);
3514 Py_XDECREF(lz->data);
3515 Py_XDECREF(lz->selectors);
3516 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003517}
3518
3519static int
3520compress_traverse(compressobject *lz, visitproc visit, void *arg)
3521{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003522 Py_VISIT(lz->data);
3523 Py_VISIT(lz->selectors);
3524 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003525}
3526
3527static PyObject *
3528compress_next(compressobject *lz)
3529{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003530 PyObject *data = lz->data, *selectors = lz->selectors;
3531 PyObject *datum, *selector;
3532 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3533 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3534 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003535
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003536 while (1) {
3537 /* Steps: get datum, get selector, evaluate selector.
3538 Order is important (to match the pure python version
3539 in terms of which input gets a chance to raise an
3540 exception first).
3541 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003542
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003543 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003544 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003545 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003546
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003547 selector = selectornext(selectors);
3548 if (selector == NULL) {
3549 Py_DECREF(datum);
3550 return NULL;
3551 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003552
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003553 ok = PyObject_IsTrue(selector);
3554 Py_DECREF(selector);
3555 if (ok == 1)
3556 return datum;
3557 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003558 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003559 return NULL;
3560 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003561}
3562
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003563static PyObject *
3564compress_reduce(compressobject *lz)
3565{
3566 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3567 lz->data, lz->selectors);
3568 }
3569
3570static PyMethodDef compress_methods[] = {
3571 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3572 reduce_doc},
3573 {NULL, NULL} /* sentinel */
3574};
3575
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003576PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003577"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003578\n\
3579Return data elements corresponding to true selector elements.\n\
3580Forms a shorter iterator from selected data elements using the\n\
3581selectors to choose the data elements.");
3582
3583static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003584 PyVarObject_HEAD_INIT(NULL, 0)
3585 "itertools.compress", /* tp_name */
3586 sizeof(compressobject), /* tp_basicsize */
3587 0, /* tp_itemsize */
3588 /* methods */
3589 (destructor)compress_dealloc, /* tp_dealloc */
3590 0, /* tp_print */
3591 0, /* tp_getattr */
3592 0, /* tp_setattr */
3593 0, /* tp_reserved */
3594 0, /* tp_repr */
3595 0, /* tp_as_number */
3596 0, /* tp_as_sequence */
3597 0, /* tp_as_mapping */
3598 0, /* tp_hash */
3599 0, /* tp_call */
3600 0, /* tp_str */
3601 PyObject_GenericGetAttr, /* tp_getattro */
3602 0, /* tp_setattro */
3603 0, /* tp_as_buffer */
3604 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3605 Py_TPFLAGS_BASETYPE, /* tp_flags */
3606 compress_doc, /* tp_doc */
3607 (traverseproc)compress_traverse, /* tp_traverse */
3608 0, /* tp_clear */
3609 0, /* tp_richcompare */
3610 0, /* tp_weaklistoffset */
3611 PyObject_SelfIter, /* tp_iter */
3612 (iternextfunc)compress_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003613 compress_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003614 0, /* tp_members */
3615 0, /* tp_getset */
3616 0, /* tp_base */
3617 0, /* tp_dict */
3618 0, /* tp_descr_get */
3619 0, /* tp_descr_set */
3620 0, /* tp_dictoffset */
3621 0, /* tp_init */
3622 0, /* tp_alloc */
3623 compress_new, /* tp_new */
3624 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003625};
3626
3627
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003628/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003629
3630typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003631 PyObject_HEAD
3632 PyObject *func;
3633 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003634} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003635
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003636static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003637
3638static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003639filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003640{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003641 PyObject *func, *seq;
3642 PyObject *it;
3643 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003644
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003645 if (type == &filterfalse_type &&
3646 !_PyArg_NoKeywords("filterfalse()", kwds))
3647 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003648
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003649 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3650 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003651
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003652 /* Get iterator. */
3653 it = PyObject_GetIter(seq);
3654 if (it == NULL)
3655 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003656
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003657 /* create filterfalseobject structure */
3658 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3659 if (lz == NULL) {
3660 Py_DECREF(it);
3661 return NULL;
3662 }
3663 Py_INCREF(func);
3664 lz->func = func;
3665 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003666
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003667 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003668}
3669
3670static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003671filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003672{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003673 PyObject_GC_UnTrack(lz);
3674 Py_XDECREF(lz->func);
3675 Py_XDECREF(lz->it);
3676 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003677}
3678
3679static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003680filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003681{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003682 Py_VISIT(lz->it);
3683 Py_VISIT(lz->func);
3684 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003685}
3686
3687static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003688filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003689{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003690 PyObject *item;
3691 PyObject *it = lz->it;
3692 long ok;
3693 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003694
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003695 iternext = *Py_TYPE(it)->tp_iternext;
3696 for (;;) {
3697 item = iternext(it);
3698 if (item == NULL)
3699 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003700
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003701 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3702 ok = PyObject_IsTrue(item);
3703 } else {
3704 PyObject *good;
3705 good = PyObject_CallFunctionObjArgs(lz->func,
3706 item, NULL);
3707 if (good == NULL) {
3708 Py_DECREF(item);
3709 return NULL;
3710 }
3711 ok = PyObject_IsTrue(good);
3712 Py_DECREF(good);
3713 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003714 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003715 return item;
3716 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003717 if (ok < 0)
3718 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003719 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003720}
3721
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003722static PyObject *
3723filterfalse_reduce(filterfalseobject *lz)
3724{
3725 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3726 lz->func, lz->it);
3727 }
3728
3729static PyMethodDef filterfalse_methods[] = {
3730 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3731 reduce_doc},
3732 {NULL, NULL} /* sentinel */
3733};
3734
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003735PyDoc_STRVAR(filterfalse_doc,
3736"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003737\n\
3738Return those items of sequence for which function(item) is false.\n\
3739If function is None, return the items that are false.");
3740
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003741static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003742 PyVarObject_HEAD_INIT(NULL, 0)
3743 "itertools.filterfalse", /* tp_name */
3744 sizeof(filterfalseobject), /* tp_basicsize */
3745 0, /* tp_itemsize */
3746 /* methods */
3747 (destructor)filterfalse_dealloc, /* tp_dealloc */
3748 0, /* tp_print */
3749 0, /* tp_getattr */
3750 0, /* tp_setattr */
3751 0, /* tp_reserved */
3752 0, /* tp_repr */
3753 0, /* tp_as_number */
3754 0, /* tp_as_sequence */
3755 0, /* tp_as_mapping */
3756 0, /* tp_hash */
3757 0, /* tp_call */
3758 0, /* tp_str */
3759 PyObject_GenericGetAttr, /* tp_getattro */
3760 0, /* tp_setattro */
3761 0, /* tp_as_buffer */
3762 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3763 Py_TPFLAGS_BASETYPE, /* tp_flags */
3764 filterfalse_doc, /* tp_doc */
3765 (traverseproc)filterfalse_traverse, /* tp_traverse */
3766 0, /* tp_clear */
3767 0, /* tp_richcompare */
3768 0, /* tp_weaklistoffset */
3769 PyObject_SelfIter, /* tp_iter */
3770 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003771 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003772 0, /* tp_members */
3773 0, /* tp_getset */
3774 0, /* tp_base */
3775 0, /* tp_dict */
3776 0, /* tp_descr_get */
3777 0, /* tp_descr_set */
3778 0, /* tp_dictoffset */
3779 0, /* tp_init */
3780 0, /* tp_alloc */
3781 filterfalse_new, /* tp_new */
3782 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003783};
3784
3785
3786/* count object ************************************************************/
3787
3788typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003789 PyObject_HEAD
3790 Py_ssize_t cnt;
3791 PyObject *long_cnt;
3792 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003793} countobject;
3794
Raymond Hettinger30729212009-02-12 06:28:27 +00003795/* Counting logic and invariants:
3796
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003797fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003798
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003799 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3800 Advances with: cnt += 1
3801 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003802
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003803slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003804
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003805 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3806 All counting is done with python objects (no overflows or underflows).
3807 Advances with: long_cnt += long_step
3808 Step may be zero -- effectively a slow version of repeat(cnt).
3809 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003810*/
3811
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003812static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003813
3814static PyObject *
3815count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3816{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003817 countobject *lz;
3818 int slow_mode = 0;
3819 Py_ssize_t cnt = 0;
3820 PyObject *long_cnt = NULL;
3821 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003822 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003823 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003824
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003825 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3826 kwlist, &long_cnt, &long_step))
3827 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003828
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003829 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3830 (long_step != NULL && !PyNumber_Check(long_step))) {
3831 PyErr_SetString(PyExc_TypeError, "a number is required");
3832 return NULL;
3833 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003834
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003835 if (long_cnt != NULL) {
3836 cnt = PyLong_AsSsize_t(long_cnt);
3837 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3838 PyErr_Clear();
3839 slow_mode = 1;
3840 }
3841 Py_INCREF(long_cnt);
3842 } else {
3843 cnt = 0;
3844 long_cnt = PyLong_FromLong(0);
3845 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003846
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003847 /* If not specified, step defaults to 1 */
3848 if (long_step == NULL) {
3849 long_step = PyLong_FromLong(1);
3850 if (long_step == NULL) {
3851 Py_DECREF(long_cnt);
3852 return NULL;
3853 }
3854 } else
3855 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003856
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003857 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003858
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003859 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003860 step = PyLong_AsLong(long_step);
3861 if (step != 1) {
3862 slow_mode = 1;
3863 if (step == -1 && PyErr_Occurred())
3864 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003865 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003866
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003867 if (slow_mode)
3868 cnt = PY_SSIZE_T_MAX;
3869 else
3870 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003871
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003872 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3873 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3874 assert(slow_mode ||
3875 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003876
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003877 /* create countobject structure */
3878 lz = (countobject *)type->tp_alloc(type, 0);
3879 if (lz == NULL) {
3880 Py_XDECREF(long_cnt);
3881 return NULL;
3882 }
3883 lz->cnt = cnt;
3884 lz->long_cnt = long_cnt;
3885 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003886
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003887 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003888}
3889
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003890static void
3891count_dealloc(countobject *lz)
3892{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003893 PyObject_GC_UnTrack(lz);
3894 Py_XDECREF(lz->long_cnt);
3895 Py_XDECREF(lz->long_step);
3896 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003897}
3898
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003899static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003900count_traverse(countobject *lz, visitproc visit, void *arg)
3901{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003902 Py_VISIT(lz->long_cnt);
3903 Py_VISIT(lz->long_step);
3904 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003905}
3906
3907static PyObject *
3908count_nextlong(countobject *lz)
3909{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003910 PyObject *long_cnt;
3911 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003912
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003913 long_cnt = lz->long_cnt;
3914 if (long_cnt == NULL) {
3915 /* Switch to slow_mode */
3916 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3917 if (long_cnt == NULL)
3918 return NULL;
3919 }
3920 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00003921
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003922 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3923 if (stepped_up == NULL)
3924 return NULL;
3925 lz->long_cnt = stepped_up;
3926 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003927}
3928
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003929static PyObject *
3930count_next(countobject *lz)
3931{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003932 if (lz->cnt == PY_SSIZE_T_MAX)
3933 return count_nextlong(lz);
3934 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003935}
3936
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003937static PyObject *
3938count_repr(countobject *lz)
3939{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003940 if (lz->cnt != PY_SSIZE_T_MAX)
3941 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003942
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003943 if (PyLong_Check(lz->long_step)) {
3944 long step = PyLong_AsLong(lz->long_step);
3945 if (step == -1 && PyErr_Occurred()) {
3946 PyErr_Clear();
3947 }
3948 if (step == 1) {
3949 /* Don't display step when it is an integer equal to 1 */
3950 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
3951 }
3952 }
3953 return PyUnicode_FromFormat("count(%R, %R)",
3954 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003955}
3956
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003957static PyObject *
3958count_reduce(countobject *lz)
3959{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003960 if (lz->cnt == PY_SSIZE_T_MAX)
3961 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3962 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003963}
3964
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003965static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003966 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003967 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003968 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003969};
3970
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003971PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00003972 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003973\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003974Return a count object whose .__next__() method returns consecutive values.\n\
3975Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00003976 def count(firstval=0, step=1):\n\
Eli Benderskyc554f722013-09-03 06:37:19 -07003977 x = firstval\n\
3978 while 1:\n\
3979 yield x\n\
3980 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003981
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003982static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003983 PyVarObject_HEAD_INIT(NULL, 0)
3984 "itertools.count", /* tp_name */
3985 sizeof(countobject), /* tp_basicsize */
3986 0, /* tp_itemsize */
3987 /* methods */
3988 (destructor)count_dealloc, /* tp_dealloc */
3989 0, /* tp_print */
3990 0, /* tp_getattr */
3991 0, /* tp_setattr */
3992 0, /* tp_reserved */
3993 (reprfunc)count_repr, /* tp_repr */
3994 0, /* tp_as_number */
3995 0, /* tp_as_sequence */
3996 0, /* tp_as_mapping */
3997 0, /* tp_hash */
3998 0, /* tp_call */
3999 0, /* tp_str */
4000 PyObject_GenericGetAttr, /* tp_getattro */
4001 0, /* tp_setattro */
4002 0, /* tp_as_buffer */
4003 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4004 Py_TPFLAGS_BASETYPE, /* tp_flags */
4005 count_doc, /* tp_doc */
4006 (traverseproc)count_traverse, /* tp_traverse */
4007 0, /* tp_clear */
4008 0, /* tp_richcompare */
4009 0, /* tp_weaklistoffset */
4010 PyObject_SelfIter, /* tp_iter */
4011 (iternextfunc)count_next, /* tp_iternext */
4012 count_methods, /* tp_methods */
4013 0, /* tp_members */
4014 0, /* tp_getset */
4015 0, /* tp_base */
4016 0, /* tp_dict */
4017 0, /* tp_descr_get */
4018 0, /* tp_descr_set */
4019 0, /* tp_dictoffset */
4020 0, /* tp_init */
4021 0, /* tp_alloc */
4022 count_new, /* tp_new */
4023 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004024};
4025
4026
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004027/* repeat object ************************************************************/
4028
4029typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004030 PyObject_HEAD
4031 PyObject *element;
4032 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004033} repeatobject;
4034
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004035static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004036
4037static PyObject *
4038repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4039{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004040 repeatobject *ro;
4041 PyObject *element;
4042 Py_ssize_t cnt = -1;
4043 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004044
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004045 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4046 &element, &cnt))
4047 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004048
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004049 if (PyTuple_Size(args) == 2 && cnt < 0)
4050 cnt = 0;
4051
4052 ro = (repeatobject *)type->tp_alloc(type, 0);
4053 if (ro == NULL)
4054 return NULL;
4055 Py_INCREF(element);
4056 ro->element = element;
4057 ro->cnt = cnt;
4058 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004059}
4060
4061static void
4062repeat_dealloc(repeatobject *ro)
4063{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004064 PyObject_GC_UnTrack(ro);
4065 Py_XDECREF(ro->element);
4066 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004067}
4068
4069static int
4070repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4071{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004072 Py_VISIT(ro->element);
4073 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004074}
4075
4076static PyObject *
4077repeat_next(repeatobject *ro)
4078{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004079 if (ro->cnt == 0)
4080 return NULL;
4081 if (ro->cnt > 0)
4082 ro->cnt--;
4083 Py_INCREF(ro->element);
4084 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004085}
4086
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004087static PyObject *
4088repeat_repr(repeatobject *ro)
4089{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004090 if (ro->cnt == -1)
4091 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4092 else
4093 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4094}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004095
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004096static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004097repeat_len(repeatobject *ro)
4098{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004099 if (ro->cnt == -1) {
4100 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4101 return NULL;
4102 }
4103 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004104}
4105
Armin Rigof5b3e362006-02-11 21:32:43 +00004106PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004107
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004108static PyObject *
4109repeat_reduce(repeatobject *ro)
4110{
4111 /* unpickle this so that a new repeat iterator is constructed with an
4112 * object, then call __setstate__ on it to set cnt
4113 */
4114 if (ro->cnt >= 0)
4115 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4116 else
4117 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4118}
4119
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004120static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004121 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004122 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004123 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004124};
4125
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004126PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004127"repeat(object [,times]) -> create an iterator which returns the object\n\
4128for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004129endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004130
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004131static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004132 PyVarObject_HEAD_INIT(NULL, 0)
4133 "itertools.repeat", /* tp_name */
4134 sizeof(repeatobject), /* tp_basicsize */
4135 0, /* tp_itemsize */
4136 /* methods */
4137 (destructor)repeat_dealloc, /* tp_dealloc */
4138 0, /* tp_print */
4139 0, /* tp_getattr */
4140 0, /* tp_setattr */
4141 0, /* tp_reserved */
4142 (reprfunc)repeat_repr, /* tp_repr */
4143 0, /* tp_as_number */
4144 0, /* tp_as_sequence */
4145 0, /* tp_as_mapping */
4146 0, /* tp_hash */
4147 0, /* tp_call */
4148 0, /* tp_str */
4149 PyObject_GenericGetAttr, /* tp_getattro */
4150 0, /* tp_setattro */
4151 0, /* tp_as_buffer */
4152 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4153 Py_TPFLAGS_BASETYPE, /* tp_flags */
4154 repeat_doc, /* tp_doc */
4155 (traverseproc)repeat_traverse, /* tp_traverse */
4156 0, /* tp_clear */
4157 0, /* tp_richcompare */
4158 0, /* tp_weaklistoffset */
4159 PyObject_SelfIter, /* tp_iter */
4160 (iternextfunc)repeat_next, /* tp_iternext */
4161 repeat_methods, /* tp_methods */
4162 0, /* tp_members */
4163 0, /* tp_getset */
4164 0, /* tp_base */
4165 0, /* tp_dict */
4166 0, /* tp_descr_get */
4167 0, /* tp_descr_set */
4168 0, /* tp_dictoffset */
4169 0, /* tp_init */
4170 0, /* tp_alloc */
4171 repeat_new, /* tp_new */
4172 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004173};
4174
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004175/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004176
4177#include "Python.h"
4178
4179typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004180 PyObject_HEAD
4181 Py_ssize_t tuplesize;
4182 Py_ssize_t numactive;
4183 PyObject *ittuple; /* tuple of iterators */
4184 PyObject *result;
4185 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004186} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004187
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004188static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004189
4190static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004191zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004192{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004193 ziplongestobject *lz;
4194 Py_ssize_t i;
4195 PyObject *ittuple; /* tuple of iterators */
4196 PyObject *result;
4197 PyObject *fillvalue = Py_None;
4198 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004199
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004200 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
4201 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4202 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
4203 PyErr_SetString(PyExc_TypeError,
4204 "zip_longest() got an unexpected keyword argument");
4205 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004206 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004207 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004208
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004209 /* args must be a tuple */
4210 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00004211
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004212 /* obtain iterators */
4213 ittuple = PyTuple_New(tuplesize);
4214 if (ittuple == NULL)
4215 return NULL;
4216 for (i=0; i < tuplesize; ++i) {
4217 PyObject *item = PyTuple_GET_ITEM(args, i);
4218 PyObject *it = PyObject_GetIter(item);
4219 if (it == NULL) {
4220 if (PyErr_ExceptionMatches(PyExc_TypeError))
4221 PyErr_Format(PyExc_TypeError,
4222 "zip_longest argument #%zd must support iteration",
4223 i+1);
4224 Py_DECREF(ittuple);
4225 return NULL;
4226 }
4227 PyTuple_SET_ITEM(ittuple, i, it);
4228 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004229
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004230 /* create a result holder */
4231 result = PyTuple_New(tuplesize);
4232 if (result == NULL) {
4233 Py_DECREF(ittuple);
4234 return NULL;
4235 }
4236 for (i=0 ; i < tuplesize ; i++) {
4237 Py_INCREF(Py_None);
4238 PyTuple_SET_ITEM(result, i, Py_None);
4239 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004240
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004241 /* create ziplongestobject structure */
4242 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4243 if (lz == NULL) {
4244 Py_DECREF(ittuple);
4245 Py_DECREF(result);
4246 return NULL;
4247 }
4248 lz->ittuple = ittuple;
4249 lz->tuplesize = tuplesize;
4250 lz->numactive = tuplesize;
4251 lz->result = result;
4252 Py_INCREF(fillvalue);
4253 lz->fillvalue = fillvalue;
4254 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004255}
4256
4257static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004258zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004259{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004260 PyObject_GC_UnTrack(lz);
4261 Py_XDECREF(lz->ittuple);
4262 Py_XDECREF(lz->result);
4263 Py_XDECREF(lz->fillvalue);
4264 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004265}
4266
4267static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004268zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004269{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004270 Py_VISIT(lz->ittuple);
4271 Py_VISIT(lz->result);
4272 Py_VISIT(lz->fillvalue);
4273 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004274}
4275
4276static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004277zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004278{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004279 Py_ssize_t i;
4280 Py_ssize_t tuplesize = lz->tuplesize;
4281 PyObject *result = lz->result;
4282 PyObject *it;
4283 PyObject *item;
4284 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004285
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004286 if (tuplesize == 0)
4287 return NULL;
4288 if (lz->numactive == 0)
4289 return NULL;
4290 if (Py_REFCNT(result) == 1) {
4291 Py_INCREF(result);
4292 for (i=0 ; i < tuplesize ; i++) {
4293 it = PyTuple_GET_ITEM(lz->ittuple, i);
4294 if (it == NULL) {
4295 Py_INCREF(lz->fillvalue);
4296 item = lz->fillvalue;
4297 } else {
4298 item = PyIter_Next(it);
4299 if (item == NULL) {
4300 lz->numactive -= 1;
4301 if (lz->numactive == 0 || PyErr_Occurred()) {
4302 lz->numactive = 0;
4303 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004304 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004305 } else {
4306 Py_INCREF(lz->fillvalue);
4307 item = lz->fillvalue;
4308 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4309 Py_DECREF(it);
4310 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004311 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004312 }
4313 olditem = PyTuple_GET_ITEM(result, i);
4314 PyTuple_SET_ITEM(result, i, item);
4315 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004316 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004317 } else {
4318 result = PyTuple_New(tuplesize);
4319 if (result == NULL)
4320 return NULL;
4321 for (i=0 ; i < tuplesize ; i++) {
4322 it = PyTuple_GET_ITEM(lz->ittuple, i);
4323 if (it == NULL) {
4324 Py_INCREF(lz->fillvalue);
4325 item = lz->fillvalue;
4326 } else {
4327 item = PyIter_Next(it);
4328 if (item == NULL) {
4329 lz->numactive -= 1;
4330 if (lz->numactive == 0 || PyErr_Occurred()) {
4331 lz->numactive = 0;
4332 Py_DECREF(result);
4333 return NULL;
4334 } else {
4335 Py_INCREF(lz->fillvalue);
4336 item = lz->fillvalue;
4337 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4338 Py_DECREF(it);
4339 }
4340 }
4341 }
4342 PyTuple_SET_ITEM(result, i, item);
4343 }
4344 }
4345 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004346}
4347
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004348static PyObject *
4349zip_longest_reduce(ziplongestobject *lz)
4350{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004351
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004352 /* Create a new tuple with empty sequences where appropriate to pickle.
4353 * Then use setstate to set the fillvalue
4354 */
4355 int i;
4356 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
4357 if (args == NULL)
4358 return NULL;
4359 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4360 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4361 if (elem == NULL) {
4362 elem = PyTuple_New(0);
4363 if (elem == NULL) {
4364 Py_DECREF(args);
4365 return NULL;
4366 }
4367 } else
4368 Py_INCREF(elem);
4369 PyTuple_SET_ITEM(args, i, elem);
4370 }
4371 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4372}
4373
4374static PyObject *
4375zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4376{
4377 Py_CLEAR(lz->fillvalue);
4378 lz->fillvalue = state;
4379 Py_INCREF(lz->fillvalue);
4380 Py_RETURN_NONE;
4381}
4382
4383static PyMethodDef zip_longest_methods[] = {
4384 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4385 reduce_doc},
4386 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4387 setstate_doc},
4388 {NULL, NULL} /* sentinel */
4389};
4390
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004391PyDoc_STRVAR(zip_longest_doc,
4392"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004393\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004394Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004395the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004396method continues until the longest iterable in the argument sequence\n\
4397is exhausted and then it raises StopIteration. When the shorter iterables\n\
4398are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4399defaults to None or can be specified by a keyword argument.\n\
4400");
4401
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004402static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004403 PyVarObject_HEAD_INIT(NULL, 0)
4404 "itertools.zip_longest", /* tp_name */
4405 sizeof(ziplongestobject), /* tp_basicsize */
4406 0, /* tp_itemsize */
4407 /* methods */
4408 (destructor)zip_longest_dealloc, /* tp_dealloc */
4409 0, /* tp_print */
4410 0, /* tp_getattr */
4411 0, /* tp_setattr */
4412 0, /* tp_reserved */
4413 0, /* tp_repr */
4414 0, /* tp_as_number */
4415 0, /* tp_as_sequence */
4416 0, /* tp_as_mapping */
4417 0, /* tp_hash */
4418 0, /* tp_call */
4419 0, /* tp_str */
4420 PyObject_GenericGetAttr, /* tp_getattro */
4421 0, /* tp_setattro */
4422 0, /* tp_as_buffer */
4423 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4424 Py_TPFLAGS_BASETYPE, /* tp_flags */
4425 zip_longest_doc, /* tp_doc */
4426 (traverseproc)zip_longest_traverse, /* tp_traverse */
4427 0, /* tp_clear */
4428 0, /* tp_richcompare */
4429 0, /* tp_weaklistoffset */
4430 PyObject_SelfIter, /* tp_iter */
4431 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004432 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004433 0, /* tp_members */
4434 0, /* tp_getset */
4435 0, /* tp_base */
4436 0, /* tp_dict */
4437 0, /* tp_descr_get */
4438 0, /* tp_descr_set */
4439 0, /* tp_dictoffset */
4440 0, /* tp_init */
4441 0, /* tp_alloc */
4442 zip_longest_new, /* tp_new */
4443 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004444};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004445
4446/* module level code ********************************************************/
4447
4448PyDoc_STRVAR(module_doc,
4449"Functional tools for creating and using iterators.\n\
4450\n\
4451Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004452count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004453cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004454repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004455\n\
4456Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004457accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004458chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettingerfb92f392013-09-09 02:01:35 -05004459chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004460compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4461dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4462groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004463filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004464islice(seq, [start,] stop [, step]) --> elements from\n\
4465 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004466starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004467tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004468takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004469zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004470\n\
4471Combinatoric generators:\n\
4472product(p, q, ... [repeat=1]) --> cartesian product\n\
4473permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004474combinations(p, r)\n\
4475combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004476");
4477
4478
Raymond Hettingerad983e72003-11-12 14:32:26 +00004479static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004480 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4481 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004482};
4483
Martin v. Löwis1a214512008-06-11 05:26:20 +00004484
4485static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004486 PyModuleDef_HEAD_INIT,
4487 "itertools",
4488 module_doc,
4489 -1,
4490 module_methods,
4491 NULL,
4492 NULL,
4493 NULL,
4494 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004495};
4496
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004497PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004498PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004499{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004500 int i;
4501 PyObject *m;
4502 char *name;
4503 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004504 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004505 &combinations_type,
4506 &cwr_type,
4507 &cycle_type,
4508 &dropwhile_type,
4509 &takewhile_type,
4510 &islice_type,
4511 &starmap_type,
4512 &chain_type,
4513 &compress_type,
4514 &filterfalse_type,
4515 &count_type,
4516 &ziplongest_type,
4517 &permutations_type,
4518 &product_type,
4519 &repeat_type,
4520 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004521 &_grouper_type,
4522 &tee_type,
4523 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004524 NULL
4525 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004526
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004527 Py_TYPE(&teedataobject_type) = &PyType_Type;
4528 m = PyModule_Create(&itertoolsmodule);
4529 if (m == NULL)
4530 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004531
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004532 for (i=0 ; typelist[i] != NULL ; i++) {
4533 if (PyType_Ready(typelist[i]) < 0)
4534 return NULL;
4535 name = strchr(typelist[i]->tp_name, '.');
4536 assert (name != NULL);
4537 Py_INCREF(typelist[i]);
4538 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4539 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004540
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004541 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004542}