blob: ff7867d5bcb03ec4e70a6dc6f9e67eedbb4a39f2 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
Raymond Hettingerca3788c2015-08-16 14:49:24 -07002#define PY_SSIZE_T_CLEAN
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003#include "Python.h"
Fred Drake08ebfec2004-10-17 19:36:57 +00004#include "structmember.h"
Raymond Hettinger96ef8112003-02-01 00:10:11 +00005
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00006/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00007 by Raymond D. Hettinger <python@rcn.com>
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07008 Copyright (c) 2003-2015 Python Software Foundation.
Raymond Hettinger96ef8112003-02-01 00:10:11 +00009 All rights reserved.
10*/
11
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000012
Raymond Hettingera6ea44a2015-08-17 23:55:28 -070013/* groupby object ************************************************************/
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000014
15typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000016 PyObject_HEAD
17 PyObject *it;
18 PyObject *keyfunc;
19 PyObject *tgtkey;
20 PyObject *currkey;
21 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000022} groupbyobject;
23
24static PyTypeObject groupby_type;
25static PyObject *_grouper_create(groupbyobject *, PyObject *);
26
27static PyObject *
28groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
29{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000030 static char *kwargs[] = {"iterable", "key", NULL};
31 groupbyobject *gbo;
32 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000033
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000034 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
35 &it, &keyfunc))
36 return NULL;
37
38 gbo = (groupbyobject *)type->tp_alloc(type, 0);
39 if (gbo == NULL)
40 return NULL;
41 gbo->tgtkey = NULL;
42 gbo->currkey = NULL;
43 gbo->currvalue = NULL;
44 gbo->keyfunc = keyfunc;
45 Py_INCREF(keyfunc);
46 gbo->it = PyObject_GetIter(it);
47 if (gbo->it == NULL) {
48 Py_DECREF(gbo);
49 return NULL;
50 }
51 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000052}
53
54static void
55groupby_dealloc(groupbyobject *gbo)
56{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000057 PyObject_GC_UnTrack(gbo);
58 Py_XDECREF(gbo->it);
59 Py_XDECREF(gbo->keyfunc);
60 Py_XDECREF(gbo->tgtkey);
61 Py_XDECREF(gbo->currkey);
62 Py_XDECREF(gbo->currvalue);
63 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000064}
65
66static int
67groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
68{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000069 Py_VISIT(gbo->it);
70 Py_VISIT(gbo->keyfunc);
71 Py_VISIT(gbo->tgtkey);
72 Py_VISIT(gbo->currkey);
73 Py_VISIT(gbo->currvalue);
74 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000075}
76
77static PyObject *
78groupby_next(groupbyobject *gbo)
79{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000080 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000081
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000082 /* skip to next iteration group */
83 for (;;) {
84 if (gbo->currkey == NULL)
85 /* pass */;
86 else if (gbo->tgtkey == NULL)
87 break;
88 else {
89 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000090
Raymond Hettingera6ea44a2015-08-17 23:55:28 -070091 rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000092 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 {
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700106 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000107 if (newkey == NULL) {
108 Py_DECREF(newvalue);
109 return NULL;
110 }
111 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000112
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000113 tmp = gbo->currkey;
114 gbo->currkey = newkey;
115 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000116
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000117 tmp = gbo->currvalue;
118 gbo->currvalue = newvalue;
119 Py_XDECREF(tmp);
120 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000121
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000122 Py_INCREF(gbo->currkey);
123 tmp = gbo->tgtkey;
124 gbo->tgtkey = gbo->currkey;
125 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000126
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000127 grouper = _grouper_create(gbo, gbo->tgtkey);
128 if (grouper == NULL)
129 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000130
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000131 r = PyTuple_Pack(2, gbo->currkey, grouper);
132 Py_DECREF(grouper);
133 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000134}
135
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000136static PyObject *
137groupby_reduce(groupbyobject *lz)
138{
139 /* reduce as a 'new' call with an optional 'setstate' if groupby
140 * has started
141 */
142 PyObject *value;
143 if (lz->tgtkey && lz->currkey && lz->currvalue)
144 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
145 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
146 else
147 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
148 lz->it, lz->keyfunc);
149
150 return value;
151}
152
153PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
154
155static PyObject *
156groupby_setstate(groupbyobject *lz, PyObject *state)
157{
158 PyObject *currkey, *currvalue, *tgtkey;
159 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey))
160 return NULL;
161 Py_CLEAR(lz->currkey);
162 lz->currkey = currkey;
163 Py_INCREF(lz->currkey);
164 Py_CLEAR(lz->currvalue);
165 lz->currvalue = currvalue;
166 Py_INCREF(lz->currvalue);
167 Py_CLEAR(lz->tgtkey);
168 lz->tgtkey = tgtkey;
169 Py_INCREF(lz->tgtkey);
170 Py_RETURN_NONE;
171}
172
173PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
174
175static PyMethodDef groupby_methods[] = {
176 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
177 reduce_doc},
178 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
179 setstate_doc},
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700180 {NULL, NULL} /* sentinel */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000181};
182
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000183PyDoc_STRVAR(groupby_doc,
184"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
185(key, sub-iterator) grouped by each value of key(value).\n");
186
187static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000188 PyVarObject_HEAD_INIT(NULL, 0)
189 "itertools.groupby", /* tp_name */
190 sizeof(groupbyobject), /* tp_basicsize */
191 0, /* tp_itemsize */
192 /* methods */
193 (destructor)groupby_dealloc, /* tp_dealloc */
194 0, /* tp_print */
195 0, /* tp_getattr */
196 0, /* tp_setattr */
197 0, /* tp_reserved */
198 0, /* tp_repr */
199 0, /* tp_as_number */
200 0, /* tp_as_sequence */
201 0, /* tp_as_mapping */
202 0, /* tp_hash */
203 0, /* tp_call */
204 0, /* tp_str */
205 PyObject_GenericGetAttr, /* tp_getattro */
206 0, /* tp_setattro */
207 0, /* tp_as_buffer */
208 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
209 Py_TPFLAGS_BASETYPE, /* tp_flags */
210 groupby_doc, /* tp_doc */
211 (traverseproc)groupby_traverse, /* tp_traverse */
212 0, /* tp_clear */
213 0, /* tp_richcompare */
214 0, /* tp_weaklistoffset */
215 PyObject_SelfIter, /* tp_iter */
216 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000217 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000218 0, /* tp_members */
219 0, /* tp_getset */
220 0, /* tp_base */
221 0, /* tp_dict */
222 0, /* tp_descr_get */
223 0, /* tp_descr_set */
224 0, /* tp_dictoffset */
225 0, /* tp_init */
226 0, /* tp_alloc */
227 groupby_new, /* tp_new */
228 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000229};
230
231
232/* _grouper object (internal) ************************************************/
233
234typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000235 PyObject_HEAD
236 PyObject *parent;
237 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000238} _grouperobject;
239
240static PyTypeObject _grouper_type;
241
242static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000243_grouper_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
244{
245 PyObject *parent, *tgtkey;
246
247 if (!PyArg_ParseTuple(args, "O!O", &groupby_type, &parent, &tgtkey))
248 return NULL;
249
250 return _grouper_create((groupbyobject*) parent, tgtkey);
251}
252
253static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000254_grouper_create(groupbyobject *parent, PyObject *tgtkey)
255{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000256 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000257
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000258 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
259 if (igo == NULL)
260 return NULL;
261 igo->parent = (PyObject *)parent;
262 Py_INCREF(parent);
263 igo->tgtkey = tgtkey;
264 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000265
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000266 PyObject_GC_Track(igo);
267 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000268}
269
270static void
271_grouper_dealloc(_grouperobject *igo)
272{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000273 PyObject_GC_UnTrack(igo);
274 Py_DECREF(igo->parent);
275 Py_DECREF(igo->tgtkey);
276 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000277}
278
279static int
280_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
281{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000282 Py_VISIT(igo->parent);
283 Py_VISIT(igo->tgtkey);
284 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000285}
286
287static PyObject *
288_grouper_next(_grouperobject *igo)
289{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000290 groupbyobject *gbo = (groupbyobject *)igo->parent;
291 PyObject *newvalue, *newkey, *r;
292 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000293
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000294 if (gbo->currvalue == NULL) {
295 newvalue = PyIter_Next(gbo->it);
296 if (newvalue == NULL)
297 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000298
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000299 if (gbo->keyfunc == Py_None) {
300 newkey = newvalue;
301 Py_INCREF(newvalue);
302 } else {
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700303 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000304 if (newkey == NULL) {
305 Py_DECREF(newvalue);
306 return NULL;
307 }
308 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000309
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000310 assert(gbo->currkey == NULL);
311 gbo->currkey = newkey;
312 gbo->currvalue = newvalue;
313 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000314
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000315 assert(gbo->currkey != NULL);
316 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
317 if (rcmp <= 0)
318 /* got any error or current group is end */
319 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000320
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000321 r = gbo->currvalue;
322 gbo->currvalue = NULL;
323 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000324
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000325 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000326}
327
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000328static PyObject *
329_grouper_reduce(_grouperobject *lz)
330{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700331 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000332}
333
334static PyMethodDef _grouper_methods[] = {
335 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
336 reduce_doc},
337 {NULL, NULL} /* sentinel */
338};
339
340
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000341static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000342 PyVarObject_HEAD_INIT(NULL, 0)
343 "itertools._grouper", /* tp_name */
344 sizeof(_grouperobject), /* tp_basicsize */
345 0, /* tp_itemsize */
346 /* methods */
347 (destructor)_grouper_dealloc, /* tp_dealloc */
348 0, /* tp_print */
349 0, /* tp_getattr */
350 0, /* tp_setattr */
351 0, /* tp_reserved */
352 0, /* tp_repr */
353 0, /* tp_as_number */
354 0, /* tp_as_sequence */
355 0, /* tp_as_mapping */
356 0, /* tp_hash */
357 0, /* tp_call */
358 0, /* tp_str */
359 PyObject_GenericGetAttr, /* tp_getattro */
360 0, /* tp_setattro */
361 0, /* tp_as_buffer */
362 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
363 0, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700364 (traverseproc)_grouper_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000365 0, /* tp_clear */
366 0, /* tp_richcompare */
367 0, /* tp_weaklistoffset */
368 PyObject_SelfIter, /* tp_iter */
369 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000370 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000371 0, /* tp_members */
372 0, /* tp_getset */
373 0, /* tp_base */
374 0, /* tp_dict */
375 0, /* tp_descr_get */
376 0, /* tp_descr_set */
377 0, /* tp_dictoffset */
378 0, /* tp_init */
379 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000380 _grouper_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000381 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000382};
383
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000384
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700385/* tee object and with supporting function and objects ***********************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000386
Raymond Hettingerad983e72003-11-12 14:32:26 +0000387/* The teedataobject pre-allocates space for LINKCELLS number of objects.
388 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000389 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000390 the other structure members including PyHEAD overhead. The larger the
391 value, the less memory overhead per object and the less time spent
392 allocating/deallocating new links. The smaller the number, the less
393 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000394*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000395#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000396
397typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000398 PyObject_HEAD
399 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700400 int numread; /* 0 <= numread <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000401 PyObject *nextlink;
402 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000403} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000404
405typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000406 PyObject_HEAD
407 teedataobject *dataobj;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700408 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000409 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000410} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000411
Raymond Hettingerad983e72003-11-12 14:32:26 +0000412static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000413
414static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000415teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000416{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000417 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000418
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000419 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
420 if (tdo == NULL)
421 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000422
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000423 tdo->numread = 0;
424 tdo->nextlink = NULL;
425 Py_INCREF(it);
426 tdo->it = it;
427 PyObject_GC_Track(tdo);
428 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000429}
430
431static PyObject *
432teedataobject_jumplink(teedataobject *tdo)
433{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000434 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000435 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000436 Py_XINCREF(tdo->nextlink);
437 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000438}
439
440static PyObject *
441teedataobject_getitem(teedataobject *tdo, int i)
442{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000443 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000444
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000445 assert(i < LINKCELLS);
446 if (i < tdo->numread)
447 value = tdo->values[i];
448 else {
449 /* this is the lead iterator, so fetch more data */
450 assert(i == tdo->numread);
451 value = PyIter_Next(tdo->it);
452 if (value == NULL)
453 return NULL;
454 tdo->numread++;
455 tdo->values[i] = value;
456 }
457 Py_INCREF(value);
458 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000459}
460
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000461static int
462teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
463{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000464 int i;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700465
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000466 Py_VISIT(tdo->it);
467 for (i = 0; i < tdo->numread; i++)
468 Py_VISIT(tdo->values[i]);
469 Py_VISIT(tdo->nextlink);
470 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000471}
472
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200473static void
474teedataobject_safe_decref(PyObject *obj)
475{
476 while (obj && Py_TYPE(obj) == &teedataobject_type &&
477 Py_REFCNT(obj) == 1) {
478 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
479 ((teedataobject *)obj)->nextlink = NULL;
480 Py_DECREF(obj);
481 obj = nextlink;
482 }
483 Py_XDECREF(obj);
484}
485
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000486static int
487teedataobject_clear(teedataobject *tdo)
488{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000489 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200490 PyObject *tmp;
491
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000492 Py_CLEAR(tdo->it);
493 for (i=0 ; i<tdo->numread ; i++)
494 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200495 tmp = tdo->nextlink;
496 tdo->nextlink = NULL;
497 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000498 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000499}
500
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000501static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000502teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000503{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000504 PyObject_GC_UnTrack(tdo);
505 teedataobject_clear(tdo);
506 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000507}
508
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000509static PyObject *
510teedataobject_reduce(teedataobject *tdo)
511{
512 int i;
513 /* create a temporary list of already iterated values */
514 PyObject *values = PyList_New(tdo->numread);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700515
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000516 if (!values)
517 return NULL;
518 for (i=0 ; i<tdo->numread ; i++) {
519 Py_INCREF(tdo->values[i]);
520 PyList_SET_ITEM(values, i, tdo->values[i]);
521 }
522 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
523 values,
524 tdo->nextlink ? tdo->nextlink : Py_None);
525}
526
527static PyTypeObject teedataobject_type;
528
529static PyObject *
530teedataobject_new(PyTypeObject *type, PyObject *args, PyObject *kw)
531{
532 teedataobject *tdo;
533 PyObject *it, *values, *next;
534 Py_ssize_t i, len;
535
536 assert(type == &teedataobject_type);
537 if (!PyArg_ParseTuple(args, "OO!O", &it, &PyList_Type, &values, &next))
538 return NULL;
539
540 tdo = (teedataobject *)teedataobject_newinternal(it);
541 if (!tdo)
542 return NULL;
543
544 len = PyList_GET_SIZE(values);
545 if (len > LINKCELLS)
546 goto err;
547 for (i=0; i<len; i++) {
548 tdo->values[i] = PyList_GET_ITEM(values, i);
549 Py_INCREF(tdo->values[i]);
550 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200551 /* len <= LINKCELLS < INT_MAX */
552 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000553
554 if (len == LINKCELLS) {
555 if (next != Py_None) {
556 if (Py_TYPE(next) != &teedataobject_type)
557 goto err;
558 assert(tdo->nextlink == NULL);
559 Py_INCREF(next);
560 tdo->nextlink = next;
561 }
562 } else {
563 if (next != Py_None)
564 goto err; /* shouldn't have a next if we are not full */
565 }
566 return (PyObject*)tdo;
567
568err:
569 Py_XDECREF(tdo);
570 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
571 return NULL;
572}
573
574static PyMethodDef teedataobject_methods[] = {
575 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
576 reduce_doc},
577 {NULL, NULL} /* sentinel */
578};
579
Raymond Hettingerad983e72003-11-12 14:32:26 +0000580PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000581
Raymond Hettingerad983e72003-11-12 14:32:26 +0000582static PyTypeObject teedataobject_type = {
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700583 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000584 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000585 sizeof(teedataobject), /* tp_basicsize */
586 0, /* tp_itemsize */
587 /* methods */
588 (destructor)teedataobject_dealloc, /* tp_dealloc */
589 0, /* tp_print */
590 0, /* tp_getattr */
591 0, /* tp_setattr */
592 0, /* tp_reserved */
593 0, /* tp_repr */
594 0, /* tp_as_number */
595 0, /* tp_as_sequence */
596 0, /* tp_as_mapping */
597 0, /* tp_hash */
598 0, /* tp_call */
599 0, /* tp_str */
600 PyObject_GenericGetAttr, /* tp_getattro */
601 0, /* tp_setattro */
602 0, /* tp_as_buffer */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700603 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000604 teedataobject_doc, /* tp_doc */
605 (traverseproc)teedataobject_traverse, /* tp_traverse */
606 (inquiry)teedataobject_clear, /* tp_clear */
607 0, /* tp_richcompare */
608 0, /* tp_weaklistoffset */
609 0, /* tp_iter */
610 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000611 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000612 0, /* tp_members */
613 0, /* tp_getset */
614 0, /* tp_base */
615 0, /* tp_dict */
616 0, /* tp_descr_get */
617 0, /* tp_descr_set */
618 0, /* tp_dictoffset */
619 0, /* tp_init */
620 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000621 teedataobject_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000622 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000623};
624
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000625
626static PyTypeObject tee_type;
627
628static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000629tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000630{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000631 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000632
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000633 if (to->index >= LINKCELLS) {
634 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200635 if (link == NULL)
636 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000637 Py_DECREF(to->dataobj);
638 to->dataobj = (teedataobject *)link;
639 to->index = 0;
640 }
641 value = teedataobject_getitem(to->dataobj, to->index);
642 if (value == NULL)
643 return NULL;
644 to->index++;
645 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000646}
647
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000648static int
649tee_traverse(teeobject *to, visitproc visit, void *arg)
650{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000651 Py_VISIT((PyObject *)to->dataobj);
652 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000653}
654
Raymond Hettingerad983e72003-11-12 14:32:26 +0000655static PyObject *
656tee_copy(teeobject *to)
657{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000658 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000659
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000660 newto = PyObject_GC_New(teeobject, &tee_type);
661 if (newto == NULL)
662 return NULL;
663 Py_INCREF(to->dataobj);
664 newto->dataobj = to->dataobj;
665 newto->index = to->index;
666 newto->weakreflist = NULL;
667 PyObject_GC_Track(newto);
668 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000669}
670
671PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
672
673static PyObject *
674tee_fromiterable(PyObject *iterable)
675{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000676 teeobject *to;
677 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000678
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000679 it = PyObject_GetIter(iterable);
680 if (it == NULL)
681 return NULL;
682 if (PyObject_TypeCheck(it, &tee_type)) {
683 to = (teeobject *)tee_copy((teeobject *)it);
684 goto done;
685 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000686
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000687 to = PyObject_GC_New(teeobject, &tee_type);
688 if (to == NULL)
689 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000690 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000691 if (!to->dataobj) {
692 PyObject_GC_Del(to);
693 to = NULL;
694 goto done;
695 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000696
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000697 to->index = 0;
698 to->weakreflist = NULL;
699 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000700done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000701 Py_XDECREF(it);
702 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000703}
704
705static PyObject *
706tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
707{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000708 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000709
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000710 if (!PyArg_UnpackTuple(args, "_tee", 1, 1, &iterable))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000711 return NULL;
712 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000713}
714
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000715static int
716tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000717{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000718 if (to->weakreflist != NULL)
719 PyObject_ClearWeakRefs((PyObject *) to);
720 Py_CLEAR(to->dataobj);
721 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000722}
723
724static void
725tee_dealloc(teeobject *to)
726{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000727 PyObject_GC_UnTrack(to);
728 tee_clear(to);
729 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000730}
731
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000732static PyObject *
733tee_reduce(teeobject *to)
734{
735 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
736}
737
738static PyObject *
739tee_setstate(teeobject *to, PyObject *state)
740{
741 teedataobject *tdo;
742 int index;
743 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index))
744 return NULL;
745 if (index < 0 || index > LINKCELLS) {
746 PyErr_SetString(PyExc_ValueError, "Index out of range");
747 return NULL;
748 }
749 Py_CLEAR(to->dataobj);
750 to->dataobj = tdo;
751 Py_INCREF(to->dataobj);
752 to->index = index;
753 Py_RETURN_NONE;
754}
755
Raymond Hettingerad983e72003-11-12 14:32:26 +0000756PyDoc_STRVAR(teeobject_doc,
757"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000758
Raymond Hettingerad983e72003-11-12 14:32:26 +0000759static PyMethodDef tee_methods[] = {
Raymond Hettingera6a2d442015-08-16 14:38:07 -0700760 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
761 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
762 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000763 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000764};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000765
766static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000767 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000768 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000769 sizeof(teeobject), /* tp_basicsize */
770 0, /* tp_itemsize */
771 /* methods */
772 (destructor)tee_dealloc, /* tp_dealloc */
773 0, /* tp_print */
774 0, /* tp_getattr */
775 0, /* tp_setattr */
776 0, /* tp_reserved */
777 0, /* tp_repr */
778 0, /* tp_as_number */
779 0, /* tp_as_sequence */
780 0, /* tp_as_mapping */
781 0, /* tp_hash */
782 0, /* tp_call */
783 0, /* tp_str */
784 0, /* tp_getattro */
785 0, /* tp_setattro */
786 0, /* tp_as_buffer */
787 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
788 teeobject_doc, /* tp_doc */
789 (traverseproc)tee_traverse, /* tp_traverse */
790 (inquiry)tee_clear, /* tp_clear */
791 0, /* tp_richcompare */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700792 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000793 PyObject_SelfIter, /* tp_iter */
794 (iternextfunc)tee_next, /* tp_iternext */
795 tee_methods, /* tp_methods */
796 0, /* tp_members */
797 0, /* tp_getset */
798 0, /* tp_base */
799 0, /* tp_dict */
800 0, /* tp_descr_get */
801 0, /* tp_descr_set */
802 0, /* tp_dictoffset */
803 0, /* tp_init */
804 0, /* tp_alloc */
805 tee_new, /* tp_new */
806 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000807};
808
Raymond Hettingerad983e72003-11-12 14:32:26 +0000809static PyObject *
810tee(PyObject *self, PyObject *args)
811{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000812 Py_ssize_t i, n=2;
813 PyObject *it, *iterable, *copyable, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200814 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000815
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000816 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
817 return NULL;
818 if (n < 0) {
819 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
820 return NULL;
821 }
822 result = PyTuple_New(n);
823 if (result == NULL)
824 return NULL;
825 if (n == 0)
826 return result;
827 it = PyObject_GetIter(iterable);
828 if (it == NULL) {
829 Py_DECREF(result);
830 return NULL;
831 }
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200832 if (!_PyObject_HasAttrId(it, &PyId___copy__)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000833 copyable = tee_fromiterable(it);
834 Py_DECREF(it);
835 if (copyable == NULL) {
836 Py_DECREF(result);
837 return NULL;
838 }
839 } else
840 copyable = it;
841 PyTuple_SET_ITEM(result, 0, copyable);
842 for (i=1 ; i<n ; i++) {
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200843
844 copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000845 if (copyable == NULL) {
846 Py_DECREF(result);
847 return NULL;
848 }
849 PyTuple_SET_ITEM(result, i, copyable);
850 }
851 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000852}
853
854PyDoc_STRVAR(tee_doc,
855"tee(iterable, n=2) --> tuple of n independent iterators.");
856
857
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700858/* cycle object **************************************************************/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000859
860typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000861 PyObject_HEAD
862 PyObject *it;
863 PyObject *saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700864 Py_ssize_t index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000865 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000866} cycleobject;
867
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000868static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000869
870static PyObject *
871cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
872{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000873 PyObject *it;
874 PyObject *iterable;
875 PyObject *saved;
876 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000877
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000878 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
879 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000880
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000881 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
882 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000883
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000884 /* Get iterator. */
885 it = PyObject_GetIter(iterable);
886 if (it == NULL)
887 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000888
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000889 saved = PyList_New(0);
890 if (saved == NULL) {
891 Py_DECREF(it);
892 return NULL;
893 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000894
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000895 /* create cycleobject structure */
896 lz = (cycleobject *)type->tp_alloc(type, 0);
897 if (lz == NULL) {
898 Py_DECREF(it);
899 Py_DECREF(saved);
900 return NULL;
901 }
902 lz->it = it;
903 lz->saved = saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700904 lz->index = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000905 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);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000914 Py_XDECREF(lz->it);
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700915 Py_XDECREF(lz->saved);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000916 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{
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700922 if (lz->it)
923 Py_VISIT(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000924 Py_VISIT(lz->saved);
925 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000926}
927
928static PyObject *
929cycle_next(cycleobject *lz)
930{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000931 PyObject *item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000932
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700933 if (lz->it != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000934 item = PyIter_Next(lz->it);
935 if (item != NULL) {
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700936 if (lz->firstpass)
937 return item;
938 if (PyList_Append(lz->saved, item)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000939 Py_DECREF(item);
940 return NULL;
941 }
942 return item;
943 }
Raymond Hettinger98958fe2015-08-15 15:09:30 -0700944 /* Note: StopIteration is already cleared by PyIter_Next() */
945 if (PyErr_Occurred())
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000946 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700947 Py_CLEAR(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000948 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700949 if (Py_SIZE(lz->saved) == 0)
950 return NULL;
951 item = PyList_GET_ITEM(lz->saved, lz->index);
952 lz->index++;
953 if (lz->index >= Py_SIZE(lz->saved))
954 lz->index = 0;
955 Py_INCREF(item);
956 return item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000957}
958
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000959static PyObject *
960cycle_reduce(cycleobject *lz)
961{
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700962 /* Create a new cycle with the iterator tuple, then set the saved state */
963 if (lz->it == NULL) {
964 PyObject *it = PyObject_GetIter(lz->saved);
965 if (it == NULL)
966 return NULL;
967 if (lz->index != 0) {
968 _Py_IDENTIFIER(__setstate__);
969 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
970 "n", lz->index);
971 if (res == NULL) {
972 Py_DECREF(it);
973 return NULL;
974 }
975 Py_DECREF(res);
976 }
977 return Py_BuildValue("O(N)(Oi)", Py_TYPE(lz), it, lz->saved, 1);
978 }
979 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), lz->it, lz->saved,
980 lz->firstpass);
Raymond Hettingerc39786d2015-08-15 15:16:12 -0700981}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000982
983static PyObject *
984cycle_setstate(cycleobject *lz, PyObject *state)
985{
986 PyObject *saved=NULL;
987 int firstpass;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700988
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700989 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000990 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700991 Py_INCREF(saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000992 Py_CLEAR(lz->saved);
993 lz->saved = saved;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000994 lz->firstpass = firstpass != 0;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700995 lz->index = 0;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000996 Py_RETURN_NONE;
997}
998
999static PyMethodDef cycle_methods[] = {
1000 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1001 reduce_doc},
1002 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1003 setstate_doc},
1004 {NULL, NULL} /* sentinel */
1005};
1006
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001007PyDoc_STRVAR(cycle_doc,
1008"cycle(iterable) --> cycle object\n\
1009\n\
1010Return elements from the iterable until it is exhausted.\n\
1011Then repeat the sequence indefinitely.");
1012
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001013static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001014 PyVarObject_HEAD_INIT(NULL, 0)
1015 "itertools.cycle", /* tp_name */
1016 sizeof(cycleobject), /* tp_basicsize */
1017 0, /* tp_itemsize */
1018 /* methods */
1019 (destructor)cycle_dealloc, /* tp_dealloc */
1020 0, /* tp_print */
1021 0, /* tp_getattr */
1022 0, /* tp_setattr */
1023 0, /* tp_reserved */
1024 0, /* tp_repr */
1025 0, /* tp_as_number */
1026 0, /* tp_as_sequence */
1027 0, /* tp_as_mapping */
1028 0, /* tp_hash */
1029 0, /* tp_call */
1030 0, /* tp_str */
1031 PyObject_GenericGetAttr, /* tp_getattro */
1032 0, /* tp_setattro */
1033 0, /* tp_as_buffer */
1034 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1035 Py_TPFLAGS_BASETYPE, /* tp_flags */
1036 cycle_doc, /* tp_doc */
1037 (traverseproc)cycle_traverse, /* tp_traverse */
1038 0, /* tp_clear */
1039 0, /* tp_richcompare */
1040 0, /* tp_weaklistoffset */
1041 PyObject_SelfIter, /* tp_iter */
1042 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001043 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001044 0, /* tp_members */
1045 0, /* tp_getset */
1046 0, /* tp_base */
1047 0, /* tp_dict */
1048 0, /* tp_descr_get */
1049 0, /* tp_descr_set */
1050 0, /* tp_dictoffset */
1051 0, /* tp_init */
1052 0, /* tp_alloc */
1053 cycle_new, /* tp_new */
1054 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001055};
1056
1057
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001058/* dropwhile object **********************************************************/
1059
1060typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001061 PyObject_HEAD
1062 PyObject *func;
1063 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001064 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001065} dropwhileobject;
1066
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001067static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001068
1069static PyObject *
1070dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1071{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001072 PyObject *func, *seq;
1073 PyObject *it;
1074 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001075
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001076 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
1077 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001078
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001079 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
1080 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001081
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001082 /* Get iterator. */
1083 it = PyObject_GetIter(seq);
1084 if (it == NULL)
1085 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001086
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001087 /* create dropwhileobject structure */
1088 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1089 if (lz == NULL) {
1090 Py_DECREF(it);
1091 return NULL;
1092 }
1093 Py_INCREF(func);
1094 lz->func = func;
1095 lz->it = it;
1096 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001097
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001098 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001099}
1100
1101static void
1102dropwhile_dealloc(dropwhileobject *lz)
1103{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001104 PyObject_GC_UnTrack(lz);
1105 Py_XDECREF(lz->func);
1106 Py_XDECREF(lz->it);
1107 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001108}
1109
1110static int
1111dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1112{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001113 Py_VISIT(lz->it);
1114 Py_VISIT(lz->func);
1115 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116}
1117
1118static PyObject *
1119dropwhile_next(dropwhileobject *lz)
1120{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001121 PyObject *item, *good;
1122 PyObject *it = lz->it;
1123 long ok;
1124 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001125
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001126 iternext = *Py_TYPE(it)->tp_iternext;
1127 for (;;) {
1128 item = iternext(it);
1129 if (item == NULL)
1130 return NULL;
1131 if (lz->start == 1)
1132 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001133
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001134 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1135 if (good == NULL) {
1136 Py_DECREF(item);
1137 return NULL;
1138 }
1139 ok = PyObject_IsTrue(good);
1140 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001141 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001142 lz->start = 1;
1143 return item;
1144 }
1145 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001146 if (ok < 0)
1147 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001148 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001149}
1150
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001151static PyObject *
1152dropwhile_reduce(dropwhileobject *lz)
1153{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001154 return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->start);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001155}
1156
1157static PyObject *
1158dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1159{
1160 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001161 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001162 return NULL;
1163 lz->start = start;
1164 Py_RETURN_NONE;
1165}
1166
1167static PyMethodDef dropwhile_methods[] = {
1168 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1169 reduce_doc},
1170 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1171 setstate_doc},
1172 {NULL, NULL} /* sentinel */
1173};
1174
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001175PyDoc_STRVAR(dropwhile_doc,
1176"dropwhile(predicate, iterable) --> dropwhile object\n\
1177\n\
1178Drop items from the iterable while predicate(item) is true.\n\
1179Afterwards, return every element until the iterable is exhausted.");
1180
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001181static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001182 PyVarObject_HEAD_INIT(NULL, 0)
1183 "itertools.dropwhile", /* tp_name */
1184 sizeof(dropwhileobject), /* tp_basicsize */
1185 0, /* tp_itemsize */
1186 /* methods */
1187 (destructor)dropwhile_dealloc, /* tp_dealloc */
1188 0, /* tp_print */
1189 0, /* tp_getattr */
1190 0, /* tp_setattr */
1191 0, /* tp_reserved */
1192 0, /* tp_repr */
1193 0, /* tp_as_number */
1194 0, /* tp_as_sequence */
1195 0, /* tp_as_mapping */
1196 0, /* tp_hash */
1197 0, /* tp_call */
1198 0, /* tp_str */
1199 PyObject_GenericGetAttr, /* tp_getattro */
1200 0, /* tp_setattro */
1201 0, /* tp_as_buffer */
1202 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1203 Py_TPFLAGS_BASETYPE, /* tp_flags */
1204 dropwhile_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001205 (traverseproc)dropwhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001206 0, /* tp_clear */
1207 0, /* tp_richcompare */
1208 0, /* tp_weaklistoffset */
1209 PyObject_SelfIter, /* tp_iter */
1210 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001211 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001212 0, /* tp_members */
1213 0, /* tp_getset */
1214 0, /* tp_base */
1215 0, /* tp_dict */
1216 0, /* tp_descr_get */
1217 0, /* tp_descr_set */
1218 0, /* tp_dictoffset */
1219 0, /* tp_init */
1220 0, /* tp_alloc */
1221 dropwhile_new, /* tp_new */
1222 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001223};
1224
1225
1226/* takewhile object **********************************************************/
1227
1228typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001229 PyObject_HEAD
1230 PyObject *func;
1231 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001232 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001233} takewhileobject;
1234
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001235static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001236
1237static PyObject *
1238takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1239{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001240 PyObject *func, *seq;
1241 PyObject *it;
1242 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001243
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001244 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1245 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001246
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001247 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1248 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001249
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001250 /* Get iterator. */
1251 it = PyObject_GetIter(seq);
1252 if (it == NULL)
1253 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001254
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001255 /* create takewhileobject structure */
1256 lz = (takewhileobject *)type->tp_alloc(type, 0);
1257 if (lz == NULL) {
1258 Py_DECREF(it);
1259 return NULL;
1260 }
1261 Py_INCREF(func);
1262 lz->func = func;
1263 lz->it = it;
1264 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001265
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001266 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001267}
1268
1269static void
1270takewhile_dealloc(takewhileobject *lz)
1271{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001272 PyObject_GC_UnTrack(lz);
1273 Py_XDECREF(lz->func);
1274 Py_XDECREF(lz->it);
1275 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001276}
1277
1278static int
1279takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1280{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001281 Py_VISIT(lz->it);
1282 Py_VISIT(lz->func);
1283 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001284}
1285
1286static PyObject *
1287takewhile_next(takewhileobject *lz)
1288{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001289 PyObject *item, *good;
1290 PyObject *it = lz->it;
1291 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001292
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001293 if (lz->stop == 1)
1294 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001295
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001296 item = (*Py_TYPE(it)->tp_iternext)(it);
1297 if (item == NULL)
1298 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001299
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001300 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1301 if (good == NULL) {
1302 Py_DECREF(item);
1303 return NULL;
1304 }
1305 ok = PyObject_IsTrue(good);
1306 Py_DECREF(good);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001307 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001308 return item;
1309 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001310 if (ok == 0)
1311 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001312 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001313}
1314
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001315static PyObject *
1316takewhile_reduce(takewhileobject *lz)
1317{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001318 return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->stop);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001319}
1320
1321static PyObject *
1322takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1323{
1324 int stop = PyObject_IsTrue(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001325
Antoine Pitrou721738f2012-08-15 23:20:39 +02001326 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001327 return NULL;
1328 lz->stop = stop;
1329 Py_RETURN_NONE;
1330}
1331
1332static PyMethodDef takewhile_reduce_methods[] = {
1333 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1334 reduce_doc},
1335 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1336 setstate_doc},
1337 {NULL, NULL} /* sentinel */
1338};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001339PyDoc_STRVAR(takewhile_doc,
1340"takewhile(predicate, iterable) --> takewhile object\n\
1341\n\
1342Return successive entries from an iterable as long as the \n\
1343predicate evaluates to true for each entry.");
1344
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001345static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001346 PyVarObject_HEAD_INIT(NULL, 0)
1347 "itertools.takewhile", /* tp_name */
1348 sizeof(takewhileobject), /* tp_basicsize */
1349 0, /* tp_itemsize */
1350 /* methods */
1351 (destructor)takewhile_dealloc, /* tp_dealloc */
1352 0, /* tp_print */
1353 0, /* tp_getattr */
1354 0, /* tp_setattr */
1355 0, /* tp_reserved */
1356 0, /* tp_repr */
1357 0, /* tp_as_number */
1358 0, /* tp_as_sequence */
1359 0, /* tp_as_mapping */
1360 0, /* tp_hash */
1361 0, /* tp_call */
1362 0, /* tp_str */
1363 PyObject_GenericGetAttr, /* tp_getattro */
1364 0, /* tp_setattro */
1365 0, /* tp_as_buffer */
1366 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1367 Py_TPFLAGS_BASETYPE, /* tp_flags */
1368 takewhile_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001369 (traverseproc)takewhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001370 0, /* tp_clear */
1371 0, /* tp_richcompare */
1372 0, /* tp_weaklistoffset */
1373 PyObject_SelfIter, /* tp_iter */
1374 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001375 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001376 0, /* tp_members */
1377 0, /* tp_getset */
1378 0, /* tp_base */
1379 0, /* tp_dict */
1380 0, /* tp_descr_get */
1381 0, /* tp_descr_set */
1382 0, /* tp_dictoffset */
1383 0, /* tp_init */
1384 0, /* tp_alloc */
1385 takewhile_new, /* tp_new */
1386 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001387};
1388
1389
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001390/* islice object *************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001391
1392typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001393 PyObject_HEAD
1394 PyObject *it;
1395 Py_ssize_t next;
1396 Py_ssize_t stop;
1397 Py_ssize_t step;
1398 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001399} isliceobject;
1400
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001401static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001402
1403static PyObject *
1404islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1405{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001406 PyObject *seq;
1407 Py_ssize_t start=0, stop=-1, step=1;
1408 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1409 Py_ssize_t numargs;
1410 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001411
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001412 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1413 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001414
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001415 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1416 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001417
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001418 numargs = PyTuple_Size(args);
1419 if (numargs == 2) {
1420 if (a1 != Py_None) {
1421 stop = PyLong_AsSsize_t(a1);
1422 if (stop == -1) {
1423 if (PyErr_Occurred())
1424 PyErr_Clear();
1425 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001426 "Stop argument for islice() must be None or "
1427 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001428 return NULL;
1429 }
1430 }
1431 } else {
1432 if (a1 != Py_None)
1433 start = PyLong_AsSsize_t(a1);
1434 if (start == -1 && PyErr_Occurred())
1435 PyErr_Clear();
1436 if (a2 != Py_None) {
1437 stop = PyLong_AsSsize_t(a2);
1438 if (stop == -1) {
1439 if (PyErr_Occurred())
1440 PyErr_Clear();
1441 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001442 "Stop argument for islice() must be None or "
1443 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001444 return NULL;
1445 }
1446 }
1447 }
1448 if (start<0 || stop<-1) {
1449 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001450 "Indices for islice() must be None or "
1451 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001452 return NULL;
1453 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001454
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001455 if (a3 != NULL) {
1456 if (a3 != Py_None)
1457 step = PyLong_AsSsize_t(a3);
1458 if (step == -1 && PyErr_Occurred())
1459 PyErr_Clear();
1460 }
1461 if (step<1) {
1462 PyErr_SetString(PyExc_ValueError,
1463 "Step for islice() must be a positive integer or None.");
1464 return NULL;
1465 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001466
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001467 /* Get iterator. */
1468 it = PyObject_GetIter(seq);
1469 if (it == NULL)
1470 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001471
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001472 /* create isliceobject structure */
1473 lz = (isliceobject *)type->tp_alloc(type, 0);
1474 if (lz == NULL) {
1475 Py_DECREF(it);
1476 return NULL;
1477 }
1478 lz->it = it;
1479 lz->next = start;
1480 lz->stop = stop;
1481 lz->step = step;
1482 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001483
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001484 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001485}
1486
1487static void
1488islice_dealloc(isliceobject *lz)
1489{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001490 PyObject_GC_UnTrack(lz);
1491 Py_XDECREF(lz->it);
1492 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001493}
1494
1495static int
1496islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1497{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001498 Py_VISIT(lz->it);
1499 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001500}
1501
1502static PyObject *
1503islice_next(isliceobject *lz)
1504{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001505 PyObject *item;
1506 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001507 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001508 Py_ssize_t oldnext;
1509 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001510
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001511 if (it == NULL)
1512 return NULL;
1513
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001514 iternext = *Py_TYPE(it)->tp_iternext;
1515 while (lz->cnt < lz->next) {
1516 item = iternext(it);
1517 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001518 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001519 Py_DECREF(item);
1520 lz->cnt++;
1521 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001522 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001523 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001524 item = iternext(it);
1525 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001526 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001527 lz->cnt++;
1528 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001529 /* The (size_t) cast below avoids the danger of undefined
1530 behaviour from signed integer overflow. */
1531 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001532 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1533 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001534 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001535
1536empty:
1537 Py_CLEAR(lz->it);
1538 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001539}
1540
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001541static PyObject *
1542islice_reduce(isliceobject *lz)
1543{
1544 /* When unpickled, generate a new object with the same bounds,
1545 * then 'setstate' with the next and count
1546 */
1547 PyObject *stop;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001548
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001549 if (lz->it == NULL) {
1550 PyObject *empty_list;
1551 PyObject *empty_it;
1552 empty_list = PyList_New(0);
1553 if (empty_list == NULL)
1554 return NULL;
1555 empty_it = PyObject_GetIter(empty_list);
1556 Py_DECREF(empty_list);
1557 if (empty_it == NULL)
1558 return NULL;
1559 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1560 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001561 if (lz->stop == -1) {
1562 stop = Py_None;
1563 Py_INCREF(stop);
1564 } else {
1565 stop = PyLong_FromSsize_t(lz->stop);
1566 if (stop == NULL)
1567 return NULL;
1568 }
1569 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1570 lz->it, lz->next, stop, lz->step,
1571 lz->cnt);
1572}
1573
1574static PyObject *
1575islice_setstate(isliceobject *lz, PyObject *state)
1576{
1577 Py_ssize_t cnt = PyLong_AsSsize_t(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001578
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001579 if (cnt == -1 && PyErr_Occurred())
1580 return NULL;
1581 lz->cnt = cnt;
1582 Py_RETURN_NONE;
1583}
1584
1585static PyMethodDef islice_methods[] = {
1586 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1587 reduce_doc},
1588 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1589 setstate_doc},
1590 {NULL, NULL} /* sentinel */
1591};
1592
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001593PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001594"islice(iterable, stop) --> islice object\n\
1595islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001596\n\
1597Return an iterator whose next() method returns selected values from an\n\
1598iterable. If start is specified, will skip all preceding elements;\n\
1599otherwise, start defaults to zero. Step defaults to one. If\n\
1600specified as another value, step determines how many values are \n\
1601skipped between successive calls. Works like a slice() on a list\n\
1602but returns an iterator.");
1603
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001604static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001605 PyVarObject_HEAD_INIT(NULL, 0)
1606 "itertools.islice", /* tp_name */
1607 sizeof(isliceobject), /* tp_basicsize */
1608 0, /* tp_itemsize */
1609 /* methods */
1610 (destructor)islice_dealloc, /* tp_dealloc */
1611 0, /* tp_print */
1612 0, /* tp_getattr */
1613 0, /* tp_setattr */
1614 0, /* tp_reserved */
1615 0, /* tp_repr */
1616 0, /* tp_as_number */
1617 0, /* tp_as_sequence */
1618 0, /* tp_as_mapping */
1619 0, /* tp_hash */
1620 0, /* tp_call */
1621 0, /* tp_str */
1622 PyObject_GenericGetAttr, /* tp_getattro */
1623 0, /* tp_setattro */
1624 0, /* tp_as_buffer */
1625 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1626 Py_TPFLAGS_BASETYPE, /* tp_flags */
1627 islice_doc, /* tp_doc */
1628 (traverseproc)islice_traverse, /* tp_traverse */
1629 0, /* tp_clear */
1630 0, /* tp_richcompare */
1631 0, /* tp_weaklistoffset */
1632 PyObject_SelfIter, /* tp_iter */
1633 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001634 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001635 0, /* tp_members */
1636 0, /* tp_getset */
1637 0, /* tp_base */
1638 0, /* tp_dict */
1639 0, /* tp_descr_get */
1640 0, /* tp_descr_set */
1641 0, /* tp_dictoffset */
1642 0, /* tp_init */
1643 0, /* tp_alloc */
1644 islice_new, /* tp_new */
1645 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001646};
1647
1648
1649/* starmap object ************************************************************/
1650
1651typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001652 PyObject_HEAD
1653 PyObject *func;
1654 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001655} starmapobject;
1656
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001657static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001658
1659static PyObject *
1660starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1661{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001662 PyObject *func, *seq;
1663 PyObject *it;
1664 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001665
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001666 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1667 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001668
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001669 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1670 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001671
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001672 /* Get iterator. */
1673 it = PyObject_GetIter(seq);
1674 if (it == NULL)
1675 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001676
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001677 /* create starmapobject structure */
1678 lz = (starmapobject *)type->tp_alloc(type, 0);
1679 if (lz == NULL) {
1680 Py_DECREF(it);
1681 return NULL;
1682 }
1683 Py_INCREF(func);
1684 lz->func = func;
1685 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001686
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001687 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001688}
1689
1690static void
1691starmap_dealloc(starmapobject *lz)
1692{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001693 PyObject_GC_UnTrack(lz);
1694 Py_XDECREF(lz->func);
1695 Py_XDECREF(lz->it);
1696 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001697}
1698
1699static int
1700starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1701{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001702 Py_VISIT(lz->it);
1703 Py_VISIT(lz->func);
1704 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001705}
1706
1707static PyObject *
1708starmap_next(starmapobject *lz)
1709{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001710 PyObject *args;
1711 PyObject *result;
1712 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001713
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001714 args = (*Py_TYPE(it)->tp_iternext)(it);
1715 if (args == NULL)
1716 return NULL;
1717 if (!PyTuple_CheckExact(args)) {
1718 PyObject *newargs = PySequence_Tuple(args);
1719 Py_DECREF(args);
1720 if (newargs == NULL)
1721 return NULL;
1722 args = newargs;
1723 }
1724 result = PyObject_Call(lz->func, args, NULL);
1725 Py_DECREF(args);
1726 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001727}
1728
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001729static PyObject *
1730starmap_reduce(starmapobject *lz)
1731{
1732 /* Just pickle the iterator */
1733 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1734}
1735
1736static PyMethodDef starmap_methods[] = {
1737 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1738 reduce_doc},
1739 {NULL, NULL} /* sentinel */
1740};
1741
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001742PyDoc_STRVAR(starmap_doc,
1743"starmap(function, sequence) --> starmap object\n\
1744\n\
1745Return an iterator whose values are returned from the function evaluated\n\
1746with a argument tuple taken from the given sequence.");
1747
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001748static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001749 PyVarObject_HEAD_INIT(NULL, 0)
1750 "itertools.starmap", /* tp_name */
1751 sizeof(starmapobject), /* tp_basicsize */
1752 0, /* tp_itemsize */
1753 /* methods */
1754 (destructor)starmap_dealloc, /* tp_dealloc */
1755 0, /* tp_print */
1756 0, /* tp_getattr */
1757 0, /* tp_setattr */
1758 0, /* tp_reserved */
1759 0, /* tp_repr */
1760 0, /* tp_as_number */
1761 0, /* tp_as_sequence */
1762 0, /* tp_as_mapping */
1763 0, /* tp_hash */
1764 0, /* tp_call */
1765 0, /* tp_str */
1766 PyObject_GenericGetAttr, /* tp_getattro */
1767 0, /* tp_setattro */
1768 0, /* tp_as_buffer */
1769 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1770 Py_TPFLAGS_BASETYPE, /* tp_flags */
1771 starmap_doc, /* tp_doc */
1772 (traverseproc)starmap_traverse, /* tp_traverse */
1773 0, /* tp_clear */
1774 0, /* tp_richcompare */
1775 0, /* tp_weaklistoffset */
1776 PyObject_SelfIter, /* tp_iter */
1777 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001778 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001779 0, /* tp_members */
1780 0, /* tp_getset */
1781 0, /* tp_base */
1782 0, /* tp_dict */
1783 0, /* tp_descr_get */
1784 0, /* tp_descr_set */
1785 0, /* tp_dictoffset */
1786 0, /* tp_init */
1787 0, /* tp_alloc */
1788 starmap_new, /* tp_new */
1789 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001790};
1791
1792
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001793/* chain object **************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001794
1795typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001796 PyObject_HEAD
1797 PyObject *source; /* Iterator over input iterables */
1798 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001799} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001800
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001801static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001802
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001803static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001804chain_new_internal(PyTypeObject *type, PyObject *source)
1805{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001806 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001807
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001808 lz = (chainobject *)type->tp_alloc(type, 0);
1809 if (lz == NULL) {
1810 Py_DECREF(source);
1811 return NULL;
1812 }
1813
1814 lz->source = source;
1815 lz->active = NULL;
1816 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001817}
1818
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001819static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001820chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001821{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001822 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001823
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001824 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1825 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001826
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001827 source = PyObject_GetIter(args);
1828 if (source == NULL)
1829 return NULL;
1830
1831 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001832}
1833
1834static PyObject *
1835chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1836{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001837 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001838
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001839 source = PyObject_GetIter(arg);
1840 if (source == NULL)
1841 return NULL;
1842
1843 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001844}
1845
1846static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001847chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001848{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001849 PyObject_GC_UnTrack(lz);
1850 Py_XDECREF(lz->active);
1851 Py_XDECREF(lz->source);
1852 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001853}
1854
Raymond Hettinger2012f172003-02-07 05:32:58 +00001855static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001856chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001857{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001858 Py_VISIT(lz->source);
1859 Py_VISIT(lz->active);
1860 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001861}
1862
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001863static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001864chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001865{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001866 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001867
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001868 if (lz->source == NULL)
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001869 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001870
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001871 if (lz->active == NULL) {
1872 PyObject *iterable = PyIter_Next(lz->source);
1873 if (iterable == NULL) {
1874 Py_CLEAR(lz->source);
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001875 return NULL; /* no more input sources */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001876 }
1877 lz->active = PyObject_GetIter(iterable);
1878 Py_DECREF(iterable);
1879 if (lz->active == NULL) {
1880 Py_CLEAR(lz->source);
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001881 return NULL; /* input not iterable */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001882 }
1883 }
Raymond Hettingerb5244a32015-08-16 14:24:20 -07001884 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001885 if (item != NULL)
1886 return item;
1887 if (PyErr_Occurred()) {
1888 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1889 PyErr_Clear();
1890 else
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001891 return NULL; /* input raised an exception */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001892 }
1893 Py_CLEAR(lz->active);
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001894 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001895}
1896
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001897static PyObject *
1898chain_reduce(chainobject *lz)
1899{
1900 if (lz->source) {
1901 /* we can't pickle function objects (itertools.from_iterable) so
1902 * we must use setstate to replace the iterable. One day we
1903 * will fix pickling of functions
1904 */
1905 if (lz->active) {
1906 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1907 } else {
1908 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1909 }
1910 } else {
1911 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1912 }
1913 return NULL;
1914}
1915
1916static PyObject *
1917chain_setstate(chainobject *lz, PyObject *state)
1918{
1919 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001920
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001921 if (! PyArg_ParseTuple(state, "O|O", &source, &active))
1922 return NULL;
1923
1924 Py_CLEAR(lz->source);
1925 lz->source = source;
1926 Py_INCREF(lz->source);
1927 Py_CLEAR(lz->active);
1928 lz->active = active;
1929 Py_XINCREF(lz->active);
1930 Py_RETURN_NONE;
1931}
1932
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001933PyDoc_STRVAR(chain_doc,
1934"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001935\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001936Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001937first iterable until it is exhausted, then elements from the next\n\
1938iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001939
Christian Heimesf16baeb2008-02-29 14:57:44 +00001940PyDoc_STRVAR(chain_from_iterable_doc,
1941"chain.from_iterable(iterable) --> chain object\n\
1942\n\
1943Alternate chain() contructor taking a single iterable argument\n\
1944that evaluates lazily.");
1945
1946static PyMethodDef chain_methods[] = {
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001947 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1948 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001949 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1950 reduce_doc},
1951 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1952 setstate_doc},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001953 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001954};
1955
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001956static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001957 PyVarObject_HEAD_INIT(NULL, 0)
1958 "itertools.chain", /* tp_name */
1959 sizeof(chainobject), /* tp_basicsize */
1960 0, /* tp_itemsize */
1961 /* methods */
1962 (destructor)chain_dealloc, /* tp_dealloc */
1963 0, /* tp_print */
1964 0, /* tp_getattr */
1965 0, /* tp_setattr */
1966 0, /* tp_reserved */
1967 0, /* tp_repr */
1968 0, /* tp_as_number */
1969 0, /* tp_as_sequence */
1970 0, /* tp_as_mapping */
1971 0, /* tp_hash */
1972 0, /* tp_call */
1973 0, /* tp_str */
1974 PyObject_GenericGetAttr, /* tp_getattro */
1975 0, /* tp_setattro */
1976 0, /* tp_as_buffer */
1977 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1978 Py_TPFLAGS_BASETYPE, /* tp_flags */
1979 chain_doc, /* tp_doc */
1980 (traverseproc)chain_traverse, /* tp_traverse */
1981 0, /* tp_clear */
1982 0, /* tp_richcompare */
1983 0, /* tp_weaklistoffset */
1984 PyObject_SelfIter, /* tp_iter */
1985 (iternextfunc)chain_next, /* tp_iternext */
1986 chain_methods, /* tp_methods */
1987 0, /* tp_members */
1988 0, /* tp_getset */
1989 0, /* tp_base */
1990 0, /* tp_dict */
1991 0, /* tp_descr_get */
1992 0, /* tp_descr_set */
1993 0, /* tp_dictoffset */
1994 0, /* tp_init */
1995 0, /* tp_alloc */
1996 chain_new, /* tp_new */
1997 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001998};
1999
2000
Christian Heimesc3f30c42008-02-22 16:37:40 +00002001/* product object ************************************************************/
2002
2003typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002004 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002005 PyObject *pools; /* tuple of pool tuples */
2006 Py_ssize_t *indices; /* one index per pool */
2007 PyObject *result; /* most recently returned result tuple */
2008 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002009} productobject;
2010
2011static PyTypeObject product_type;
2012
2013static PyObject *
2014product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2015{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002016 productobject *lz;
2017 Py_ssize_t nargs, npools, repeat=1;
2018 PyObject *pools = NULL;
2019 Py_ssize_t *indices = NULL;
2020 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002021
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002022 if (kwds != NULL) {
2023 char *kwlist[] = {"repeat", 0};
2024 PyObject *tmpargs = PyTuple_New(0);
2025 if (tmpargs == NULL)
2026 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002027 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2028 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002029 Py_DECREF(tmpargs);
2030 return NULL;
2031 }
2032 Py_DECREF(tmpargs);
2033 if (repeat < 0) {
2034 PyErr_SetString(PyExc_ValueError,
2035 "repeat argument cannot be negative");
2036 return NULL;
2037 }
2038 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002039
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002040 assert(PyTuple_CheckExact(args));
2041 if (repeat == 0) {
2042 nargs = 0;
2043 } else {
2044 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002045 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002046 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2047 return NULL;
2048 }
2049 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002050 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002051
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002052 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002053 if (indices == NULL) {
2054 PyErr_NoMemory();
2055 goto error;
2056 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002057
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002058 pools = PyTuple_New(npools);
2059 if (pools == NULL)
2060 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002061
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002062 for (i=0; i < nargs ; ++i) {
2063 PyObject *item = PyTuple_GET_ITEM(args, i);
2064 PyObject *pool = PySequence_Tuple(item);
2065 if (pool == NULL)
2066 goto error;
2067 PyTuple_SET_ITEM(pools, i, pool);
2068 indices[i] = 0;
2069 }
2070 for ( ; i < npools; ++i) {
2071 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2072 Py_INCREF(pool);
2073 PyTuple_SET_ITEM(pools, i, pool);
2074 indices[i] = 0;
2075 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002076
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002077 /* create productobject structure */
2078 lz = (productobject *)type->tp_alloc(type, 0);
2079 if (lz == NULL)
2080 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002081
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002082 lz->pools = pools;
2083 lz->indices = indices;
2084 lz->result = NULL;
2085 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002086
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002087 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002088
2089error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002090 if (indices != NULL)
2091 PyMem_Free(indices);
2092 Py_XDECREF(pools);
2093 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002094}
2095
2096static void
2097product_dealloc(productobject *lz)
2098{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002099 PyObject_GC_UnTrack(lz);
2100 Py_XDECREF(lz->pools);
2101 Py_XDECREF(lz->result);
2102 if (lz->indices != NULL)
2103 PyMem_Free(lz->indices);
2104 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002105}
2106
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002107static PyObject *
2108product_sizeof(productobject *lz, void *unused)
2109{
2110 Py_ssize_t res;
2111
2112 res = sizeof(productobject);
2113 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2114 return PyLong_FromSsize_t(res);
2115}
2116
2117PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2118
Christian Heimesc3f30c42008-02-22 16:37:40 +00002119static int
2120product_traverse(productobject *lz, visitproc visit, void *arg)
2121{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002122 Py_VISIT(lz->pools);
2123 Py_VISIT(lz->result);
2124 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002125}
2126
2127static PyObject *
2128product_next(productobject *lz)
2129{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002130 PyObject *pool;
2131 PyObject *elem;
2132 PyObject *oldelem;
2133 PyObject *pools = lz->pools;
2134 PyObject *result = lz->result;
2135 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2136 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002137
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002138 if (lz->stopped)
2139 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002140
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002141 if (result == NULL) {
2142 /* On the first pass, return an initial tuple filled with the
2143 first element from each pool. */
2144 result = PyTuple_New(npools);
2145 if (result == NULL)
2146 goto empty;
2147 lz->result = result;
2148 for (i=0; i < npools; i++) {
2149 pool = PyTuple_GET_ITEM(pools, i);
2150 if (PyTuple_GET_SIZE(pool) == 0)
2151 goto empty;
2152 elem = PyTuple_GET_ITEM(pool, 0);
2153 Py_INCREF(elem);
2154 PyTuple_SET_ITEM(result, i, elem);
2155 }
2156 } else {
2157 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002158
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002159 /* Copy the previous result tuple or re-use it if available */
2160 if (Py_REFCNT(result) > 1) {
2161 PyObject *old_result = result;
2162 result = PyTuple_New(npools);
2163 if (result == NULL)
2164 goto empty;
2165 lz->result = result;
2166 for (i=0; i < npools; i++) {
2167 elem = PyTuple_GET_ITEM(old_result, i);
2168 Py_INCREF(elem);
2169 PyTuple_SET_ITEM(result, i, elem);
2170 }
2171 Py_DECREF(old_result);
2172 }
2173 /* Now, we've got the only copy so we can update it in-place */
2174 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002175
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002176 /* Update the pool indices right-to-left. Only advance to the
2177 next pool when the previous one rolls-over */
2178 for (i=npools-1 ; i >= 0 ; i--) {
2179 pool = PyTuple_GET_ITEM(pools, i);
2180 indices[i]++;
2181 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2182 /* Roll-over and advance to next pool */
2183 indices[i] = 0;
2184 elem = PyTuple_GET_ITEM(pool, 0);
2185 Py_INCREF(elem);
2186 oldelem = PyTuple_GET_ITEM(result, i);
2187 PyTuple_SET_ITEM(result, i, elem);
2188 Py_DECREF(oldelem);
2189 } else {
2190 /* No rollover. Just increment and stop here. */
2191 elem = PyTuple_GET_ITEM(pool, indices[i]);
2192 Py_INCREF(elem);
2193 oldelem = PyTuple_GET_ITEM(result, i);
2194 PyTuple_SET_ITEM(result, i, elem);
2195 Py_DECREF(oldelem);
2196 break;
2197 }
2198 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002199
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002200 /* If i is negative, then the indices have all rolled-over
2201 and we're done. */
2202 if (i < 0)
2203 goto empty;
2204 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002205
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002206 Py_INCREF(result);
2207 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002208
2209empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002210 lz->stopped = 1;
2211 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002212}
2213
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002214static PyObject *
2215product_reduce(productobject *lz)
2216{
2217 if (lz->stopped) {
2218 return Py_BuildValue("O(())", Py_TYPE(lz));
2219 } else if (lz->result == NULL) {
2220 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2221 } else {
2222 PyObject *indices;
2223 Py_ssize_t n, i;
2224
2225 /* we must pickle the indices use them for setstate, and
2226 * additionally indicate that the iterator has started
2227 */
2228 n = PyTuple_GET_SIZE(lz->pools);
2229 indices = PyTuple_New(n);
2230 if (indices == NULL)
2231 return NULL;
2232 for (i=0; i<n; i++){
2233 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2234 if (!index) {
2235 Py_DECREF(indices);
2236 return NULL;
2237 }
2238 PyTuple_SET_ITEM(indices, i, index);
2239 }
2240 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2241 }
2242}
2243
2244static PyObject *
2245product_setstate(productobject *lz, PyObject *state)
2246{
2247 PyObject *result;
2248 Py_ssize_t n, i;
2249
2250 n = PyTuple_GET_SIZE(lz->pools);
2251 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2252 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2253 return NULL;
2254 }
2255 for (i=0; i<n; i++)
2256 {
2257 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2258 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002259 PyObject* pool;
2260 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002261 if (index < 0 && PyErr_Occurred())
2262 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002263 pool = PyTuple_GET_ITEM(lz->pools, i);
2264 poolsize = PyTuple_GET_SIZE(pool);
2265 if (poolsize == 0) {
2266 lz->stopped = 1;
2267 Py_RETURN_NONE;
2268 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002269 /* clamp the index */
2270 if (index < 0)
2271 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002272 else if (index > poolsize-1)
2273 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002274 lz->indices[i] = index;
2275 }
2276
2277 result = PyTuple_New(n);
2278 if (!result)
2279 return NULL;
2280 for (i=0; i<n; i++) {
2281 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2282 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2283 Py_INCREF(element);
2284 PyTuple_SET_ITEM(result, i, element);
2285 }
2286 Py_CLEAR(lz->result);
2287 lz->result = result;
2288 Py_RETURN_NONE;
2289}
2290
2291static PyMethodDef product_methods[] = {
2292 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2293 reduce_doc},
2294 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2295 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002296 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2297 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002298 {NULL, NULL} /* sentinel */
2299};
2300
Christian Heimesc3f30c42008-02-22 16:37:40 +00002301PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002302"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002303\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002304Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002305For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2306The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2307cycle in a manner similar to an odometer (with the rightmost element changing\n\
2308on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002309To compute the product of an iterable with itself, specify the number\n\
2310of repetitions with the optional repeat keyword argument. For example,\n\
2311product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002312product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2313product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2314
2315static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002316 PyVarObject_HEAD_INIT(NULL, 0)
2317 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002318 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002319 0, /* tp_itemsize */
2320 /* methods */
2321 (destructor)product_dealloc, /* tp_dealloc */
2322 0, /* tp_print */
2323 0, /* tp_getattr */
2324 0, /* tp_setattr */
2325 0, /* tp_reserved */
2326 0, /* tp_repr */
2327 0, /* tp_as_number */
2328 0, /* tp_as_sequence */
2329 0, /* tp_as_mapping */
2330 0, /* tp_hash */
2331 0, /* tp_call */
2332 0, /* tp_str */
2333 PyObject_GenericGetAttr, /* tp_getattro */
2334 0, /* tp_setattro */
2335 0, /* tp_as_buffer */
2336 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2337 Py_TPFLAGS_BASETYPE, /* tp_flags */
2338 product_doc, /* tp_doc */
2339 (traverseproc)product_traverse, /* tp_traverse */
2340 0, /* tp_clear */
2341 0, /* tp_richcompare */
2342 0, /* tp_weaklistoffset */
2343 PyObject_SelfIter, /* tp_iter */
2344 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002345 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002346 0, /* tp_members */
2347 0, /* tp_getset */
2348 0, /* tp_base */
2349 0, /* tp_dict */
2350 0, /* tp_descr_get */
2351 0, /* tp_descr_set */
2352 0, /* tp_dictoffset */
2353 0, /* tp_init */
2354 0, /* tp_alloc */
2355 product_new, /* tp_new */
2356 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002357};
2358
2359
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002360/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002361
2362typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002363 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002364 PyObject *pool; /* input converted to a tuple */
2365 Py_ssize_t *indices; /* one index per result element */
2366 PyObject *result; /* most recently returned result tuple */
2367 Py_ssize_t r; /* size of result tuple */
2368 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002369} combinationsobject;
2370
2371static PyTypeObject combinations_type;
2372
2373static PyObject *
2374combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2375{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002376 combinationsobject *co;
2377 Py_ssize_t n;
2378 Py_ssize_t r;
2379 PyObject *pool = NULL;
2380 PyObject *iterable = NULL;
2381 Py_ssize_t *indices = NULL;
2382 Py_ssize_t i;
2383 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002384
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002385 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2386 &iterable, &r))
2387 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002388
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002389 pool = PySequence_Tuple(iterable);
2390 if (pool == NULL)
2391 goto error;
2392 n = PyTuple_GET_SIZE(pool);
2393 if (r < 0) {
2394 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2395 goto error;
2396 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002397
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002398 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002399 if (indices == NULL) {
2400 PyErr_NoMemory();
2401 goto error;
2402 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002403
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002404 for (i=0 ; i<r ; i++)
2405 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002406
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002407 /* create combinationsobject structure */
2408 co = (combinationsobject *)type->tp_alloc(type, 0);
2409 if (co == NULL)
2410 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002411
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002412 co->pool = pool;
2413 co->indices = indices;
2414 co->result = NULL;
2415 co->r = r;
2416 co->stopped = r > n ? 1 : 0;
2417
2418 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002419
2420error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002421 if (indices != NULL)
2422 PyMem_Free(indices);
2423 Py_XDECREF(pool);
2424 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002425}
2426
2427static void
2428combinations_dealloc(combinationsobject *co)
2429{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002430 PyObject_GC_UnTrack(co);
2431 Py_XDECREF(co->pool);
2432 Py_XDECREF(co->result);
2433 if (co->indices != NULL)
2434 PyMem_Free(co->indices);
2435 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002436}
2437
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002438static PyObject *
2439combinations_sizeof(combinationsobject *co, void *unused)
2440{
2441 Py_ssize_t res;
2442
2443 res = sizeof(combinationsobject);
2444 res += co->r * sizeof(Py_ssize_t);
2445 return PyLong_FromSsize_t(res);
2446}
2447
Christian Heimes380f7f22008-02-28 11:19:05 +00002448static int
2449combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2450{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002451 Py_VISIT(co->pool);
2452 Py_VISIT(co->result);
2453 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002454}
2455
2456static PyObject *
2457combinations_next(combinationsobject *co)
2458{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002459 PyObject *elem;
2460 PyObject *oldelem;
2461 PyObject *pool = co->pool;
2462 Py_ssize_t *indices = co->indices;
2463 PyObject *result = co->result;
2464 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2465 Py_ssize_t r = co->r;
2466 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002467
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002468 if (co->stopped)
2469 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002470
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002471 if (result == NULL) {
2472 /* On the first pass, initialize result tuple using the indices */
2473 result = PyTuple_New(r);
2474 if (result == NULL)
2475 goto empty;
2476 co->result = result;
2477 for (i=0; i<r ; i++) {
2478 index = indices[i];
2479 elem = PyTuple_GET_ITEM(pool, index);
2480 Py_INCREF(elem);
2481 PyTuple_SET_ITEM(result, i, elem);
2482 }
2483 } else {
2484 /* Copy the previous result tuple or re-use it if available */
2485 if (Py_REFCNT(result) > 1) {
2486 PyObject *old_result = result;
2487 result = PyTuple_New(r);
2488 if (result == NULL)
2489 goto empty;
2490 co->result = result;
2491 for (i=0; i<r ; i++) {
2492 elem = PyTuple_GET_ITEM(old_result, i);
2493 Py_INCREF(elem);
2494 PyTuple_SET_ITEM(result, i, elem);
2495 }
2496 Py_DECREF(old_result);
2497 }
2498 /* Now, we've got the only copy so we can update it in-place
2499 * CPython's empty tuple is a singleton and cached in
2500 * PyTuple's freelist.
2501 */
2502 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002503
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002504 /* Scan indices right-to-left until finding one that is not
2505 at its maximum (i + n - r). */
2506 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2507 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002508
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002509 /* If i is negative, then the indices are all at
2510 their maximum value and we're done. */
2511 if (i < 0)
2512 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002513
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002514 /* Increment the current index which we know is not at its
2515 maximum. Then move back to the right setting each index
2516 to its lowest possible value (one higher than the index
2517 to its left -- this maintains the sort order invariant). */
2518 indices[i]++;
2519 for (j=i+1 ; j<r ; j++)
2520 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002521
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002522 /* Update the result tuple for the new indices
2523 starting with i, the leftmost index that changed */
2524 for ( ; i<r ; i++) {
2525 index = indices[i];
2526 elem = PyTuple_GET_ITEM(pool, index);
2527 Py_INCREF(elem);
2528 oldelem = PyTuple_GET_ITEM(result, i);
2529 PyTuple_SET_ITEM(result, i, elem);
2530 Py_DECREF(oldelem);
2531 }
2532 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002533
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002534 Py_INCREF(result);
2535 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002536
2537empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002538 co->stopped = 1;
2539 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002540}
2541
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002542static PyObject *
2543combinations_reduce(combinationsobject *lz)
2544{
2545 if (lz->result == NULL) {
2546 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2547 } else if (lz->stopped) {
2548 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2549 } else {
2550 PyObject *indices;
2551 Py_ssize_t i;
2552
2553 /* we must pickle the indices and use them for setstate */
2554 indices = PyTuple_New(lz->r);
2555 if (!indices)
2556 return NULL;
2557 for (i=0; i<lz->r; i++)
2558 {
2559 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2560 if (!index) {
2561 Py_DECREF(indices);
2562 return NULL;
2563 }
2564 PyTuple_SET_ITEM(indices, i, index);
2565 }
2566
2567 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2568 }
2569}
2570
2571static PyObject *
2572combinations_setstate(combinationsobject *lz, PyObject *state)
2573{
2574 PyObject *result;
2575 Py_ssize_t i;
2576 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2577
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002578 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002579 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2580 return NULL;
2581 }
2582
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002583 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002584 Py_ssize_t max;
2585 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2586 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002587
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002588 if (index == -1 && PyErr_Occurred())
2589 return NULL; /* not an integer */
2590 max = i + n - lz->r;
2591 /* clamp the index (beware of negative max) */
2592 if (index > max)
2593 index = max;
2594 if (index < 0)
2595 index = 0;
2596 lz->indices[i] = index;
2597 }
2598
2599 result = PyTuple_New(lz->r);
2600 if (result == NULL)
2601 return NULL;
2602 for (i=0; i<lz->r; i++) {
2603 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2604 Py_INCREF(element);
2605 PyTuple_SET_ITEM(result, i, element);
2606 }
2607
2608 Py_CLEAR(lz->result);
2609 lz->result = result;
2610 Py_RETURN_NONE;
2611}
2612
2613static PyMethodDef combinations_methods[] = {
2614 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2615 reduce_doc},
2616 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2617 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002618 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2619 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002620 {NULL, NULL} /* sentinel */
2621};
2622
Christian Heimes380f7f22008-02-28 11:19:05 +00002623PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002624"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002625\n\
2626Return successive r-length combinations of elements in the iterable.\n\n\
2627combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2628
2629static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002630 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002631 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002632 sizeof(combinationsobject), /* tp_basicsize */
2633 0, /* tp_itemsize */
2634 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002635 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002636 0, /* tp_print */
2637 0, /* tp_getattr */
2638 0, /* tp_setattr */
2639 0, /* tp_reserved */
2640 0, /* tp_repr */
2641 0, /* tp_as_number */
2642 0, /* tp_as_sequence */
2643 0, /* tp_as_mapping */
2644 0, /* tp_hash */
2645 0, /* tp_call */
2646 0, /* tp_str */
2647 PyObject_GenericGetAttr, /* tp_getattro */
2648 0, /* tp_setattro */
2649 0, /* tp_as_buffer */
2650 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2651 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002652 combinations_doc, /* tp_doc */
2653 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002654 0, /* tp_clear */
2655 0, /* tp_richcompare */
2656 0, /* tp_weaklistoffset */
2657 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002658 (iternextfunc)combinations_next, /* tp_iternext */
2659 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002660 0, /* tp_members */
2661 0, /* tp_getset */
2662 0, /* tp_base */
2663 0, /* tp_dict */
2664 0, /* tp_descr_get */
2665 0, /* tp_descr_set */
2666 0, /* tp_dictoffset */
2667 0, /* tp_init */
2668 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002669 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002670 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002671};
2672
2673
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002674/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002675
2676/* Equivalent to:
2677
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002678 def combinations_with_replacement(iterable, r):
2679 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2680 # number items returned: (n+r-1)! / r! / (n-1)!
2681 pool = tuple(iterable)
2682 n = len(pool)
2683 indices = [0] * r
2684 yield tuple(pool[i] for i in indices)
2685 while 1:
2686 for i in reversed(range(r)):
2687 if indices[i] != n - 1:
2688 break
2689 else:
2690 return
2691 indices[i:] = [indices[i] + 1] * (r - i)
2692 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002693
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002694 def combinations_with_replacement2(iterable, r):
2695 'Alternate version that filters from product()'
2696 pool = tuple(iterable)
2697 n = len(pool)
2698 for indices in product(range(n), repeat=r):
2699 if sorted(indices) == list(indices):
2700 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002701*/
2702typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002703 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002704 PyObject *pool; /* input converted to a tuple */
2705 Py_ssize_t *indices; /* one index per result element */
2706 PyObject *result; /* most recently returned result tuple */
2707 Py_ssize_t r; /* size of result tuple */
2708 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002709} cwrobject;
2710
2711static PyTypeObject cwr_type;
2712
2713static PyObject *
2714cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2715{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002716 cwrobject *co;
2717 Py_ssize_t n;
2718 Py_ssize_t r;
2719 PyObject *pool = NULL;
2720 PyObject *iterable = NULL;
2721 Py_ssize_t *indices = NULL;
2722 Py_ssize_t i;
2723 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002724
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002725 if (!PyArg_ParseTupleAndKeywords(args, kwds,
2726 "On:combinations_with_replacement",
2727 kwargs, &iterable, &r))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002728 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002729
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002730 pool = PySequence_Tuple(iterable);
2731 if (pool == NULL)
2732 goto error;
2733 n = PyTuple_GET_SIZE(pool);
2734 if (r < 0) {
2735 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2736 goto error;
2737 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002738
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002739 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002740 if (indices == NULL) {
2741 PyErr_NoMemory();
2742 goto error;
2743 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002744
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002745 for (i=0 ; i<r ; i++)
2746 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002747
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002748 /* create cwrobject structure */
2749 co = (cwrobject *)type->tp_alloc(type, 0);
2750 if (co == NULL)
2751 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002752
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002753 co->pool = pool;
2754 co->indices = indices;
2755 co->result = NULL;
2756 co->r = r;
2757 co->stopped = !n && r;
2758
2759 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002760
2761error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002762 if (indices != NULL)
2763 PyMem_Free(indices);
2764 Py_XDECREF(pool);
2765 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002766}
2767
2768static void
2769cwr_dealloc(cwrobject *co)
2770{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002771 PyObject_GC_UnTrack(co);
2772 Py_XDECREF(co->pool);
2773 Py_XDECREF(co->result);
2774 if (co->indices != NULL)
2775 PyMem_Free(co->indices);
2776 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002777}
2778
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002779static PyObject *
2780cwr_sizeof(cwrobject *co, void *unused)
2781{
2782 Py_ssize_t res;
2783
2784 res = sizeof(cwrobject);
2785 res += co->r * sizeof(Py_ssize_t);
2786 return PyLong_FromSsize_t(res);
2787}
2788
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002789static int
2790cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2791{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002792 Py_VISIT(co->pool);
2793 Py_VISIT(co->result);
2794 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002795}
2796
2797static PyObject *
2798cwr_next(cwrobject *co)
2799{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002800 PyObject *elem;
2801 PyObject *oldelem;
2802 PyObject *pool = co->pool;
2803 Py_ssize_t *indices = co->indices;
2804 PyObject *result = co->result;
2805 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2806 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002807 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002808
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002809 if (co->stopped)
2810 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002811
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002812 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002813 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002814 result = PyTuple_New(r);
2815 if (result == NULL)
2816 goto empty;
2817 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002818 if (n > 0) {
2819 elem = PyTuple_GET_ITEM(pool, 0);
2820 for (i=0; i<r ; i++) {
2821 assert(indices[i] == 0);
2822 Py_INCREF(elem);
2823 PyTuple_SET_ITEM(result, i, elem);
2824 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002825 }
2826 } else {
2827 /* Copy the previous result tuple or re-use it if available */
2828 if (Py_REFCNT(result) > 1) {
2829 PyObject *old_result = result;
2830 result = PyTuple_New(r);
2831 if (result == NULL)
2832 goto empty;
2833 co->result = result;
2834 for (i=0; i<r ; i++) {
2835 elem = PyTuple_GET_ITEM(old_result, i);
2836 Py_INCREF(elem);
2837 PyTuple_SET_ITEM(result, i, elem);
2838 }
2839 Py_DECREF(old_result);
2840 }
2841 /* Now, we've got the only copy so we can update it in-place CPython's
2842 empty tuple is a singleton and cached in PyTuple's freelist. */
2843 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002844
Tim Peters9edb1682013-09-03 11:49:31 -05002845 /* Scan indices right-to-left until finding one that is not
2846 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002847 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2848 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002849
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002850 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002851 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002852 if (i < 0)
2853 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002854
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002855 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002856 maximum. Then set all to the right to the same value. */
2857 index = indices[i] + 1;
2858 assert(index < n);
2859 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002860 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002861 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002862 Py_INCREF(elem);
2863 oldelem = PyTuple_GET_ITEM(result, i);
2864 PyTuple_SET_ITEM(result, i, elem);
2865 Py_DECREF(oldelem);
2866 }
2867 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002868
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002869 Py_INCREF(result);
2870 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002871
2872empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002873 co->stopped = 1;
2874 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002875}
2876
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002877static PyObject *
2878cwr_reduce(cwrobject *lz)
2879{
2880 if (lz->result == NULL) {
2881 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2882 } else if (lz->stopped) {
2883 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2884 } else {
2885 PyObject *indices;
2886 Py_ssize_t i;
2887
2888 /* we must pickle the indices and use them for setstate */
2889 indices = PyTuple_New(lz->r);
2890 if (!indices)
2891 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002892 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002893 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2894 if (!index) {
2895 Py_DECREF(indices);
2896 return NULL;
2897 }
2898 PyTuple_SET_ITEM(indices, i, index);
2899 }
2900
2901 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2902 }
2903}
2904
2905static PyObject *
2906cwr_setstate(cwrobject *lz, PyObject *state)
2907{
2908 PyObject *result;
2909 Py_ssize_t n, i;
2910
2911 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2912 {
2913 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2914 return NULL;
2915 }
2916
2917 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002918 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002919 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2920 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002921
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002922 if (index < 0 && PyErr_Occurred())
2923 return NULL; /* not an integer */
2924 /* clamp the index */
2925 if (index < 0)
2926 index = 0;
2927 else if (index > n-1)
2928 index = n-1;
2929 lz->indices[i] = index;
2930 }
2931 result = PyTuple_New(lz->r);
2932 if (result == NULL)
2933 return NULL;
2934 for (i=0; i<lz->r; i++) {
2935 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2936 Py_INCREF(element);
2937 PyTuple_SET_ITEM(result, i, element);
2938 }
2939 Py_CLEAR(lz->result);
2940 lz->result = result;
2941 Py_RETURN_NONE;
2942}
2943
2944static PyMethodDef cwr_methods[] = {
2945 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2946 reduce_doc},
2947 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2948 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002949 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
2950 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002951 {NULL, NULL} /* sentinel */
2952};
2953
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002954PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002955"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002956\n\
2957Return successive r-length combinations of elements in the iterable\n\
2958allowing individual elements to have successive repeats.\n\
2959combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2960
2961static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002962 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002963 "itertools.combinations_with_replacement", /* tp_name */
2964 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002965 0, /* tp_itemsize */
2966 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002967 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002968 0, /* tp_print */
2969 0, /* tp_getattr */
2970 0, /* tp_setattr */
2971 0, /* tp_reserved */
2972 0, /* tp_repr */
2973 0, /* tp_as_number */
2974 0, /* tp_as_sequence */
2975 0, /* tp_as_mapping */
2976 0, /* tp_hash */
2977 0, /* tp_call */
2978 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002979 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002980 0, /* tp_setattro */
2981 0, /* tp_as_buffer */
2982 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002983 Py_TPFLAGS_BASETYPE, /* tp_flags */
2984 cwr_doc, /* tp_doc */
2985 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002986 0, /* tp_clear */
2987 0, /* tp_richcompare */
2988 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002989 PyObject_SelfIter, /* tp_iter */
2990 (iternextfunc)cwr_next, /* tp_iternext */
2991 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002992 0, /* tp_members */
2993 0, /* tp_getset */
2994 0, /* tp_base */
2995 0, /* tp_dict */
2996 0, /* tp_descr_get */
2997 0, /* tp_descr_set */
2998 0, /* tp_dictoffset */
2999 0, /* tp_init */
3000 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003001 cwr_new, /* tp_new */
3002 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003003};
3004
3005
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003006/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003007
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003008def permutations(iterable, r=None):
3009 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
3010 pool = tuple(iterable)
3011 n = len(pool)
3012 r = n if r is None else r
3013 indices = range(n)
3014 cycles = range(n-r+1, n+1)[::-1]
3015 yield tuple(pool[i] for i in indices[:r])
3016 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003017 for i in reversed(range(r)):
3018 cycles[i] -= 1
3019 if cycles[i] == 0:
3020 indices[i:] = indices[i+1:] + indices[i:i+1]
3021 cycles[i] = n - i
3022 else:
3023 j = cycles[i]
3024 indices[i], indices[-j] = indices[-j], indices[i]
3025 yield tuple(pool[i] for i in indices[:r])
3026 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003027 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003028 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003029*/
3030
3031typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003032 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003033 PyObject *pool; /* input converted to a tuple */
3034 Py_ssize_t *indices; /* one index per element in the pool */
3035 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3036 PyObject *result; /* most recently returned result tuple */
3037 Py_ssize_t r; /* size of result tuple */
3038 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003039} permutationsobject;
3040
3041static PyTypeObject permutations_type;
3042
3043static PyObject *
3044permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3045{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003046 permutationsobject *po;
3047 Py_ssize_t n;
3048 Py_ssize_t r;
3049 PyObject *robj = Py_None;
3050 PyObject *pool = NULL;
3051 PyObject *iterable = NULL;
3052 Py_ssize_t *indices = NULL;
3053 Py_ssize_t *cycles = NULL;
3054 Py_ssize_t i;
3055 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003056
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003057 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
3058 &iterable, &robj))
3059 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003060
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003061 pool = PySequence_Tuple(iterable);
3062 if (pool == NULL)
3063 goto error;
3064 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003065
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003066 r = n;
3067 if (robj != Py_None) {
3068 if (!PyLong_Check(robj)) {
3069 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3070 goto error;
3071 }
3072 r = PyLong_AsSsize_t(robj);
3073 if (r == -1 && PyErr_Occurred())
3074 goto error;
3075 }
3076 if (r < 0) {
3077 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3078 goto error;
3079 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003080
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003081 indices = PyMem_New(Py_ssize_t, n);
3082 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003083 if (indices == NULL || cycles == NULL) {
3084 PyErr_NoMemory();
3085 goto error;
3086 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003087
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003088 for (i=0 ; i<n ; i++)
3089 indices[i] = i;
3090 for (i=0 ; i<r ; i++)
3091 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003092
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003093 /* create permutationsobject structure */
3094 po = (permutationsobject *)type->tp_alloc(type, 0);
3095 if (po == NULL)
3096 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003097
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003098 po->pool = pool;
3099 po->indices = indices;
3100 po->cycles = cycles;
3101 po->result = NULL;
3102 po->r = r;
3103 po->stopped = r > n ? 1 : 0;
3104
3105 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003106
3107error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003108 if (indices != NULL)
3109 PyMem_Free(indices);
3110 if (cycles != NULL)
3111 PyMem_Free(cycles);
3112 Py_XDECREF(pool);
3113 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003114}
3115
3116static void
3117permutations_dealloc(permutationsobject *po)
3118{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003119 PyObject_GC_UnTrack(po);
3120 Py_XDECREF(po->pool);
3121 Py_XDECREF(po->result);
3122 PyMem_Free(po->indices);
3123 PyMem_Free(po->cycles);
3124 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003125}
3126
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003127static PyObject *
3128permutations_sizeof(permutationsobject *po, void *unused)
3129{
3130 Py_ssize_t res;
3131
3132 res = sizeof(permutationsobject);
3133 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3134 res += po->r * sizeof(Py_ssize_t);
3135 return PyLong_FromSsize_t(res);
3136}
3137
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003138static int
3139permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3140{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003141 Py_VISIT(po->pool);
3142 Py_VISIT(po->result);
3143 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003144}
3145
3146static PyObject *
3147permutations_next(permutationsobject *po)
3148{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003149 PyObject *elem;
3150 PyObject *oldelem;
3151 PyObject *pool = po->pool;
3152 Py_ssize_t *indices = po->indices;
3153 Py_ssize_t *cycles = po->cycles;
3154 PyObject *result = po->result;
3155 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3156 Py_ssize_t r = po->r;
3157 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003158
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003159 if (po->stopped)
3160 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003161
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003162 if (result == NULL) {
3163 /* On the first pass, initialize result tuple using the indices */
3164 result = PyTuple_New(r);
3165 if (result == NULL)
3166 goto empty;
3167 po->result = result;
3168 for (i=0; i<r ; i++) {
3169 index = indices[i];
3170 elem = PyTuple_GET_ITEM(pool, index);
3171 Py_INCREF(elem);
3172 PyTuple_SET_ITEM(result, i, elem);
3173 }
3174 } else {
3175 if (n == 0)
3176 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003177
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003178 /* Copy the previous result tuple or re-use it if available */
3179 if (Py_REFCNT(result) > 1) {
3180 PyObject *old_result = result;
3181 result = PyTuple_New(r);
3182 if (result == NULL)
3183 goto empty;
3184 po->result = result;
3185 for (i=0; i<r ; i++) {
3186 elem = PyTuple_GET_ITEM(old_result, i);
3187 Py_INCREF(elem);
3188 PyTuple_SET_ITEM(result, i, elem);
3189 }
3190 Py_DECREF(old_result);
3191 }
3192 /* Now, we've got the only copy so we can update it in-place */
3193 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003194
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003195 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3196 for (i=r-1 ; i>=0 ; i--) {
3197 cycles[i] -= 1;
3198 if (cycles[i] == 0) {
3199 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3200 index = indices[i];
3201 for (j=i ; j<n-1 ; j++)
3202 indices[j] = indices[j+1];
3203 indices[n-1] = index;
3204 cycles[i] = n - i;
3205 } else {
3206 j = cycles[i];
3207 index = indices[i];
3208 indices[i] = indices[n-j];
3209 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003210
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003211 for (k=i; k<r ; k++) {
3212 /* start with i, the leftmost element that changed */
3213 /* yield tuple(pool[k] for k in indices[:r]) */
3214 index = indices[k];
3215 elem = PyTuple_GET_ITEM(pool, index);
3216 Py_INCREF(elem);
3217 oldelem = PyTuple_GET_ITEM(result, k);
3218 PyTuple_SET_ITEM(result, k, elem);
3219 Py_DECREF(oldelem);
3220 }
3221 break;
3222 }
3223 }
3224 /* If i is negative, then the cycles have all
3225 rolled-over and we're done. */
3226 if (i < 0)
3227 goto empty;
3228 }
3229 Py_INCREF(result);
3230 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003231
3232empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003233 po->stopped = 1;
3234 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003235}
3236
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003237static PyObject *
3238permutations_reduce(permutationsobject *po)
3239{
3240 if (po->result == NULL) {
3241 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3242 } else if (po->stopped) {
3243 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3244 } else {
3245 PyObject *indices=NULL, *cycles=NULL;
3246 Py_ssize_t n, i;
3247
3248 /* we must pickle the indices and cycles and use them for setstate */
3249 n = PyTuple_GET_SIZE(po->pool);
3250 indices = PyTuple_New(n);
3251 if (indices == NULL)
3252 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003253 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003254 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3255 if (!index)
3256 goto err;
3257 PyTuple_SET_ITEM(indices, i, index);
3258 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003259
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003260 cycles = PyTuple_New(po->r);
3261 if (cycles == NULL)
3262 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003263 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003264 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3265 if (!index)
3266 goto err;
3267 PyTuple_SET_ITEM(cycles, i, index);
3268 }
3269 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3270 po->pool, po->r,
3271 indices, cycles);
3272 err:
3273 Py_XDECREF(indices);
3274 Py_XDECREF(cycles);
3275 return NULL;
3276 }
3277}
3278
3279static PyObject *
3280permutations_setstate(permutationsobject *po, PyObject *state)
3281{
3282 PyObject *indices, *cycles, *result;
3283 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003284
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003285 if (!PyArg_ParseTuple(state, "O!O!",
3286 &PyTuple_Type, &indices,
3287 &PyTuple_Type, &cycles))
3288 return NULL;
3289
3290 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003291 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003292 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3293 return NULL;
3294 }
3295
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003296 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003297 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3298 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3299 if (index < 0 && PyErr_Occurred())
3300 return NULL; /* not an integer */
3301 /* clamp the index */
3302 if (index < 0)
3303 index = 0;
3304 else if (index > n-1)
3305 index = n-1;
3306 po->indices[i] = index;
3307 }
3308
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003309 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003310 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3311 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3312 if (index < 0 && PyErr_Occurred())
3313 return NULL; /* not an integer */
3314 if (index < 1)
3315 index = 1;
3316 else if (index > n-i)
3317 index = n-i;
3318 po->cycles[i] = index;
3319 }
3320 result = PyTuple_New(po->r);
3321 if (result == NULL)
3322 return NULL;
3323 for (i=0; i<po->r; i++) {
3324 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3325 Py_INCREF(element);
3326 PyTuple_SET_ITEM(result, i, element);
3327 }
3328 Py_CLEAR(po->result);
3329 po->result = result;
3330 Py_RETURN_NONE;
3331}
3332
3333static PyMethodDef permuations_methods[] = {
3334 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3335 reduce_doc},
3336 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3337 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003338 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3339 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003340 {NULL, NULL} /* sentinel */
3341};
3342
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003343PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003344"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003345\n\
3346Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003347permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003348
3349static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003350 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003351 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003352 sizeof(permutationsobject), /* tp_basicsize */
3353 0, /* tp_itemsize */
3354 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003355 (destructor)permutations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003356 0, /* tp_print */
3357 0, /* tp_getattr */
3358 0, /* tp_setattr */
3359 0, /* tp_reserved */
3360 0, /* tp_repr */
3361 0, /* tp_as_number */
3362 0, /* tp_as_sequence */
3363 0, /* tp_as_mapping */
3364 0, /* tp_hash */
3365 0, /* tp_call */
3366 0, /* tp_str */
3367 PyObject_GenericGetAttr, /* tp_getattro */
3368 0, /* tp_setattro */
3369 0, /* tp_as_buffer */
3370 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3371 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003372 permutations_doc, /* tp_doc */
3373 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003374 0, /* tp_clear */
3375 0, /* tp_richcompare */
3376 0, /* tp_weaklistoffset */
3377 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003378 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003379 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003380 0, /* tp_members */
3381 0, /* tp_getset */
3382 0, /* tp_base */
3383 0, /* tp_dict */
3384 0, /* tp_descr_get */
3385 0, /* tp_descr_set */
3386 0, /* tp_dictoffset */
3387 0, /* tp_init */
3388 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003389 permutations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003390 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003391};
3392
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003393/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003394
3395typedef struct {
3396 PyObject_HEAD
3397 PyObject *total;
3398 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003399 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003400} accumulateobject;
3401
3402static PyTypeObject accumulate_type;
3403
3404static PyObject *
3405accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3406{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003407 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003408 PyObject *iterable;
3409 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003410 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003411 accumulateobject *lz;
3412
Raymond Hettinger5d446132011-03-27 18:52:10 -07003413 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3414 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003415 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003416
3417 /* Get iterator. */
3418 it = PyObject_GetIter(iterable);
3419 if (it == NULL)
3420 return NULL;
3421
Raymond Hettinger482ba772010-12-01 22:48:00 +00003422 /* create accumulateobject structure */
3423 lz = (accumulateobject *)type->tp_alloc(type, 0);
3424 if (lz == NULL) {
3425 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003426 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003427 }
3428
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003429 if (binop != Py_None) {
3430 Py_XINCREF(binop);
3431 lz->binop = binop;
3432 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003433 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003434 lz->it = it;
3435 return (PyObject *)lz;
3436}
3437
3438static void
3439accumulate_dealloc(accumulateobject *lz)
3440{
3441 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003442 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003443 Py_XDECREF(lz->total);
3444 Py_XDECREF(lz->it);
3445 Py_TYPE(lz)->tp_free(lz);
3446}
3447
3448static int
3449accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3450{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003451 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003452 Py_VISIT(lz->it);
3453 Py_VISIT(lz->total);
3454 return 0;
3455}
3456
3457static PyObject *
3458accumulate_next(accumulateobject *lz)
3459{
3460 PyObject *val, *oldtotal, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003461
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003462 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003463 if (val == NULL)
3464 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003465
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003466 if (lz->total == NULL) {
3467 Py_INCREF(val);
3468 lz->total = val;
3469 return lz->total;
3470 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003471
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003472 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003473 newtotal = PyNumber_Add(lz->total, val);
3474 else
3475 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003476 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003477 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003478 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003479
3480 oldtotal = lz->total;
3481 lz->total = newtotal;
3482 Py_DECREF(oldtotal);
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003483
Raymond Hettinger482ba772010-12-01 22:48:00 +00003484 Py_INCREF(newtotal);
3485 return newtotal;
3486}
3487
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003488static PyObject *
3489accumulate_reduce(accumulateobject *lz)
3490{
3491 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3492 lz->it, lz->binop?lz->binop:Py_None,
3493 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003494}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003495
3496static PyObject *
3497accumulate_setstate(accumulateobject *lz, PyObject *state)
3498{
3499 Py_CLEAR(lz->total);
3500 lz->total = state;
3501 Py_INCREF(lz->total);
3502 Py_RETURN_NONE;
3503}
3504
3505static PyMethodDef accumulate_methods[] = {
3506 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3507 reduce_doc},
3508 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3509 setstate_doc},
3510 {NULL, NULL} /* sentinel */
3511};
3512
Raymond Hettinger482ba772010-12-01 22:48:00 +00003513PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003514"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003515\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003516Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003517
3518static PyTypeObject accumulate_type = {
3519 PyVarObject_HEAD_INIT(NULL, 0)
3520 "itertools.accumulate", /* tp_name */
3521 sizeof(accumulateobject), /* tp_basicsize */
3522 0, /* tp_itemsize */
3523 /* methods */
3524 (destructor)accumulate_dealloc, /* tp_dealloc */
3525 0, /* tp_print */
3526 0, /* tp_getattr */
3527 0, /* tp_setattr */
3528 0, /* tp_reserved */
3529 0, /* tp_repr */
3530 0, /* tp_as_number */
3531 0, /* tp_as_sequence */
3532 0, /* tp_as_mapping */
3533 0, /* tp_hash */
3534 0, /* tp_call */
3535 0, /* tp_str */
3536 PyObject_GenericGetAttr, /* tp_getattro */
3537 0, /* tp_setattro */
3538 0, /* tp_as_buffer */
3539 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3540 Py_TPFLAGS_BASETYPE, /* tp_flags */
3541 accumulate_doc, /* tp_doc */
3542 (traverseproc)accumulate_traverse, /* tp_traverse */
3543 0, /* tp_clear */
3544 0, /* tp_richcompare */
3545 0, /* tp_weaklistoffset */
3546 PyObject_SelfIter, /* tp_iter */
3547 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003548 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003549 0, /* tp_members */
3550 0, /* tp_getset */
3551 0, /* tp_base */
3552 0, /* tp_dict */
3553 0, /* tp_descr_get */
3554 0, /* tp_descr_set */
3555 0, /* tp_dictoffset */
3556 0, /* tp_init */
3557 0, /* tp_alloc */
3558 accumulate_new, /* tp_new */
3559 PyObject_GC_Del, /* tp_free */
3560};
3561
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003562
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003563/* compress object ************************************************************/
3564
3565/* Equivalent to:
3566
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003567 def compress(data, selectors):
3568 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3569 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003570*/
3571
3572typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003573 PyObject_HEAD
3574 PyObject *data;
3575 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003576} compressobject;
3577
3578static PyTypeObject compress_type;
3579
3580static PyObject *
3581compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3582{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003583 PyObject *seq1, *seq2;
3584 PyObject *data=NULL, *selectors=NULL;
3585 compressobject *lz;
3586 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003587
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003588 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3589 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003590
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003591 data = PyObject_GetIter(seq1);
3592 if (data == NULL)
3593 goto fail;
3594 selectors = PyObject_GetIter(seq2);
3595 if (selectors == NULL)
3596 goto fail;
3597
3598 /* create compressobject structure */
3599 lz = (compressobject *)type->tp_alloc(type, 0);
3600 if (lz == NULL)
3601 goto fail;
3602 lz->data = data;
3603 lz->selectors = selectors;
3604 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003605
3606fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003607 Py_XDECREF(data);
3608 Py_XDECREF(selectors);
3609 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003610}
3611
3612static void
3613compress_dealloc(compressobject *lz)
3614{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003615 PyObject_GC_UnTrack(lz);
3616 Py_XDECREF(lz->data);
3617 Py_XDECREF(lz->selectors);
3618 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003619}
3620
3621static int
3622compress_traverse(compressobject *lz, visitproc visit, void *arg)
3623{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003624 Py_VISIT(lz->data);
3625 Py_VISIT(lz->selectors);
3626 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003627}
3628
3629static PyObject *
3630compress_next(compressobject *lz)
3631{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003632 PyObject *data = lz->data, *selectors = lz->selectors;
3633 PyObject *datum, *selector;
3634 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3635 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3636 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003637
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003638 while (1) {
3639 /* Steps: get datum, get selector, evaluate selector.
3640 Order is important (to match the pure python version
3641 in terms of which input gets a chance to raise an
3642 exception first).
3643 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003644
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003645 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003646 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003647 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003648
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003649 selector = selectornext(selectors);
3650 if (selector == NULL) {
3651 Py_DECREF(datum);
3652 return NULL;
3653 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003654
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003655 ok = PyObject_IsTrue(selector);
3656 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003657 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003658 return datum;
3659 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003660 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003661 return NULL;
3662 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003663}
3664
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003665static PyObject *
3666compress_reduce(compressobject *lz)
3667{
3668 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3669 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003670}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003671
3672static PyMethodDef compress_methods[] = {
3673 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3674 reduce_doc},
3675 {NULL, NULL} /* sentinel */
3676};
3677
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003678PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003679"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003680\n\
3681Return data elements corresponding to true selector elements.\n\
3682Forms a shorter iterator from selected data elements using the\n\
3683selectors to choose the data elements.");
3684
3685static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003686 PyVarObject_HEAD_INIT(NULL, 0)
3687 "itertools.compress", /* tp_name */
3688 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003689 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003690 /* methods */
3691 (destructor)compress_dealloc, /* tp_dealloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003692 0, /* tp_print */
3693 0, /* tp_getattr */
3694 0, /* tp_setattr */
3695 0, /* tp_reserved */
3696 0, /* tp_repr */
3697 0, /* tp_as_number */
3698 0, /* tp_as_sequence */
3699 0, /* tp_as_mapping */
3700 0, /* tp_hash */
3701 0, /* tp_call */
3702 0, /* tp_str */
3703 PyObject_GenericGetAttr, /* tp_getattro */
3704 0, /* tp_setattro */
3705 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003706 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003707 Py_TPFLAGS_BASETYPE, /* tp_flags */
3708 compress_doc, /* tp_doc */
3709 (traverseproc)compress_traverse, /* tp_traverse */
3710 0, /* tp_clear */
3711 0, /* tp_richcompare */
3712 0, /* tp_weaklistoffset */
3713 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003714 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003715 compress_methods, /* tp_methods */
3716 0, /* tp_members */
3717 0, /* tp_getset */
3718 0, /* tp_base */
3719 0, /* tp_dict */
3720 0, /* tp_descr_get */
3721 0, /* tp_descr_set */
3722 0, /* tp_dictoffset */
3723 0, /* tp_init */
3724 0, /* tp_alloc */
3725 compress_new, /* tp_new */
3726 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003727};
3728
3729
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003730/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003731
3732typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003733 PyObject_HEAD
3734 PyObject *func;
3735 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003736} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003737
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003738static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003739
3740static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003741filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003742{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003743 PyObject *func, *seq;
3744 PyObject *it;
3745 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003746
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003747 if (type == &filterfalse_type &&
3748 !_PyArg_NoKeywords("filterfalse()", kwds))
3749 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003750
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003751 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3752 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003753
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003754 /* Get iterator. */
3755 it = PyObject_GetIter(seq);
3756 if (it == NULL)
3757 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003758
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003759 /* create filterfalseobject structure */
3760 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3761 if (lz == NULL) {
3762 Py_DECREF(it);
3763 return NULL;
3764 }
3765 Py_INCREF(func);
3766 lz->func = func;
3767 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003768
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003769 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003770}
3771
3772static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003773filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003774{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003775 PyObject_GC_UnTrack(lz);
3776 Py_XDECREF(lz->func);
3777 Py_XDECREF(lz->it);
3778 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003779}
3780
3781static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003782filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003783{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003784 Py_VISIT(lz->it);
3785 Py_VISIT(lz->func);
3786 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003787}
3788
3789static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003790filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003791{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003792 PyObject *item;
3793 PyObject *it = lz->it;
3794 long ok;
3795 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003796
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003797 iternext = *Py_TYPE(it)->tp_iternext;
3798 for (;;) {
3799 item = iternext(it);
3800 if (item == NULL)
3801 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003802
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003803 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3804 ok = PyObject_IsTrue(item);
3805 } else {
3806 PyObject *good;
Raymond Hettingera6a2d442015-08-16 14:38:07 -07003807 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003808 if (good == NULL) {
3809 Py_DECREF(item);
3810 return NULL;
3811 }
3812 ok = PyObject_IsTrue(good);
3813 Py_DECREF(good);
3814 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003815 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003816 return item;
3817 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003818 if (ok < 0)
3819 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003820 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003821}
3822
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003823static PyObject *
3824filterfalse_reduce(filterfalseobject *lz)
3825{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003826 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3827}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003828
3829static PyMethodDef filterfalse_methods[] = {
3830 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3831 reduce_doc},
3832 {NULL, NULL} /* sentinel */
3833};
3834
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003835PyDoc_STRVAR(filterfalse_doc,
3836"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003837\n\
3838Return those items of sequence for which function(item) is false.\n\
3839If function is None, return the items that are false.");
3840
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003841static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003842 PyVarObject_HEAD_INIT(NULL, 0)
3843 "itertools.filterfalse", /* tp_name */
3844 sizeof(filterfalseobject), /* tp_basicsize */
3845 0, /* tp_itemsize */
3846 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003847 (destructor)filterfalse_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003848 0, /* tp_print */
3849 0, /* tp_getattr */
3850 0, /* tp_setattr */
3851 0, /* tp_reserved */
3852 0, /* tp_repr */
3853 0, /* tp_as_number */
3854 0, /* tp_as_sequence */
3855 0, /* tp_as_mapping */
3856 0, /* tp_hash */
3857 0, /* tp_call */
3858 0, /* tp_str */
3859 PyObject_GenericGetAttr, /* tp_getattro */
3860 0, /* tp_setattro */
3861 0, /* tp_as_buffer */
3862 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3863 Py_TPFLAGS_BASETYPE, /* tp_flags */
3864 filterfalse_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003865 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003866 0, /* tp_clear */
3867 0, /* tp_richcompare */
3868 0, /* tp_weaklistoffset */
3869 PyObject_SelfIter, /* tp_iter */
3870 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003871 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003872 0, /* tp_members */
3873 0, /* tp_getset */
3874 0, /* tp_base */
3875 0, /* tp_dict */
3876 0, /* tp_descr_get */
3877 0, /* tp_descr_set */
3878 0, /* tp_dictoffset */
3879 0, /* tp_init */
3880 0, /* tp_alloc */
3881 filterfalse_new, /* tp_new */
3882 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003883};
3884
3885
3886/* count object ************************************************************/
3887
3888typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003889 PyObject_HEAD
3890 Py_ssize_t cnt;
3891 PyObject *long_cnt;
3892 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003893} countobject;
3894
Raymond Hettinger30729212009-02-12 06:28:27 +00003895/* Counting logic and invariants:
3896
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003897fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003898
Serhiy Storchaka483405b2015-02-17 10:14:30 +02003899 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003900 Advances with: cnt += 1
3901 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003902
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003903slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003904
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003905 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3906 All counting is done with python objects (no overflows or underflows).
3907 Advances with: long_cnt += long_step
3908 Step may be zero -- effectively a slow version of repeat(cnt).
3909 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003910*/
3911
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003912static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003913
3914static PyObject *
3915count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3916{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003917 countobject *lz;
3918 int slow_mode = 0;
3919 Py_ssize_t cnt = 0;
3920 PyObject *long_cnt = NULL;
3921 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003922 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003923 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003924
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003925 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3926 kwlist, &long_cnt, &long_step))
3927 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003928
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003929 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3930 (long_step != NULL && !PyNumber_Check(long_step))) {
3931 PyErr_SetString(PyExc_TypeError, "a number is required");
3932 return NULL;
3933 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003934
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003935 if (long_cnt != NULL) {
3936 cnt = PyLong_AsSsize_t(long_cnt);
3937 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3938 PyErr_Clear();
3939 slow_mode = 1;
3940 }
3941 Py_INCREF(long_cnt);
3942 } else {
3943 cnt = 0;
3944 long_cnt = PyLong_FromLong(0);
3945 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003946
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003947 /* If not specified, step defaults to 1 */
3948 if (long_step == NULL) {
3949 long_step = PyLong_FromLong(1);
3950 if (long_step == NULL) {
3951 Py_DECREF(long_cnt);
3952 return NULL;
3953 }
3954 } else
3955 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003956
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003957 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003958
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003959 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003960 step = PyLong_AsLong(long_step);
3961 if (step != 1) {
3962 slow_mode = 1;
3963 if (step == -1 && PyErr_Occurred())
3964 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003965 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003966
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003967 if (slow_mode)
3968 cnt = PY_SSIZE_T_MAX;
3969 else
3970 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003971
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003972 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3973 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3974 assert(slow_mode ||
3975 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003976
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003977 /* create countobject structure */
3978 lz = (countobject *)type->tp_alloc(type, 0);
3979 if (lz == NULL) {
3980 Py_XDECREF(long_cnt);
3981 return NULL;
3982 }
3983 lz->cnt = cnt;
3984 lz->long_cnt = long_cnt;
3985 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003986
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003987 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003988}
3989
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003990static void
3991count_dealloc(countobject *lz)
3992{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003993 PyObject_GC_UnTrack(lz);
3994 Py_XDECREF(lz->long_cnt);
3995 Py_XDECREF(lz->long_step);
3996 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003997}
3998
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003999static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004000count_traverse(countobject *lz, visitproc visit, void *arg)
4001{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004002 Py_VISIT(lz->long_cnt);
4003 Py_VISIT(lz->long_step);
4004 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004005}
4006
4007static PyObject *
4008count_nextlong(countobject *lz)
4009{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004010 PyObject *long_cnt;
4011 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004012
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004013 long_cnt = lz->long_cnt;
4014 if (long_cnt == NULL) {
4015 /* Switch to slow_mode */
4016 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4017 if (long_cnt == NULL)
4018 return NULL;
4019 }
4020 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004021
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004022 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4023 if (stepped_up == NULL)
4024 return NULL;
4025 lz->long_cnt = stepped_up;
4026 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004027}
4028
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004029static PyObject *
4030count_next(countobject *lz)
4031{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004032 if (lz->cnt == PY_SSIZE_T_MAX)
4033 return count_nextlong(lz);
4034 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004035}
4036
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004037static PyObject *
4038count_repr(countobject *lz)
4039{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004040 if (lz->cnt != PY_SSIZE_T_MAX)
4041 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004042
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004043 if (PyLong_Check(lz->long_step)) {
4044 long step = PyLong_AsLong(lz->long_step);
4045 if (step == -1 && PyErr_Occurred()) {
4046 PyErr_Clear();
4047 }
4048 if (step == 1) {
4049 /* Don't display step when it is an integer equal to 1 */
4050 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
4051 }
4052 }
4053 return PyUnicode_FromFormat("count(%R, %R)",
4054 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004055}
4056
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004057static PyObject *
4058count_reduce(countobject *lz)
4059{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004060 if (lz->cnt == PY_SSIZE_T_MAX)
4061 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4062 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004063}
4064
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004065static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004066 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004067 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004068 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004069};
4070
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004071PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00004072 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004073\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004074Return a count object whose .__next__() method returns consecutive values.\n\
4075Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00004076 def count(firstval=0, step=1):\n\
Eli Benderskyc554f722013-09-03 06:37:19 -07004077 x = firstval\n\
4078 while 1:\n\
4079 yield x\n\
4080 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004081
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004082static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004083 PyVarObject_HEAD_INIT(NULL, 0)
4084 "itertools.count", /* tp_name */
4085 sizeof(countobject), /* tp_basicsize */
4086 0, /* tp_itemsize */
4087 /* methods */
4088 (destructor)count_dealloc, /* tp_dealloc */
4089 0, /* tp_print */
4090 0, /* tp_getattr */
4091 0, /* tp_setattr */
4092 0, /* tp_reserved */
4093 (reprfunc)count_repr, /* tp_repr */
4094 0, /* tp_as_number */
4095 0, /* tp_as_sequence */
4096 0, /* tp_as_mapping */
4097 0, /* tp_hash */
4098 0, /* tp_call */
4099 0, /* tp_str */
4100 PyObject_GenericGetAttr, /* tp_getattro */
4101 0, /* tp_setattro */
4102 0, /* tp_as_buffer */
4103 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004104 Py_TPFLAGS_BASETYPE, /* tp_flags */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004105 count_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004106 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004107 0, /* tp_clear */
4108 0, /* tp_richcompare */
4109 0, /* tp_weaklistoffset */
4110 PyObject_SelfIter, /* tp_iter */
4111 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004112 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004113 0, /* tp_members */
4114 0, /* tp_getset */
4115 0, /* tp_base */
4116 0, /* tp_dict */
4117 0, /* tp_descr_get */
4118 0, /* tp_descr_set */
4119 0, /* tp_dictoffset */
4120 0, /* tp_init */
4121 0, /* tp_alloc */
4122 count_new, /* tp_new */
4123 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004124};
4125
4126
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004127/* repeat object ************************************************************/
4128
4129typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004130 PyObject_HEAD
4131 PyObject *element;
4132 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004133} repeatobject;
4134
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004135static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004136
4137static PyObject *
4138repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4139{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004140 repeatobject *ro;
4141 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004142 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004143 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004144
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004145 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4146 &element, &cnt))
4147 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004148
Raymond Hettinger97d35552014-06-24 21:36:58 -07004149 if (kwds != NULL)
4150 n_kwds = PyDict_Size(kwds);
4151 /* Does user supply times argument? */
4152 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004153 cnt = 0;
4154
4155 ro = (repeatobject *)type->tp_alloc(type, 0);
4156 if (ro == NULL)
4157 return NULL;
4158 Py_INCREF(element);
4159 ro->element = element;
4160 ro->cnt = cnt;
4161 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004162}
4163
4164static void
4165repeat_dealloc(repeatobject *ro)
4166{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004167 PyObject_GC_UnTrack(ro);
4168 Py_XDECREF(ro->element);
4169 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004170}
4171
4172static int
4173repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4174{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004175 Py_VISIT(ro->element);
4176 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004177}
4178
4179static PyObject *
4180repeat_next(repeatobject *ro)
4181{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004182 if (ro->cnt == 0)
4183 return NULL;
4184 if (ro->cnt > 0)
4185 ro->cnt--;
4186 Py_INCREF(ro->element);
4187 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004188}
4189
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004190static PyObject *
4191repeat_repr(repeatobject *ro)
4192{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004193 if (ro->cnt == -1)
4194 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4195 else
4196 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4197}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004198
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004199static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004200repeat_len(repeatobject *ro)
4201{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004202 if (ro->cnt == -1) {
4203 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4204 return NULL;
4205 }
4206 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004207}
4208
Armin Rigof5b3e362006-02-11 21:32:43 +00004209PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004210
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004211static PyObject *
4212repeat_reduce(repeatobject *ro)
4213{
4214 /* unpickle this so that a new repeat iterator is constructed with an
4215 * object, then call __setstate__ on it to set cnt
4216 */
4217 if (ro->cnt >= 0)
4218 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4219 else
4220 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4221}
4222
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004223static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004224 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004225 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004226 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004227};
4228
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004229PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004230"repeat(object [,times]) -> create an iterator which returns the object\n\
4231for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004232endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004233
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004234static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004235 PyVarObject_HEAD_INIT(NULL, 0)
4236 "itertools.repeat", /* tp_name */
4237 sizeof(repeatobject), /* tp_basicsize */
4238 0, /* tp_itemsize */
4239 /* methods */
4240 (destructor)repeat_dealloc, /* tp_dealloc */
4241 0, /* tp_print */
4242 0, /* tp_getattr */
4243 0, /* tp_setattr */
4244 0, /* tp_reserved */
4245 (reprfunc)repeat_repr, /* tp_repr */
4246 0, /* tp_as_number */
4247 0, /* tp_as_sequence */
4248 0, /* tp_as_mapping */
4249 0, /* tp_hash */
4250 0, /* tp_call */
4251 0, /* tp_str */
4252 PyObject_GenericGetAttr, /* tp_getattro */
4253 0, /* tp_setattro */
4254 0, /* tp_as_buffer */
4255 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4256 Py_TPFLAGS_BASETYPE, /* tp_flags */
4257 repeat_doc, /* tp_doc */
4258 (traverseproc)repeat_traverse, /* tp_traverse */
4259 0, /* tp_clear */
4260 0, /* tp_richcompare */
4261 0, /* tp_weaklistoffset */
4262 PyObject_SelfIter, /* tp_iter */
4263 (iternextfunc)repeat_next, /* tp_iternext */
4264 repeat_methods, /* tp_methods */
4265 0, /* tp_members */
4266 0, /* tp_getset */
4267 0, /* tp_base */
4268 0, /* tp_dict */
4269 0, /* tp_descr_get */
4270 0, /* tp_descr_set */
4271 0, /* tp_dictoffset */
4272 0, /* tp_init */
4273 0, /* tp_alloc */
4274 repeat_new, /* tp_new */
4275 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004276};
4277
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004278/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004279
4280typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004281 PyObject_HEAD
4282 Py_ssize_t tuplesize;
4283 Py_ssize_t numactive;
4284 PyObject *ittuple; /* tuple of iterators */
4285 PyObject *result;
4286 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004287} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004288
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004289static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004290
4291static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004292zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004293{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004294 ziplongestobject *lz;
4295 Py_ssize_t i;
4296 PyObject *ittuple; /* tuple of iterators */
4297 PyObject *result;
4298 PyObject *fillvalue = Py_None;
4299 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004300
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004301 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
4302 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4303 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
4304 PyErr_SetString(PyExc_TypeError,
4305 "zip_longest() got an unexpected keyword argument");
4306 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004307 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004308 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004309
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004310 /* args must be a tuple */
4311 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00004312
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004313 /* obtain iterators */
4314 ittuple = PyTuple_New(tuplesize);
4315 if (ittuple == NULL)
4316 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004317 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004318 PyObject *item = PyTuple_GET_ITEM(args, i);
4319 PyObject *it = PyObject_GetIter(item);
4320 if (it == NULL) {
4321 if (PyErr_ExceptionMatches(PyExc_TypeError))
4322 PyErr_Format(PyExc_TypeError,
4323 "zip_longest argument #%zd must support iteration",
4324 i+1);
4325 Py_DECREF(ittuple);
4326 return NULL;
4327 }
4328 PyTuple_SET_ITEM(ittuple, i, it);
4329 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004330
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004331 /* create a result holder */
4332 result = PyTuple_New(tuplesize);
4333 if (result == NULL) {
4334 Py_DECREF(ittuple);
4335 return NULL;
4336 }
4337 for (i=0 ; i < tuplesize ; i++) {
4338 Py_INCREF(Py_None);
4339 PyTuple_SET_ITEM(result, i, Py_None);
4340 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004341
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004342 /* create ziplongestobject structure */
4343 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4344 if (lz == NULL) {
4345 Py_DECREF(ittuple);
4346 Py_DECREF(result);
4347 return NULL;
4348 }
4349 lz->ittuple = ittuple;
4350 lz->tuplesize = tuplesize;
4351 lz->numactive = tuplesize;
4352 lz->result = result;
4353 Py_INCREF(fillvalue);
4354 lz->fillvalue = fillvalue;
4355 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004356}
4357
4358static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004359zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004360{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004361 PyObject_GC_UnTrack(lz);
4362 Py_XDECREF(lz->ittuple);
4363 Py_XDECREF(lz->result);
4364 Py_XDECREF(lz->fillvalue);
4365 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004366}
4367
4368static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004369zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004370{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004371 Py_VISIT(lz->ittuple);
4372 Py_VISIT(lz->result);
4373 Py_VISIT(lz->fillvalue);
4374 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004375}
4376
4377static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004378zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004379{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004380 Py_ssize_t i;
4381 Py_ssize_t tuplesize = lz->tuplesize;
4382 PyObject *result = lz->result;
4383 PyObject *it;
4384 PyObject *item;
4385 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004386
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004387 if (tuplesize == 0)
4388 return NULL;
4389 if (lz->numactive == 0)
4390 return NULL;
4391 if (Py_REFCNT(result) == 1) {
4392 Py_INCREF(result);
4393 for (i=0 ; i < tuplesize ; i++) {
4394 it = PyTuple_GET_ITEM(lz->ittuple, i);
4395 if (it == NULL) {
4396 Py_INCREF(lz->fillvalue);
4397 item = lz->fillvalue;
4398 } else {
4399 item = PyIter_Next(it);
4400 if (item == NULL) {
4401 lz->numactive -= 1;
4402 if (lz->numactive == 0 || PyErr_Occurred()) {
4403 lz->numactive = 0;
4404 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004405 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004406 } else {
4407 Py_INCREF(lz->fillvalue);
4408 item = lz->fillvalue;
4409 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4410 Py_DECREF(it);
4411 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004412 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004413 }
4414 olditem = PyTuple_GET_ITEM(result, i);
4415 PyTuple_SET_ITEM(result, i, item);
4416 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004417 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004418 } else {
4419 result = PyTuple_New(tuplesize);
4420 if (result == NULL)
4421 return NULL;
4422 for (i=0 ; i < tuplesize ; i++) {
4423 it = PyTuple_GET_ITEM(lz->ittuple, i);
4424 if (it == NULL) {
4425 Py_INCREF(lz->fillvalue);
4426 item = lz->fillvalue;
4427 } else {
4428 item = PyIter_Next(it);
4429 if (item == NULL) {
4430 lz->numactive -= 1;
4431 if (lz->numactive == 0 || PyErr_Occurred()) {
4432 lz->numactive = 0;
4433 Py_DECREF(result);
4434 return NULL;
4435 } else {
4436 Py_INCREF(lz->fillvalue);
4437 item = lz->fillvalue;
4438 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4439 Py_DECREF(it);
4440 }
4441 }
4442 }
4443 PyTuple_SET_ITEM(result, i, item);
4444 }
4445 }
4446 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004447}
4448
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004449static PyObject *
4450zip_longest_reduce(ziplongestobject *lz)
4451{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004452
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004453 /* Create a new tuple with empty sequences where appropriate to pickle.
4454 * Then use setstate to set the fillvalue
4455 */
4456 int i;
4457 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004458
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004459 if (args == NULL)
4460 return NULL;
4461 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4462 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4463 if (elem == NULL) {
4464 elem = PyTuple_New(0);
4465 if (elem == NULL) {
4466 Py_DECREF(args);
4467 return NULL;
4468 }
4469 } else
4470 Py_INCREF(elem);
4471 PyTuple_SET_ITEM(args, i, elem);
4472 }
4473 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4474}
4475
4476static PyObject *
4477zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4478{
4479 Py_CLEAR(lz->fillvalue);
4480 lz->fillvalue = state;
4481 Py_INCREF(lz->fillvalue);
4482 Py_RETURN_NONE;
4483}
4484
4485static PyMethodDef zip_longest_methods[] = {
4486 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4487 reduce_doc},
4488 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4489 setstate_doc},
4490 {NULL, NULL} /* sentinel */
4491};
4492
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004493PyDoc_STRVAR(zip_longest_doc,
4494"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004495\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004496Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004497the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004498method continues until the longest iterable in the argument sequence\n\
4499is exhausted and then it raises StopIteration. When the shorter iterables\n\
4500are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4501defaults to None or can be specified by a keyword argument.\n\
4502");
4503
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004504static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004505 PyVarObject_HEAD_INIT(NULL, 0)
4506 "itertools.zip_longest", /* tp_name */
4507 sizeof(ziplongestobject), /* tp_basicsize */
4508 0, /* tp_itemsize */
4509 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004510 (destructor)zip_longest_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004511 0, /* tp_print */
4512 0, /* tp_getattr */
4513 0, /* tp_setattr */
4514 0, /* tp_reserved */
4515 0, /* tp_repr */
4516 0, /* tp_as_number */
4517 0, /* tp_as_sequence */
4518 0, /* tp_as_mapping */
4519 0, /* tp_hash */
4520 0, /* tp_call */
4521 0, /* tp_str */
4522 PyObject_GenericGetAttr, /* tp_getattro */
4523 0, /* tp_setattro */
4524 0, /* tp_as_buffer */
4525 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4526 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004527 zip_longest_doc, /* tp_doc */
4528 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004529 0, /* tp_clear */
4530 0, /* tp_richcompare */
4531 0, /* tp_weaklistoffset */
4532 PyObject_SelfIter, /* tp_iter */
4533 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004534 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004535 0, /* tp_members */
4536 0, /* tp_getset */
4537 0, /* tp_base */
4538 0, /* tp_dict */
4539 0, /* tp_descr_get */
4540 0, /* tp_descr_set */
4541 0, /* tp_dictoffset */
4542 0, /* tp_init */
4543 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004544 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004545 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004546};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004547
4548/* module level code ********************************************************/
4549
4550PyDoc_STRVAR(module_doc,
4551"Functional tools for creating and using iterators.\n\
4552\n\
4553Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004554count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004555cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004556repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004557\n\
4558Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004559accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004560chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger8df58f72013-09-09 01:29:40 -05004561chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004562compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4563dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4564groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004565filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004566islice(seq, [start,] stop [, step]) --> elements from\n\
4567 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004568starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004569tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004570takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004571zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004572\n\
4573Combinatoric generators:\n\
4574product(p, q, ... [repeat=1]) --> cartesian product\n\
4575permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004576combinations(p, r)\n\
4577combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004578");
4579
4580
Raymond Hettingerad983e72003-11-12 14:32:26 +00004581static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004582 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4583 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004584};
4585
Martin v. Löwis1a214512008-06-11 05:26:20 +00004586
4587static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004588 PyModuleDef_HEAD_INIT,
4589 "itertools",
4590 module_doc,
4591 -1,
4592 module_methods,
4593 NULL,
4594 NULL,
4595 NULL,
4596 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004597};
4598
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004599PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004600PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004601{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004602 int i;
4603 PyObject *m;
4604 char *name;
4605 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004606 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004607 &combinations_type,
4608 &cwr_type,
4609 &cycle_type,
4610 &dropwhile_type,
4611 &takewhile_type,
4612 &islice_type,
4613 &starmap_type,
4614 &chain_type,
4615 &compress_type,
4616 &filterfalse_type,
4617 &count_type,
4618 &ziplongest_type,
4619 &permutations_type,
4620 &product_type,
4621 &repeat_type,
4622 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004623 &_grouper_type,
4624 &tee_type,
4625 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004626 NULL
4627 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004628
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004629 Py_TYPE(&teedataobject_type) = &PyType_Type;
4630 m = PyModule_Create(&itertoolsmodule);
4631 if (m == NULL)
4632 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004633
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004634 for (i=0 ; typelist[i] != NULL ; i++) {
4635 if (PyType_Ready(typelist[i]) < 0)
4636 return NULL;
4637 name = strchr(typelist[i]->tp_name, '.');
4638 assert (name != NULL);
4639 Py_INCREF(typelist[i]);
4640 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4641 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004642
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004643 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004644}