blob: 103029d251e0dae7bfd28863043d50299d6fec32 [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"
Sergey Fedoseev234531b2019-02-25 21:59:12 +05004#include "pycore_tupleobject.h"
Fred Drake08ebfec2004-10-17 19:36:57 +00005#include "structmember.h"
Raymond Hettinger96ef8112003-02-01 00:10:11 +00006
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00007/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00008 by Raymond D. Hettinger <python@rcn.com>
Raymond Hettinger96ef8112003-02-01 00:10:11 +00009*/
10
Tal Einat3286ce42018-09-10 21:33:08 +030011/*[clinic input]
12module itertools
13class itertools.groupby "groupbyobject *" "&groupby_type"
14class itertools._grouper "_grouperobject *" "&_grouper_type"
Tal Einatc4bccd32018-09-12 00:49:13 +030015class itertools.teedataobject "teedataobject *" "&teedataobject_type"
16class itertools._tee "teeobject *" "&tee_type"
17class itertools.cycle "cycleobject *" "&cycle_type"
18class itertools.dropwhile "dropwhileobject *" "&dropwhile_type"
19class itertools.takewhile "takewhileobject *" "&takewhile_type"
20class itertools.starmap "starmapobject *" "&starmap_type"
21class itertools.chain "chainobject *" "&chain_type"
22class itertools.combinations "combinationsobject *" "&combinations_type"
23class itertools.combinations_with_replacement "cwr_object *" "&cwr_type"
24class itertools.permutations "permutationsobject *" "&permutations_type"
25class itertools.accumulate "accumulateobject *" "&accumulate_type"
26class itertools.compress "compressobject *" "&compress_type"
27class itertools.filterfalse "filterfalseobject *" "&filterfalse_type"
28class itertools.count "countobject *" "&count_type"
Tal Einat3286ce42018-09-10 21:33:08 +030029[clinic start generated code]*/
Tal Einatc4bccd32018-09-12 00:49:13 +030030/*[clinic end generated code: output=da39a3ee5e6b4b0d input=ea05c93c6d94726a]*/
Tal Einat3286ce42018-09-10 21:33:08 +030031
32static PyTypeObject groupby_type;
33static PyTypeObject _grouper_type;
Tal Einatc4bccd32018-09-12 00:49:13 +030034static PyTypeObject teedataobject_type;
35static PyTypeObject tee_type;
36static PyTypeObject cycle_type;
37static PyTypeObject dropwhile_type;
38static PyTypeObject takewhile_type;
39static PyTypeObject starmap_type;
40static PyTypeObject combinations_type;
41static PyTypeObject cwr_type;
42static PyTypeObject permutations_type;
43static PyTypeObject accumulate_type;
44static PyTypeObject compress_type;
45static PyTypeObject filterfalse_type;
46static PyTypeObject count_type;
47
Tal Einat3286ce42018-09-10 21:33:08 +030048#include "clinic/itertoolsmodule.c.h"
49
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000050
Raymond Hettingera6ea44a2015-08-17 23:55:28 -070051/* groupby object ************************************************************/
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000052
53typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000054 PyObject_HEAD
55 PyObject *it;
56 PyObject *keyfunc;
57 PyObject *tgtkey;
58 PyObject *currkey;
59 PyObject *currvalue;
Serhiy Storchakac247caf2017-09-24 13:36:11 +030060 const void *currgrouper; /* borrowed reference */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000061} groupbyobject;
62
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000063static PyObject *_grouper_create(groupbyobject *, PyObject *);
64
Tal Einat3286ce42018-09-10 21:33:08 +030065/*[clinic input]
66@classmethod
67itertools.groupby.__new__
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000068
Tal Einat3286ce42018-09-10 21:33:08 +030069 iterable as it: object
70 Elements to divide into groups according to the key function.
71 key as keyfunc: object = None
72 A function for computing the group category for each element.
73 If the key function is not specified or is None, the element itself
74 is used for grouping.
75
76make an iterator that returns consecutive keys and groups from the iterable
77[clinic start generated code]*/
78
79static PyObject *
80itertools_groupby_impl(PyTypeObject *type, PyObject *it, PyObject *keyfunc)
81/*[clinic end generated code: output=cbb1ae3a90fd4141 input=6b3d123e87ff65a1]*/
82{
83 groupbyobject *gbo;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000084
85 gbo = (groupbyobject *)type->tp_alloc(type, 0);
86 if (gbo == NULL)
87 return NULL;
88 gbo->tgtkey = NULL;
89 gbo->currkey = NULL;
90 gbo->currvalue = NULL;
91 gbo->keyfunc = keyfunc;
92 Py_INCREF(keyfunc);
93 gbo->it = PyObject_GetIter(it);
94 if (gbo->it == NULL) {
95 Py_DECREF(gbo);
96 return NULL;
97 }
98 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000099}
100
101static void
102groupby_dealloc(groupbyobject *gbo)
103{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000104 PyObject_GC_UnTrack(gbo);
105 Py_XDECREF(gbo->it);
106 Py_XDECREF(gbo->keyfunc);
107 Py_XDECREF(gbo->tgtkey);
108 Py_XDECREF(gbo->currkey);
109 Py_XDECREF(gbo->currvalue);
110 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000111}
112
113static int
114groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
115{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000116 Py_VISIT(gbo->it);
117 Py_VISIT(gbo->keyfunc);
118 Py_VISIT(gbo->tgtkey);
119 Py_VISIT(gbo->currkey);
120 Py_VISIT(gbo->currvalue);
121 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122}
123
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300124Py_LOCAL_INLINE(int)
125groupby_step(groupbyobject *gbo)
126{
127 PyObject *newvalue, *newkey, *oldvalue;
128
129 newvalue = PyIter_Next(gbo->it);
130 if (newvalue == NULL)
131 return -1;
132
133 if (gbo->keyfunc == Py_None) {
134 newkey = newvalue;
135 Py_INCREF(newvalue);
136 } else {
137 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL);
138 if (newkey == NULL) {
139 Py_DECREF(newvalue);
140 return -1;
141 }
142 }
143
144 oldvalue = gbo->currvalue;
145 gbo->currvalue = newvalue;
146 Py_XSETREF(gbo->currkey, newkey);
147 Py_XDECREF(oldvalue);
148 return 0;
149}
150
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000151static PyObject *
152groupby_next(groupbyobject *gbo)
153{
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300154 PyObject *r, *grouper;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000155
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300156 gbo->currgrouper = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000157 /* skip to next iteration group */
158 for (;;) {
159 if (gbo->currkey == NULL)
160 /* pass */;
161 else if (gbo->tgtkey == NULL)
162 break;
163 else {
164 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000165
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700166 rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000167 if (rcmp == -1)
168 return NULL;
169 else if (rcmp == 0)
170 break;
171 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000172
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300173 if (groupby_step(gbo) < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000174 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000175 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000176 Py_INCREF(gbo->currkey);
Serhiy Storchakaec397562016-04-06 09:50:03 +0300177 Py_XSETREF(gbo->tgtkey, gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000178
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000179 grouper = _grouper_create(gbo, gbo->tgtkey);
180 if (grouper == NULL)
181 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000182
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000183 r = PyTuple_Pack(2, gbo->currkey, grouper);
184 Py_DECREF(grouper);
185 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000186}
187
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000188static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530189groupby_reduce(groupbyobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000190{
191 /* reduce as a 'new' call with an optional 'setstate' if groupby
192 * has started
193 */
194 PyObject *value;
195 if (lz->tgtkey && lz->currkey && lz->currvalue)
196 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
197 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
198 else
199 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
200 lz->it, lz->keyfunc);
201
202 return value;
203}
204
205PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
206
207static PyObject *
208groupby_setstate(groupbyobject *lz, PyObject *state)
209{
210 PyObject *currkey, *currvalue, *tgtkey;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300211 if (!PyTuple_Check(state)) {
212 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000213 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300214 }
215 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey)) {
216 return NULL;
217 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200218 Py_INCREF(currkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300219 Py_XSETREF(lz->currkey, currkey);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200220 Py_INCREF(currvalue);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300221 Py_XSETREF(lz->currvalue, currvalue);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200222 Py_INCREF(tgtkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300223 Py_XSETREF(lz->tgtkey, tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000224 Py_RETURN_NONE;
225}
226
227PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
228
229static PyMethodDef groupby_methods[] = {
230 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
231 reduce_doc},
232 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
233 setstate_doc},
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700234 {NULL, NULL} /* sentinel */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000235};
236
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000237static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000238 PyVarObject_HEAD_INIT(NULL, 0)
239 "itertools.groupby", /* tp_name */
240 sizeof(groupbyobject), /* tp_basicsize */
241 0, /* tp_itemsize */
242 /* methods */
243 (destructor)groupby_dealloc, /* tp_dealloc */
244 0, /* tp_print */
245 0, /* tp_getattr */
246 0, /* tp_setattr */
247 0, /* tp_reserved */
248 0, /* tp_repr */
249 0, /* tp_as_number */
250 0, /* tp_as_sequence */
251 0, /* tp_as_mapping */
252 0, /* tp_hash */
253 0, /* tp_call */
254 0, /* tp_str */
255 PyObject_GenericGetAttr, /* tp_getattro */
256 0, /* tp_setattro */
257 0, /* tp_as_buffer */
258 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
259 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einat3286ce42018-09-10 21:33:08 +0300260 itertools_groupby__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000261 (traverseproc)groupby_traverse, /* tp_traverse */
262 0, /* tp_clear */
263 0, /* tp_richcompare */
264 0, /* tp_weaklistoffset */
265 PyObject_SelfIter, /* tp_iter */
266 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000267 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000268 0, /* tp_members */
269 0, /* tp_getset */
270 0, /* tp_base */
271 0, /* tp_dict */
272 0, /* tp_descr_get */
273 0, /* tp_descr_set */
274 0, /* tp_dictoffset */
275 0, /* tp_init */
276 0, /* tp_alloc */
Tal Einat3286ce42018-09-10 21:33:08 +0300277 itertools_groupby, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000278 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000279};
280
281
282/* _grouper object (internal) ************************************************/
283
284typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000285 PyObject_HEAD
286 PyObject *parent;
287 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000288} _grouperobject;
289
Tal Einat3286ce42018-09-10 21:33:08 +0300290/*[clinic input]
291@classmethod
292itertools._grouper.__new__
293
294 parent: object(subclass_of='&groupby_type')
295 tgtkey: object
296 /
297[clinic start generated code]*/
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000298
299static PyObject *
Tal Einat3286ce42018-09-10 21:33:08 +0300300itertools__grouper_impl(PyTypeObject *type, PyObject *parent,
301 PyObject *tgtkey)
302/*[clinic end generated code: output=462efb1cdebb5914 input=dc180d7771fc8c59]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000303{
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000304 return _grouper_create((groupbyobject*) parent, tgtkey);
305}
306
307static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000308_grouper_create(groupbyobject *parent, PyObject *tgtkey)
309{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000310 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000311
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000312 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
313 if (igo == NULL)
314 return NULL;
315 igo->parent = (PyObject *)parent;
316 Py_INCREF(parent);
317 igo->tgtkey = tgtkey;
318 Py_INCREF(tgtkey);
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300319 parent->currgrouper = igo; /* borrowed reference */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000320
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000321 PyObject_GC_Track(igo);
322 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000323}
324
325static void
326_grouper_dealloc(_grouperobject *igo)
327{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000328 PyObject_GC_UnTrack(igo);
329 Py_DECREF(igo->parent);
330 Py_DECREF(igo->tgtkey);
331 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000332}
333
334static int
335_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
336{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000337 Py_VISIT(igo->parent);
338 Py_VISIT(igo->tgtkey);
339 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000340}
341
342static PyObject *
343_grouper_next(_grouperobject *igo)
344{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000345 groupbyobject *gbo = (groupbyobject *)igo->parent;
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300346 PyObject *r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000347 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000348
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300349 if (gbo->currgrouper != igo)
350 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000351 if (gbo->currvalue == NULL) {
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300352 if (groupby_step(gbo) < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000353 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000354 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000355
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000356 assert(gbo->currkey != NULL);
357 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
358 if (rcmp <= 0)
359 /* got any error or current group is end */
360 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000361
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000362 r = gbo->currvalue;
363 gbo->currvalue = NULL;
364 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000365
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000366 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000367}
368
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000369static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530370_grouper_reduce(_grouperobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000371{
Serhiy Storchakabb86bf42018-12-11 08:28:18 +0200372 _Py_IDENTIFIER(iter);
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300373 if (((groupbyobject *)lz->parent)->currgrouper != lz) {
Serhiy Storchakabb86bf42018-12-11 08:28:18 +0200374 return Py_BuildValue("N(())", _PyEval_GetBuiltinId(&PyId_iter));
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300375 }
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700376 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000377}
378
379static PyMethodDef _grouper_methods[] = {
380 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
381 reduce_doc},
382 {NULL, NULL} /* sentinel */
383};
384
385
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000386static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000387 PyVarObject_HEAD_INIT(NULL, 0)
388 "itertools._grouper", /* tp_name */
389 sizeof(_grouperobject), /* tp_basicsize */
390 0, /* tp_itemsize */
391 /* methods */
392 (destructor)_grouper_dealloc, /* tp_dealloc */
393 0, /* tp_print */
394 0, /* tp_getattr */
395 0, /* tp_setattr */
396 0, /* tp_reserved */
397 0, /* tp_repr */
398 0, /* tp_as_number */
399 0, /* tp_as_sequence */
400 0, /* tp_as_mapping */
401 0, /* tp_hash */
402 0, /* tp_call */
403 0, /* tp_str */
404 PyObject_GenericGetAttr, /* tp_getattro */
405 0, /* tp_setattro */
406 0, /* tp_as_buffer */
407 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
408 0, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700409 (traverseproc)_grouper_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000410 0, /* tp_clear */
411 0, /* tp_richcompare */
412 0, /* tp_weaklistoffset */
413 PyObject_SelfIter, /* tp_iter */
414 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000415 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000416 0, /* tp_members */
417 0, /* tp_getset */
418 0, /* tp_base */
419 0, /* tp_dict */
420 0, /* tp_descr_get */
421 0, /* tp_descr_set */
422 0, /* tp_dictoffset */
423 0, /* tp_init */
424 0, /* tp_alloc */
Tal Einat3286ce42018-09-10 21:33:08 +0300425 itertools__grouper, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000426 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000427};
428
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000429
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700430/* tee object and with supporting function and objects ***********************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000431
Raymond Hettingerad983e72003-11-12 14:32:26 +0000432/* The teedataobject pre-allocates space for LINKCELLS number of objects.
433 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000434 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000435 the other structure members including PyHEAD overhead. The larger the
436 value, the less memory overhead per object and the less time spent
437 allocating/deallocating new links. The smaller the number, the less
438 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000439*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000440#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000441
442typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000443 PyObject_HEAD
444 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700445 int numread; /* 0 <= numread <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000446 PyObject *nextlink;
447 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000448} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000449
450typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000451 PyObject_HEAD
452 teedataobject *dataobj;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700453 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000454 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000455} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000456
Raymond Hettingerad983e72003-11-12 14:32:26 +0000457static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000458
459static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000460teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000461{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000462 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000463
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000464 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
465 if (tdo == NULL)
466 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000467
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000468 tdo->numread = 0;
469 tdo->nextlink = NULL;
470 Py_INCREF(it);
471 tdo->it = it;
472 PyObject_GC_Track(tdo);
473 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000474}
475
476static PyObject *
477teedataobject_jumplink(teedataobject *tdo)
478{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000479 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000480 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000481 Py_XINCREF(tdo->nextlink);
482 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000483}
484
485static PyObject *
486teedataobject_getitem(teedataobject *tdo, int i)
487{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000488 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000489
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000490 assert(i < LINKCELLS);
491 if (i < tdo->numread)
492 value = tdo->values[i];
493 else {
494 /* this is the lead iterator, so fetch more data */
495 assert(i == tdo->numread);
496 value = PyIter_Next(tdo->it);
497 if (value == NULL)
498 return NULL;
499 tdo->numread++;
500 tdo->values[i] = value;
501 }
502 Py_INCREF(value);
503 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000504}
505
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000506static int
507teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
508{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000509 int i;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700510
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000511 Py_VISIT(tdo->it);
512 for (i = 0; i < tdo->numread; i++)
513 Py_VISIT(tdo->values[i]);
514 Py_VISIT(tdo->nextlink);
515 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000516}
517
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200518static void
519teedataobject_safe_decref(PyObject *obj)
520{
521 while (obj && Py_TYPE(obj) == &teedataobject_type &&
522 Py_REFCNT(obj) == 1) {
523 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
524 ((teedataobject *)obj)->nextlink = NULL;
525 Py_DECREF(obj);
526 obj = nextlink;
527 }
528 Py_XDECREF(obj);
529}
530
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000531static int
532teedataobject_clear(teedataobject *tdo)
533{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000534 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200535 PyObject *tmp;
536
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000537 Py_CLEAR(tdo->it);
538 for (i=0 ; i<tdo->numread ; i++)
539 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200540 tmp = tdo->nextlink;
541 tdo->nextlink = NULL;
542 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000543 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000544}
545
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000546static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000547teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000548{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000549 PyObject_GC_UnTrack(tdo);
550 teedataobject_clear(tdo);
551 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000552}
553
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000554static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530555teedataobject_reduce(teedataobject *tdo, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000556{
557 int i;
558 /* create a temporary list of already iterated values */
559 PyObject *values = PyList_New(tdo->numread);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700560
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000561 if (!values)
562 return NULL;
563 for (i=0 ; i<tdo->numread ; i++) {
564 Py_INCREF(tdo->values[i]);
565 PyList_SET_ITEM(values, i, tdo->values[i]);
566 }
567 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
568 values,
569 tdo->nextlink ? tdo->nextlink : Py_None);
570}
571
Tal Einatc4bccd32018-09-12 00:49:13 +0300572/*[clinic input]
573@classmethod
574itertools.teedataobject.__new__
575 iterable as it: object
576 values: object(subclass_of='&PyList_Type')
577 next: object
578 /
579Data container common to multiple tee objects.
580[clinic start generated code]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000581
582static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300583itertools_teedataobject_impl(PyTypeObject *type, PyObject *it,
584 PyObject *values, PyObject *next)
585/*[clinic end generated code: output=3343ceb07e08df5e input=be60f2fabd2b72ba]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000586{
587 teedataobject *tdo;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000588 Py_ssize_t i, len;
589
590 assert(type == &teedataobject_type);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000591
592 tdo = (teedataobject *)teedataobject_newinternal(it);
593 if (!tdo)
594 return NULL;
595
596 len = PyList_GET_SIZE(values);
597 if (len > LINKCELLS)
598 goto err;
599 for (i=0; i<len; i++) {
600 tdo->values[i] = PyList_GET_ITEM(values, i);
601 Py_INCREF(tdo->values[i]);
602 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200603 /* len <= LINKCELLS < INT_MAX */
604 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000605
606 if (len == LINKCELLS) {
607 if (next != Py_None) {
608 if (Py_TYPE(next) != &teedataobject_type)
609 goto err;
610 assert(tdo->nextlink == NULL);
611 Py_INCREF(next);
612 tdo->nextlink = next;
613 }
614 } else {
615 if (next != Py_None)
616 goto err; /* shouldn't have a next if we are not full */
617 }
618 return (PyObject*)tdo;
619
620err:
621 Py_XDECREF(tdo);
622 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
623 return NULL;
624}
625
626static PyMethodDef teedataobject_methods[] = {
627 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
628 reduce_doc},
629 {NULL, NULL} /* sentinel */
630};
631
Raymond Hettingerad983e72003-11-12 14:32:26 +0000632static PyTypeObject teedataobject_type = {
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700633 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000634 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000635 sizeof(teedataobject), /* tp_basicsize */
636 0, /* tp_itemsize */
637 /* methods */
638 (destructor)teedataobject_dealloc, /* tp_dealloc */
639 0, /* tp_print */
640 0, /* tp_getattr */
641 0, /* tp_setattr */
642 0, /* tp_reserved */
643 0, /* tp_repr */
644 0, /* tp_as_number */
645 0, /* tp_as_sequence */
646 0, /* tp_as_mapping */
647 0, /* tp_hash */
648 0, /* tp_call */
649 0, /* tp_str */
650 PyObject_GenericGetAttr, /* tp_getattro */
651 0, /* tp_setattro */
652 0, /* tp_as_buffer */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700653 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300654 itertools_teedataobject__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000655 (traverseproc)teedataobject_traverse, /* tp_traverse */
656 (inquiry)teedataobject_clear, /* tp_clear */
657 0, /* tp_richcompare */
658 0, /* tp_weaklistoffset */
659 0, /* tp_iter */
660 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000661 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000662 0, /* tp_members */
663 0, /* tp_getset */
664 0, /* tp_base */
665 0, /* tp_dict */
666 0, /* tp_descr_get */
667 0, /* tp_descr_set */
668 0, /* tp_dictoffset */
669 0, /* tp_init */
670 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300671 itertools_teedataobject, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000672 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000673};
674
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000675
676static PyTypeObject tee_type;
677
678static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000679tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000680{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000681 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000682
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000683 if (to->index >= LINKCELLS) {
684 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200685 if (link == NULL)
686 return NULL;
Serhiy Storchaka57a01d32016-04-10 18:05:40 +0300687 Py_SETREF(to->dataobj, (teedataobject *)link);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000688 to->index = 0;
689 }
690 value = teedataobject_getitem(to->dataobj, to->index);
691 if (value == NULL)
692 return NULL;
693 to->index++;
694 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000695}
696
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000697static int
698tee_traverse(teeobject *to, visitproc visit, void *arg)
699{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000700 Py_VISIT((PyObject *)to->dataobj);
701 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000702}
703
Raymond Hettingerad983e72003-11-12 14:32:26 +0000704static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530705tee_copy(teeobject *to, PyObject *Py_UNUSED(ignored))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000706{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000707 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000708
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000709 newto = PyObject_GC_New(teeobject, &tee_type);
710 if (newto == NULL)
711 return NULL;
712 Py_INCREF(to->dataobj);
713 newto->dataobj = to->dataobj;
714 newto->index = to->index;
715 newto->weakreflist = NULL;
716 PyObject_GC_Track(newto);
717 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000718}
719
720PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
721
722static PyObject *
723tee_fromiterable(PyObject *iterable)
724{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000725 teeobject *to;
726 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000727
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000728 it = PyObject_GetIter(iterable);
729 if (it == NULL)
730 return NULL;
731 if (PyObject_TypeCheck(it, &tee_type)) {
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530732 to = (teeobject *)tee_copy((teeobject *)it, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000733 goto done;
734 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000735
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000736 to = PyObject_GC_New(teeobject, &tee_type);
737 if (to == NULL)
738 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000739 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000740 if (!to->dataobj) {
741 PyObject_GC_Del(to);
742 to = NULL;
743 goto done;
744 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000745
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000746 to->index = 0;
747 to->weakreflist = NULL;
748 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000749done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000750 Py_XDECREF(it);
751 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000752}
753
Tal Einatc4bccd32018-09-12 00:49:13 +0300754/*[clinic input]
755@classmethod
756itertools._tee.__new__
757 iterable: object
758 /
759Iterator wrapped to make it copyable.
760[clinic start generated code]*/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000761
Tal Einatc4bccd32018-09-12 00:49:13 +0300762static PyObject *
763itertools__tee_impl(PyTypeObject *type, PyObject *iterable)
764/*[clinic end generated code: output=b02d3fd26c810c3f input=adc0779d2afe37a2]*/
765{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000766 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000767}
768
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000769static int
770tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000771{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000772 if (to->weakreflist != NULL)
773 PyObject_ClearWeakRefs((PyObject *) to);
774 Py_CLEAR(to->dataobj);
775 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000776}
777
778static void
779tee_dealloc(teeobject *to)
780{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000781 PyObject_GC_UnTrack(to);
782 tee_clear(to);
783 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000784}
785
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000786static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530787tee_reduce(teeobject *to, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000788{
789 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
790}
791
792static PyObject *
793tee_setstate(teeobject *to, PyObject *state)
794{
795 teedataobject *tdo;
796 int index;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300797 if (!PyTuple_Check(state)) {
798 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000799 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300800 }
801 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
802 return NULL;
803 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000804 if (index < 0 || index > LINKCELLS) {
805 PyErr_SetString(PyExc_ValueError, "Index out of range");
806 return NULL;
807 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200808 Py_INCREF(tdo);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300809 Py_XSETREF(to->dataobj, tdo);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000810 to->index = index;
811 Py_RETURN_NONE;
812}
813
Raymond Hettingerad983e72003-11-12 14:32:26 +0000814static PyMethodDef tee_methods[] = {
Raymond Hettingera6a2d442015-08-16 14:38:07 -0700815 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
816 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
817 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000818 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000819};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000820
821static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000822 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000823 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000824 sizeof(teeobject), /* tp_basicsize */
825 0, /* tp_itemsize */
826 /* methods */
827 (destructor)tee_dealloc, /* tp_dealloc */
828 0, /* tp_print */
829 0, /* tp_getattr */
830 0, /* tp_setattr */
831 0, /* tp_reserved */
832 0, /* tp_repr */
833 0, /* tp_as_number */
834 0, /* tp_as_sequence */
835 0, /* tp_as_mapping */
836 0, /* tp_hash */
837 0, /* tp_call */
838 0, /* tp_str */
839 0, /* tp_getattro */
840 0, /* tp_setattro */
841 0, /* tp_as_buffer */
842 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300843 itertools__tee__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000844 (traverseproc)tee_traverse, /* tp_traverse */
845 (inquiry)tee_clear, /* tp_clear */
846 0, /* tp_richcompare */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700847 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000848 PyObject_SelfIter, /* tp_iter */
849 (iternextfunc)tee_next, /* tp_iternext */
850 tee_methods, /* tp_methods */
851 0, /* tp_members */
852 0, /* tp_getset */
853 0, /* tp_base */
854 0, /* tp_dict */
855 0, /* tp_descr_get */
856 0, /* tp_descr_set */
857 0, /* tp_dictoffset */
858 0, /* tp_init */
859 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300860 itertools__tee, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000861 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000862};
863
Tal Einatc4bccd32018-09-12 00:49:13 +0300864/*[clinic input]
865itertools.tee
866 iterable: object
867 n: Py_ssize_t = 2
868 /
869Returns a tuple of n independent iterators.
870[clinic start generated code]*/
871
Raymond Hettingerad983e72003-11-12 14:32:26 +0000872static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300873itertools_tee_impl(PyObject *module, PyObject *iterable, Py_ssize_t n)
874/*[clinic end generated code: output=1c64519cd859c2f0 input=c99a1472c425d66d]*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000875{
Tal Einatc4bccd32018-09-12 00:49:13 +0300876 Py_ssize_t i;
877 PyObject *it, *copyable, *copyfunc, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200878 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000879
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000880 if (n < 0) {
881 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
882 return NULL;
883 }
884 result = PyTuple_New(n);
885 if (result == NULL)
886 return NULL;
887 if (n == 0)
888 return result;
889 it = PyObject_GetIter(iterable);
890 if (it == NULL) {
891 Py_DECREF(result);
892 return NULL;
893 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200894
Serhiy Storchakaf320be72018-01-25 10:49:40 +0200895 if (_PyObject_LookupAttrId(it, &PyId___copy__, &copyfunc) < 0) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200896 Py_DECREF(it);
897 Py_DECREF(result);
898 return NULL;
899 }
Serhiy Storchakaf320be72018-01-25 10:49:40 +0200900 if (copyfunc != NULL) {
901 copyable = it;
902 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200903 else {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000904 copyable = tee_fromiterable(it);
905 Py_DECREF(it);
906 if (copyable == NULL) {
907 Py_DECREF(result);
908 return NULL;
909 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200910 copyfunc = _PyObject_GetAttrId(copyable, &PyId___copy__);
911 if (copyfunc == NULL) {
912 Py_DECREF(copyable);
913 Py_DECREF(result);
914 return NULL;
915 }
916 }
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200917
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200918 PyTuple_SET_ITEM(result, 0, copyable);
919 for (i = 1; i < n; i++) {
920 copyable = _PyObject_CallNoArg(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000921 if (copyable == NULL) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200922 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000923 Py_DECREF(result);
924 return NULL;
925 }
926 PyTuple_SET_ITEM(result, i, copyable);
927 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200928 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000929 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000930}
931
Raymond Hettingerad983e72003-11-12 14:32:26 +0000932
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700933/* cycle object **************************************************************/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000934
935typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000936 PyObject_HEAD
937 PyObject *it;
938 PyObject *saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700939 Py_ssize_t index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000940 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000941} cycleobject;
942
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000943static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000944
Tal Einatc4bccd32018-09-12 00:49:13 +0300945/*[clinic input]
946@classmethod
947itertools.cycle.__new__
948 iterable: object
949 /
950Return elements from the iterable until it is exhausted. Then repeat the sequence indefinitely.
951[clinic start generated code]*/
952
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000953static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300954itertools_cycle_impl(PyTypeObject *type, PyObject *iterable)
955/*[clinic end generated code: output=f60e5ec17a45b35c input=9d1d84bcf66e908b]*/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000956{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000957 PyObject *it;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000958 PyObject *saved;
959 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000960
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000961 /* Get iterator. */
962 it = PyObject_GetIter(iterable);
963 if (it == NULL)
964 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000965
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000966 saved = PyList_New(0);
967 if (saved == NULL) {
968 Py_DECREF(it);
969 return NULL;
970 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000971
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000972 /* create cycleobject structure */
973 lz = (cycleobject *)type->tp_alloc(type, 0);
974 if (lz == NULL) {
975 Py_DECREF(it);
976 Py_DECREF(saved);
977 return NULL;
978 }
979 lz->it = it;
980 lz->saved = saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700981 lz->index = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000982 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000983
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000984 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000985}
986
987static void
988cycle_dealloc(cycleobject *lz)
989{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000990 PyObject_GC_UnTrack(lz);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000991 Py_XDECREF(lz->it);
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700992 Py_XDECREF(lz->saved);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000993 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000994}
995
996static int
997cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
998{
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700999 if (lz->it)
1000 Py_VISIT(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001001 Py_VISIT(lz->saved);
1002 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001003}
1004
1005static PyObject *
1006cycle_next(cycleobject *lz)
1007{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001008 PyObject *item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001009
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001010 if (lz->it != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001011 item = PyIter_Next(lz->it);
1012 if (item != NULL) {
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001013 if (lz->firstpass)
1014 return item;
1015 if (PyList_Append(lz->saved, item)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001016 Py_DECREF(item);
1017 return NULL;
1018 }
1019 return item;
1020 }
Raymond Hettinger98958fe2015-08-15 15:09:30 -07001021 /* Note: StopIteration is already cleared by PyIter_Next() */
1022 if (PyErr_Occurred())
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001023 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001024 Py_CLEAR(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001025 }
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001026 if (PyList_GET_SIZE(lz->saved) == 0)
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001027 return NULL;
1028 item = PyList_GET_ITEM(lz->saved, lz->index);
1029 lz->index++;
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001030 if (lz->index >= PyList_GET_SIZE(lz->saved))
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001031 lz->index = 0;
1032 Py_INCREF(item);
1033 return item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001034}
1035
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001036static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301037cycle_reduce(cycleobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001038{
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001039 /* Create a new cycle with the iterator tuple, then set the saved state */
1040 if (lz->it == NULL) {
1041 PyObject *it = PyObject_GetIter(lz->saved);
1042 if (it == NULL)
1043 return NULL;
1044 if (lz->index != 0) {
1045 _Py_IDENTIFIER(__setstate__);
1046 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
1047 "n", lz->index);
1048 if (res == NULL) {
1049 Py_DECREF(it);
1050 return NULL;
1051 }
1052 Py_DECREF(res);
1053 }
1054 return Py_BuildValue("O(N)(Oi)", Py_TYPE(lz), it, lz->saved, 1);
1055 }
1056 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), lz->it, lz->saved,
1057 lz->firstpass);
Raymond Hettingerc39786d2015-08-15 15:16:12 -07001058}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001059
1060static PyObject *
1061cycle_setstate(cycleobject *lz, PyObject *state)
1062{
1063 PyObject *saved=NULL;
1064 int firstpass;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001065 if (!PyTuple_Check(state)) {
1066 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001067 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001068 }
1069 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
1070 return NULL;
1071 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001072 Py_INCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001073 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001074 lz->firstpass = firstpass != 0;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001075 lz->index = 0;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001076 Py_RETURN_NONE;
1077}
1078
1079static PyMethodDef cycle_methods[] = {
1080 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1081 reduce_doc},
1082 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1083 setstate_doc},
1084 {NULL, NULL} /* sentinel */
1085};
1086
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001087static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001088 PyVarObject_HEAD_INIT(NULL, 0)
1089 "itertools.cycle", /* tp_name */
1090 sizeof(cycleobject), /* tp_basicsize */
1091 0, /* tp_itemsize */
1092 /* methods */
1093 (destructor)cycle_dealloc, /* tp_dealloc */
1094 0, /* tp_print */
1095 0, /* tp_getattr */
1096 0, /* tp_setattr */
1097 0, /* tp_reserved */
1098 0, /* tp_repr */
1099 0, /* tp_as_number */
1100 0, /* tp_as_sequence */
1101 0, /* tp_as_mapping */
1102 0, /* tp_hash */
1103 0, /* tp_call */
1104 0, /* tp_str */
1105 PyObject_GenericGetAttr, /* tp_getattro */
1106 0, /* tp_setattro */
1107 0, /* tp_as_buffer */
1108 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1109 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001110 itertools_cycle__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001111 (traverseproc)cycle_traverse, /* tp_traverse */
1112 0, /* tp_clear */
1113 0, /* tp_richcompare */
1114 0, /* tp_weaklistoffset */
1115 PyObject_SelfIter, /* tp_iter */
1116 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001117 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001118 0, /* tp_members */
1119 0, /* tp_getset */
1120 0, /* tp_base */
1121 0, /* tp_dict */
1122 0, /* tp_descr_get */
1123 0, /* tp_descr_set */
1124 0, /* tp_dictoffset */
1125 0, /* tp_init */
1126 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001127 itertools_cycle, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001128 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001129};
1130
1131
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001132/* dropwhile object **********************************************************/
1133
1134typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001135 PyObject_HEAD
1136 PyObject *func;
1137 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001138 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001139} dropwhileobject;
1140
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001141static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001142
Tal Einatc4bccd32018-09-12 00:49:13 +03001143/*[clinic input]
1144@classmethod
1145itertools.dropwhile.__new__
1146 predicate as func: object
1147 iterable as seq: object
1148 /
1149Drop items from the iterable while predicate(item) is true.
1150
1151Afterwards, return every element until the iterable is exhausted.
1152[clinic start generated code]*/
1153
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001154static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001155itertools_dropwhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1156/*[clinic end generated code: output=92f9d0d89af149e4 input=d39737147c9f0a26]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001157{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001158 PyObject *it;
1159 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001160
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001161 /* Get iterator. */
1162 it = PyObject_GetIter(seq);
1163 if (it == NULL)
1164 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001165
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001166 /* create dropwhileobject structure */
1167 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1168 if (lz == NULL) {
1169 Py_DECREF(it);
1170 return NULL;
1171 }
1172 Py_INCREF(func);
1173 lz->func = func;
1174 lz->it = it;
1175 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001176
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001177 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001178}
1179
1180static void
1181dropwhile_dealloc(dropwhileobject *lz)
1182{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001183 PyObject_GC_UnTrack(lz);
1184 Py_XDECREF(lz->func);
1185 Py_XDECREF(lz->it);
1186 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001187}
1188
1189static int
1190dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1191{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001192 Py_VISIT(lz->it);
1193 Py_VISIT(lz->func);
1194 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001195}
1196
1197static PyObject *
1198dropwhile_next(dropwhileobject *lz)
1199{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001200 PyObject *item, *good;
1201 PyObject *it = lz->it;
1202 long ok;
1203 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001205 iternext = *Py_TYPE(it)->tp_iternext;
1206 for (;;) {
1207 item = iternext(it);
1208 if (item == NULL)
1209 return NULL;
1210 if (lz->start == 1)
1211 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001212
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01001213 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001214 if (good == NULL) {
1215 Py_DECREF(item);
1216 return NULL;
1217 }
1218 ok = PyObject_IsTrue(good);
1219 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001220 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001221 lz->start = 1;
1222 return item;
1223 }
1224 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001225 if (ok < 0)
1226 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001227 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001228}
1229
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001230static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301231dropwhile_reduce(dropwhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001232{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001233 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 +00001234}
1235
1236static PyObject *
1237dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1238{
1239 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001240 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001241 return NULL;
1242 lz->start = start;
1243 Py_RETURN_NONE;
1244}
1245
1246static PyMethodDef dropwhile_methods[] = {
1247 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1248 reduce_doc},
1249 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1250 setstate_doc},
1251 {NULL, NULL} /* sentinel */
1252};
1253
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001254static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001255 PyVarObject_HEAD_INIT(NULL, 0)
1256 "itertools.dropwhile", /* tp_name */
1257 sizeof(dropwhileobject), /* tp_basicsize */
1258 0, /* tp_itemsize */
1259 /* methods */
1260 (destructor)dropwhile_dealloc, /* tp_dealloc */
1261 0, /* tp_print */
1262 0, /* tp_getattr */
1263 0, /* tp_setattr */
1264 0, /* tp_reserved */
1265 0, /* tp_repr */
1266 0, /* tp_as_number */
1267 0, /* tp_as_sequence */
1268 0, /* tp_as_mapping */
1269 0, /* tp_hash */
1270 0, /* tp_call */
1271 0, /* tp_str */
1272 PyObject_GenericGetAttr, /* tp_getattro */
1273 0, /* tp_setattro */
1274 0, /* tp_as_buffer */
1275 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1276 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001277 itertools_dropwhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001278 (traverseproc)dropwhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001279 0, /* tp_clear */
1280 0, /* tp_richcompare */
1281 0, /* tp_weaklistoffset */
1282 PyObject_SelfIter, /* tp_iter */
1283 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001284 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001285 0, /* tp_members */
1286 0, /* tp_getset */
1287 0, /* tp_base */
1288 0, /* tp_dict */
1289 0, /* tp_descr_get */
1290 0, /* tp_descr_set */
1291 0, /* tp_dictoffset */
1292 0, /* tp_init */
1293 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001294 itertools_dropwhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001295 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001296};
1297
1298
1299/* takewhile object **********************************************************/
1300
1301typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001302 PyObject_HEAD
1303 PyObject *func;
1304 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001305 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001306} takewhileobject;
1307
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001308static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001309
Tal Einatc4bccd32018-09-12 00:49:13 +03001310/*[clinic input]
1311@classmethod
1312itertools.takewhile.__new__
1313 predicate as func: object
1314 iterable as seq: object
1315 /
1316Return successive entries from an iterable as long as the predicate evaluates to true for each entry.
1317[clinic start generated code]*/
1318
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001319static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001320itertools_takewhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1321/*[clinic end generated code: output=bb179ea7864e2ef6 input=ba5255f7519aa119]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001322{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001323 PyObject *it;
1324 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001325
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001326 /* Get iterator. */
1327 it = PyObject_GetIter(seq);
1328 if (it == NULL)
1329 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001330
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001331 /* create takewhileobject structure */
1332 lz = (takewhileobject *)type->tp_alloc(type, 0);
1333 if (lz == NULL) {
1334 Py_DECREF(it);
1335 return NULL;
1336 }
1337 Py_INCREF(func);
1338 lz->func = func;
1339 lz->it = it;
1340 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001341
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001342 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001343}
1344
1345static void
1346takewhile_dealloc(takewhileobject *lz)
1347{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001348 PyObject_GC_UnTrack(lz);
1349 Py_XDECREF(lz->func);
1350 Py_XDECREF(lz->it);
1351 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001352}
1353
1354static int
1355takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1356{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001357 Py_VISIT(lz->it);
1358 Py_VISIT(lz->func);
1359 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001360}
1361
1362static PyObject *
1363takewhile_next(takewhileobject *lz)
1364{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001365 PyObject *item, *good;
1366 PyObject *it = lz->it;
1367 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001368
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001369 if (lz->stop == 1)
1370 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001371
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001372 item = (*Py_TYPE(it)->tp_iternext)(it);
1373 if (item == NULL)
1374 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001375
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01001376 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001377 if (good == NULL) {
1378 Py_DECREF(item);
1379 return NULL;
1380 }
1381 ok = PyObject_IsTrue(good);
1382 Py_DECREF(good);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001383 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001384 return item;
1385 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001386 if (ok == 0)
1387 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001388 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001389}
1390
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001391static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301392takewhile_reduce(takewhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001393{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001394 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 +00001395}
1396
1397static PyObject *
1398takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1399{
1400 int stop = PyObject_IsTrue(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001401
Antoine Pitrou721738f2012-08-15 23:20:39 +02001402 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001403 return NULL;
1404 lz->stop = stop;
1405 Py_RETURN_NONE;
1406}
1407
1408static PyMethodDef takewhile_reduce_methods[] = {
1409 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1410 reduce_doc},
1411 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1412 setstate_doc},
1413 {NULL, NULL} /* sentinel */
1414};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001415
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001416static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001417 PyVarObject_HEAD_INIT(NULL, 0)
1418 "itertools.takewhile", /* tp_name */
1419 sizeof(takewhileobject), /* tp_basicsize */
1420 0, /* tp_itemsize */
1421 /* methods */
1422 (destructor)takewhile_dealloc, /* tp_dealloc */
1423 0, /* tp_print */
1424 0, /* tp_getattr */
1425 0, /* tp_setattr */
1426 0, /* tp_reserved */
1427 0, /* tp_repr */
1428 0, /* tp_as_number */
1429 0, /* tp_as_sequence */
1430 0, /* tp_as_mapping */
1431 0, /* tp_hash */
1432 0, /* tp_call */
1433 0, /* tp_str */
1434 PyObject_GenericGetAttr, /* tp_getattro */
1435 0, /* tp_setattro */
1436 0, /* tp_as_buffer */
1437 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1438 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001439 itertools_takewhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001440 (traverseproc)takewhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001441 0, /* tp_clear */
1442 0, /* tp_richcompare */
1443 0, /* tp_weaklistoffset */
1444 PyObject_SelfIter, /* tp_iter */
1445 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001446 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001447 0, /* tp_members */
1448 0, /* tp_getset */
1449 0, /* tp_base */
1450 0, /* tp_dict */
1451 0, /* tp_descr_get */
1452 0, /* tp_descr_set */
1453 0, /* tp_dictoffset */
1454 0, /* tp_init */
1455 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001456 itertools_takewhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001457 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001458};
1459
1460
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001461/* islice object *************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001462
1463typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001464 PyObject_HEAD
1465 PyObject *it;
1466 Py_ssize_t next;
1467 Py_ssize_t stop;
1468 Py_ssize_t step;
1469 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001470} isliceobject;
1471
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001472static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001473
1474static PyObject *
1475islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1476{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001477 PyObject *seq;
1478 Py_ssize_t start=0, stop=-1, step=1;
1479 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1480 Py_ssize_t numargs;
1481 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001482
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001483 if (type == &islice_type && !_PyArg_NoKeywords("islice", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001484 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001485
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001486 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1487 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001488
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001489 numargs = PyTuple_Size(args);
1490 if (numargs == 2) {
1491 if (a1 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001492 stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001493 if (stop == -1) {
1494 if (PyErr_Occurred())
1495 PyErr_Clear();
1496 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001497 "Stop argument for islice() must be None or "
1498 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001499 return NULL;
1500 }
1501 }
1502 } else {
1503 if (a1 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001504 start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001505 if (start == -1 && PyErr_Occurred())
1506 PyErr_Clear();
1507 if (a2 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001508 stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001509 if (stop == -1) {
1510 if (PyErr_Occurred())
1511 PyErr_Clear();
1512 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001513 "Stop argument for islice() must be None or "
1514 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001515 return NULL;
1516 }
1517 }
1518 }
1519 if (start<0 || stop<-1) {
1520 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001521 "Indices for islice() must be None or "
1522 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001523 return NULL;
1524 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001525
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001526 if (a3 != NULL) {
1527 if (a3 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001528 step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001529 if (step == -1 && PyErr_Occurred())
1530 PyErr_Clear();
1531 }
1532 if (step<1) {
1533 PyErr_SetString(PyExc_ValueError,
1534 "Step for islice() must be a positive integer or None.");
1535 return NULL;
1536 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001537
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001538 /* Get iterator. */
1539 it = PyObject_GetIter(seq);
1540 if (it == NULL)
1541 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001542
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001543 /* create isliceobject structure */
1544 lz = (isliceobject *)type->tp_alloc(type, 0);
1545 if (lz == NULL) {
1546 Py_DECREF(it);
1547 return NULL;
1548 }
1549 lz->it = it;
1550 lz->next = start;
1551 lz->stop = stop;
1552 lz->step = step;
1553 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001554
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001555 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001556}
1557
1558static void
1559islice_dealloc(isliceobject *lz)
1560{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001561 PyObject_GC_UnTrack(lz);
1562 Py_XDECREF(lz->it);
1563 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001564}
1565
1566static int
1567islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1568{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001569 Py_VISIT(lz->it);
1570 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001571}
1572
1573static PyObject *
1574islice_next(isliceobject *lz)
1575{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001576 PyObject *item;
1577 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001578 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001579 Py_ssize_t oldnext;
1580 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001581
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001582 if (it == NULL)
1583 return NULL;
1584
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001585 iternext = *Py_TYPE(it)->tp_iternext;
1586 while (lz->cnt < lz->next) {
1587 item = iternext(it);
1588 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001589 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001590 Py_DECREF(item);
1591 lz->cnt++;
1592 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001593 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001594 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001595 item = iternext(it);
1596 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001597 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001598 lz->cnt++;
1599 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001600 /* The (size_t) cast below avoids the danger of undefined
1601 behaviour from signed integer overflow. */
1602 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001603 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1604 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001605 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001606
1607empty:
1608 Py_CLEAR(lz->it);
1609 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001610}
1611
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001612static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301613islice_reduce(isliceobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001614{
1615 /* When unpickled, generate a new object with the same bounds,
1616 * then 'setstate' with the next and count
1617 */
1618 PyObject *stop;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001619
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001620 if (lz->it == NULL) {
1621 PyObject *empty_list;
1622 PyObject *empty_it;
1623 empty_list = PyList_New(0);
1624 if (empty_list == NULL)
1625 return NULL;
1626 empty_it = PyObject_GetIter(empty_list);
1627 Py_DECREF(empty_list);
1628 if (empty_it == NULL)
1629 return NULL;
1630 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1631 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001632 if (lz->stop == -1) {
1633 stop = Py_None;
1634 Py_INCREF(stop);
1635 } else {
1636 stop = PyLong_FromSsize_t(lz->stop);
1637 if (stop == NULL)
1638 return NULL;
1639 }
1640 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1641 lz->it, lz->next, stop, lz->step,
1642 lz->cnt);
1643}
1644
1645static PyObject *
1646islice_setstate(isliceobject *lz, PyObject *state)
1647{
1648 Py_ssize_t cnt = PyLong_AsSsize_t(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001649
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001650 if (cnt == -1 && PyErr_Occurred())
1651 return NULL;
1652 lz->cnt = cnt;
1653 Py_RETURN_NONE;
1654}
1655
1656static PyMethodDef islice_methods[] = {
1657 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1658 reduce_doc},
1659 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1660 setstate_doc},
1661 {NULL, NULL} /* sentinel */
1662};
1663
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001664PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001665"islice(iterable, stop) --> islice object\n\
1666islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001667\n\
1668Return an iterator whose next() method returns selected values from an\n\
1669iterable. If start is specified, will skip all preceding elements;\n\
1670otherwise, start defaults to zero. Step defaults to one. If\n\
oldkaa0735f2018-02-02 16:52:55 +08001671specified as another value, step determines how many values are\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001672skipped between successive calls. Works like a slice() on a list\n\
1673but returns an iterator.");
1674
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001675static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001676 PyVarObject_HEAD_INIT(NULL, 0)
1677 "itertools.islice", /* tp_name */
1678 sizeof(isliceobject), /* tp_basicsize */
1679 0, /* tp_itemsize */
1680 /* methods */
1681 (destructor)islice_dealloc, /* tp_dealloc */
1682 0, /* tp_print */
1683 0, /* tp_getattr */
1684 0, /* tp_setattr */
1685 0, /* tp_reserved */
1686 0, /* tp_repr */
1687 0, /* tp_as_number */
1688 0, /* tp_as_sequence */
1689 0, /* tp_as_mapping */
1690 0, /* tp_hash */
1691 0, /* tp_call */
1692 0, /* tp_str */
1693 PyObject_GenericGetAttr, /* tp_getattro */
1694 0, /* tp_setattro */
1695 0, /* tp_as_buffer */
1696 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1697 Py_TPFLAGS_BASETYPE, /* tp_flags */
1698 islice_doc, /* tp_doc */
1699 (traverseproc)islice_traverse, /* tp_traverse */
1700 0, /* tp_clear */
1701 0, /* tp_richcompare */
1702 0, /* tp_weaklistoffset */
1703 PyObject_SelfIter, /* tp_iter */
1704 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001705 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001706 0, /* tp_members */
1707 0, /* tp_getset */
1708 0, /* tp_base */
1709 0, /* tp_dict */
1710 0, /* tp_descr_get */
1711 0, /* tp_descr_set */
1712 0, /* tp_dictoffset */
1713 0, /* tp_init */
1714 0, /* tp_alloc */
1715 islice_new, /* tp_new */
1716 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001717};
1718
1719
1720/* starmap object ************************************************************/
1721
1722typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001723 PyObject_HEAD
1724 PyObject *func;
1725 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001726} starmapobject;
1727
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001728static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001729
Tal Einatc4bccd32018-09-12 00:49:13 +03001730/*[clinic input]
1731@classmethod
1732itertools.starmap.__new__
1733 function as func: object
1734 iterable as seq: object
1735 /
1736Return an iterator whose values are returned from the function evaluated with an argument tuple taken from the given sequence.
1737[clinic start generated code]*/
1738
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001739static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001740itertools_starmap_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1741/*[clinic end generated code: output=79eeb81d452c6e8d input=844766df6a0d4dad]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001742{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001743 PyObject *it;
1744 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001745
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001746 /* Get iterator. */
1747 it = PyObject_GetIter(seq);
1748 if (it == NULL)
1749 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001750
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001751 /* create starmapobject structure */
1752 lz = (starmapobject *)type->tp_alloc(type, 0);
1753 if (lz == NULL) {
1754 Py_DECREF(it);
1755 return NULL;
1756 }
1757 Py_INCREF(func);
1758 lz->func = func;
1759 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001760
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001761 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001762}
1763
1764static void
1765starmap_dealloc(starmapobject *lz)
1766{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001767 PyObject_GC_UnTrack(lz);
1768 Py_XDECREF(lz->func);
1769 Py_XDECREF(lz->it);
1770 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001771}
1772
1773static int
1774starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1775{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001776 Py_VISIT(lz->it);
1777 Py_VISIT(lz->func);
1778 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001779}
1780
1781static PyObject *
1782starmap_next(starmapobject *lz)
1783{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001784 PyObject *args;
1785 PyObject *result;
1786 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001787
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001788 args = (*Py_TYPE(it)->tp_iternext)(it);
1789 if (args == NULL)
1790 return NULL;
1791 if (!PyTuple_CheckExact(args)) {
1792 PyObject *newargs = PySequence_Tuple(args);
1793 Py_DECREF(args);
1794 if (newargs == NULL)
1795 return NULL;
1796 args = newargs;
1797 }
1798 result = PyObject_Call(lz->func, args, NULL);
1799 Py_DECREF(args);
1800 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001801}
1802
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001803static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301804starmap_reduce(starmapobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001805{
1806 /* Just pickle the iterator */
1807 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1808}
1809
1810static PyMethodDef starmap_methods[] = {
1811 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1812 reduce_doc},
1813 {NULL, NULL} /* sentinel */
1814};
1815
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001816static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001817 PyVarObject_HEAD_INIT(NULL, 0)
1818 "itertools.starmap", /* tp_name */
1819 sizeof(starmapobject), /* tp_basicsize */
1820 0, /* tp_itemsize */
1821 /* methods */
1822 (destructor)starmap_dealloc, /* tp_dealloc */
1823 0, /* tp_print */
1824 0, /* tp_getattr */
1825 0, /* tp_setattr */
1826 0, /* tp_reserved */
1827 0, /* tp_repr */
1828 0, /* tp_as_number */
1829 0, /* tp_as_sequence */
1830 0, /* tp_as_mapping */
1831 0, /* tp_hash */
1832 0, /* tp_call */
1833 0, /* tp_str */
1834 PyObject_GenericGetAttr, /* tp_getattro */
1835 0, /* tp_setattro */
1836 0, /* tp_as_buffer */
1837 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1838 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001839 itertools_starmap__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001840 (traverseproc)starmap_traverse, /* tp_traverse */
1841 0, /* tp_clear */
1842 0, /* tp_richcompare */
1843 0, /* tp_weaklistoffset */
1844 PyObject_SelfIter, /* tp_iter */
1845 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001846 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001847 0, /* tp_members */
1848 0, /* tp_getset */
1849 0, /* tp_base */
1850 0, /* tp_dict */
1851 0, /* tp_descr_get */
1852 0, /* tp_descr_set */
1853 0, /* tp_dictoffset */
1854 0, /* tp_init */
1855 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001856 itertools_starmap, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001857 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001858};
1859
1860
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001861/* chain object **************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001862
1863typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001864 PyObject_HEAD
1865 PyObject *source; /* Iterator over input iterables */
1866 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001867} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001868
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001869static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001870
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001871static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001872chain_new_internal(PyTypeObject *type, PyObject *source)
1873{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001874 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001875
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001876 lz = (chainobject *)type->tp_alloc(type, 0);
1877 if (lz == NULL) {
1878 Py_DECREF(source);
1879 return NULL;
1880 }
1881
1882 lz->source = source;
1883 lz->active = NULL;
1884 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001885}
1886
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001887static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001888chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001889{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001890 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001891
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001892 if (type == &chain_type && !_PyArg_NoKeywords("chain", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001893 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001894
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001895 source = PyObject_GetIter(args);
1896 if (source == NULL)
1897 return NULL;
1898
1899 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001900}
1901
Tal Einatc4bccd32018-09-12 00:49:13 +03001902/*[clinic input]
1903@classmethod
1904itertools.chain.from_iterable
1905 iterable as arg: object
1906 /
1907Alternative chain() constructor taking a single iterable argument that evaluates lazily.
1908[clinic start generated code]*/
1909
Christian Heimesf16baeb2008-02-29 14:57:44 +00001910static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001911itertools_chain_from_iterable(PyTypeObject *type, PyObject *arg)
1912/*[clinic end generated code: output=667ae7a7f7b68654 input=72c39e3a2ca3be85]*/
Christian Heimesf16baeb2008-02-29 14:57:44 +00001913{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001914 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001915
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001916 source = PyObject_GetIter(arg);
1917 if (source == NULL)
1918 return NULL;
1919
1920 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001921}
1922
1923static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001924chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001925{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001926 PyObject_GC_UnTrack(lz);
1927 Py_XDECREF(lz->active);
1928 Py_XDECREF(lz->source);
1929 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001930}
1931
Raymond Hettinger2012f172003-02-07 05:32:58 +00001932static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001933chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001934{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001935 Py_VISIT(lz->source);
1936 Py_VISIT(lz->active);
1937 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001938}
1939
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001940static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001941chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001942{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001943 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001944
T. Wouters5466d4a2017-03-30 09:58:35 -07001945 /* lz->source is the iterator of iterables. If it's NULL, we've already
1946 * consumed them all. lz->active is the current iterator. If it's NULL,
1947 * we should grab a new one from lz->source. */
1948 while (lz->source != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001949 if (lz->active == NULL) {
T. Wouters5466d4a2017-03-30 09:58:35 -07001950 PyObject *iterable = PyIter_Next(lz->source);
1951 if (iterable == NULL) {
1952 Py_CLEAR(lz->source);
1953 return NULL; /* no more input sources */
1954 }
1955 lz->active = PyObject_GetIter(iterable);
1956 Py_DECREF(iterable);
1957 if (lz->active == NULL) {
1958 Py_CLEAR(lz->source);
1959 return NULL; /* input not iterable */
1960 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001961 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001962 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1963 if (item != NULL)
1964 return item;
1965 if (PyErr_Occurred()) {
1966 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1967 PyErr_Clear();
1968 else
1969 return NULL; /* input raised an exception */
1970 }
1971 /* lz->active is consumed, try with the next iterable. */
1972 Py_CLEAR(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001973 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001974 /* Everything had been consumed already. */
1975 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001976}
1977
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001978static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301979chain_reduce(chainobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001980{
1981 if (lz->source) {
1982 /* we can't pickle function objects (itertools.from_iterable) so
1983 * we must use setstate to replace the iterable. One day we
1984 * will fix pickling of functions
1985 */
1986 if (lz->active) {
1987 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1988 } else {
1989 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1990 }
1991 } else {
1992 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1993 }
1994 return NULL;
1995}
1996
1997static PyObject *
1998chain_setstate(chainobject *lz, PyObject *state)
1999{
2000 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002001
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002002 if (!PyTuple_Check(state)) {
2003 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002004 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002005 }
2006 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
2007 return NULL;
2008 }
2009 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
2010 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
2011 return NULL;
2012 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002013
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002014 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002015 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002016 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002017 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002018 Py_RETURN_NONE;
2019}
2020
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002021PyDoc_STRVAR(chain_doc,
2022"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002023\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002024Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002025first iterable until it is exhausted, then elements from the next\n\
2026iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002027
Christian Heimesf16baeb2008-02-29 14:57:44 +00002028static PyMethodDef chain_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03002029 ITERTOOLS_CHAIN_FROM_ITERABLE_METHODDEF
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002030 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
2031 reduce_doc},
2032 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
2033 setstate_doc},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07002034 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00002035};
2036
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002037static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002038 PyVarObject_HEAD_INIT(NULL, 0)
2039 "itertools.chain", /* tp_name */
2040 sizeof(chainobject), /* tp_basicsize */
2041 0, /* tp_itemsize */
2042 /* methods */
2043 (destructor)chain_dealloc, /* tp_dealloc */
2044 0, /* tp_print */
2045 0, /* tp_getattr */
2046 0, /* tp_setattr */
2047 0, /* tp_reserved */
2048 0, /* tp_repr */
2049 0, /* tp_as_number */
2050 0, /* tp_as_sequence */
2051 0, /* tp_as_mapping */
2052 0, /* tp_hash */
2053 0, /* tp_call */
2054 0, /* tp_str */
2055 PyObject_GenericGetAttr, /* tp_getattro */
2056 0, /* tp_setattro */
2057 0, /* tp_as_buffer */
2058 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2059 Py_TPFLAGS_BASETYPE, /* tp_flags */
2060 chain_doc, /* tp_doc */
2061 (traverseproc)chain_traverse, /* tp_traverse */
2062 0, /* tp_clear */
2063 0, /* tp_richcompare */
2064 0, /* tp_weaklistoffset */
2065 PyObject_SelfIter, /* tp_iter */
2066 (iternextfunc)chain_next, /* tp_iternext */
2067 chain_methods, /* tp_methods */
2068 0, /* tp_members */
2069 0, /* tp_getset */
2070 0, /* tp_base */
2071 0, /* tp_dict */
2072 0, /* tp_descr_get */
2073 0, /* tp_descr_set */
2074 0, /* tp_dictoffset */
2075 0, /* tp_init */
2076 0, /* tp_alloc */
2077 chain_new, /* tp_new */
2078 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002079};
2080
2081
Christian Heimesc3f30c42008-02-22 16:37:40 +00002082/* product object ************************************************************/
2083
2084typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002085 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002086 PyObject *pools; /* tuple of pool tuples */
2087 Py_ssize_t *indices; /* one index per pool */
2088 PyObject *result; /* most recently returned result tuple */
2089 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002090} productobject;
2091
2092static PyTypeObject product_type;
2093
2094static PyObject *
2095product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2096{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002097 productobject *lz;
2098 Py_ssize_t nargs, npools, repeat=1;
2099 PyObject *pools = NULL;
2100 Py_ssize_t *indices = NULL;
2101 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002102
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002103 if (kwds != NULL) {
2104 char *kwlist[] = {"repeat", 0};
2105 PyObject *tmpargs = PyTuple_New(0);
2106 if (tmpargs == NULL)
2107 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002108 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2109 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002110 Py_DECREF(tmpargs);
2111 return NULL;
2112 }
2113 Py_DECREF(tmpargs);
2114 if (repeat < 0) {
2115 PyErr_SetString(PyExc_ValueError,
2116 "repeat argument cannot be negative");
2117 return NULL;
2118 }
2119 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002120
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002121 assert(PyTuple_CheckExact(args));
2122 if (repeat == 0) {
2123 nargs = 0;
2124 } else {
2125 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002126 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002127 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2128 return NULL;
2129 }
2130 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002131 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002132
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002133 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002134 if (indices == NULL) {
2135 PyErr_NoMemory();
2136 goto error;
2137 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002138
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002139 pools = PyTuple_New(npools);
2140 if (pools == NULL)
2141 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002142
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002143 for (i=0; i < nargs ; ++i) {
2144 PyObject *item = PyTuple_GET_ITEM(args, i);
2145 PyObject *pool = PySequence_Tuple(item);
2146 if (pool == NULL)
2147 goto error;
2148 PyTuple_SET_ITEM(pools, i, pool);
2149 indices[i] = 0;
2150 }
2151 for ( ; i < npools; ++i) {
2152 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2153 Py_INCREF(pool);
2154 PyTuple_SET_ITEM(pools, i, pool);
2155 indices[i] = 0;
2156 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002157
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002158 /* create productobject structure */
2159 lz = (productobject *)type->tp_alloc(type, 0);
2160 if (lz == NULL)
2161 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002162
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002163 lz->pools = pools;
2164 lz->indices = indices;
2165 lz->result = NULL;
2166 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002167
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002168 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002169
2170error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002171 if (indices != NULL)
2172 PyMem_Free(indices);
2173 Py_XDECREF(pools);
2174 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002175}
2176
2177static void
2178product_dealloc(productobject *lz)
2179{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002180 PyObject_GC_UnTrack(lz);
2181 Py_XDECREF(lz->pools);
2182 Py_XDECREF(lz->result);
2183 if (lz->indices != NULL)
2184 PyMem_Free(lz->indices);
2185 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002186}
2187
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002188static PyObject *
2189product_sizeof(productobject *lz, void *unused)
2190{
2191 Py_ssize_t res;
2192
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002193 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002194 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2195 return PyLong_FromSsize_t(res);
2196}
2197
2198PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2199
Christian Heimesc3f30c42008-02-22 16:37:40 +00002200static int
2201product_traverse(productobject *lz, visitproc visit, void *arg)
2202{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002203 Py_VISIT(lz->pools);
2204 Py_VISIT(lz->result);
2205 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002206}
2207
2208static PyObject *
2209product_next(productobject *lz)
2210{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002211 PyObject *pool;
2212 PyObject *elem;
2213 PyObject *oldelem;
2214 PyObject *pools = lz->pools;
2215 PyObject *result = lz->result;
2216 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2217 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002218
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002219 if (lz->stopped)
2220 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002221
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002222 if (result == NULL) {
2223 /* On the first pass, return an initial tuple filled with the
2224 first element from each pool. */
2225 result = PyTuple_New(npools);
2226 if (result == NULL)
2227 goto empty;
2228 lz->result = result;
2229 for (i=0; i < npools; i++) {
2230 pool = PyTuple_GET_ITEM(pools, i);
2231 if (PyTuple_GET_SIZE(pool) == 0)
2232 goto empty;
2233 elem = PyTuple_GET_ITEM(pool, 0);
2234 Py_INCREF(elem);
2235 PyTuple_SET_ITEM(result, i, elem);
2236 }
2237 } else {
2238 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002239
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002240 /* Copy the previous result tuple or re-use it if available */
2241 if (Py_REFCNT(result) > 1) {
2242 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002243 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002244 if (result == NULL)
2245 goto empty;
2246 lz->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002247 Py_DECREF(old_result);
2248 }
2249 /* Now, we've got the only copy so we can update it in-place */
2250 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002251
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002252 /* Update the pool indices right-to-left. Only advance to the
2253 next pool when the previous one rolls-over */
2254 for (i=npools-1 ; i >= 0 ; i--) {
2255 pool = PyTuple_GET_ITEM(pools, i);
2256 indices[i]++;
2257 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2258 /* Roll-over and advance to next pool */
2259 indices[i] = 0;
2260 elem = PyTuple_GET_ITEM(pool, 0);
2261 Py_INCREF(elem);
2262 oldelem = PyTuple_GET_ITEM(result, i);
2263 PyTuple_SET_ITEM(result, i, elem);
2264 Py_DECREF(oldelem);
2265 } else {
2266 /* No rollover. Just increment and stop here. */
2267 elem = PyTuple_GET_ITEM(pool, indices[i]);
2268 Py_INCREF(elem);
2269 oldelem = PyTuple_GET_ITEM(result, i);
2270 PyTuple_SET_ITEM(result, i, elem);
2271 Py_DECREF(oldelem);
2272 break;
2273 }
2274 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002275
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002276 /* If i is negative, then the indices have all rolled-over
2277 and we're done. */
2278 if (i < 0)
2279 goto empty;
2280 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002281
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002282 Py_INCREF(result);
2283 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002284
2285empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002286 lz->stopped = 1;
2287 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002288}
2289
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002290static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302291product_reduce(productobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002292{
2293 if (lz->stopped) {
2294 return Py_BuildValue("O(())", Py_TYPE(lz));
2295 } else if (lz->result == NULL) {
2296 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2297 } else {
2298 PyObject *indices;
2299 Py_ssize_t n, i;
2300
2301 /* we must pickle the indices use them for setstate, and
2302 * additionally indicate that the iterator has started
2303 */
2304 n = PyTuple_GET_SIZE(lz->pools);
2305 indices = PyTuple_New(n);
2306 if (indices == NULL)
2307 return NULL;
2308 for (i=0; i<n; i++){
2309 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2310 if (!index) {
2311 Py_DECREF(indices);
2312 return NULL;
2313 }
2314 PyTuple_SET_ITEM(indices, i, index);
2315 }
2316 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2317 }
2318}
2319
2320static PyObject *
2321product_setstate(productobject *lz, PyObject *state)
2322{
2323 PyObject *result;
2324 Py_ssize_t n, i;
2325
2326 n = PyTuple_GET_SIZE(lz->pools);
2327 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2328 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2329 return NULL;
2330 }
2331 for (i=0; i<n; i++)
2332 {
2333 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2334 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002335 PyObject* pool;
2336 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002337 if (index < 0 && PyErr_Occurred())
2338 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002339 pool = PyTuple_GET_ITEM(lz->pools, i);
2340 poolsize = PyTuple_GET_SIZE(pool);
2341 if (poolsize == 0) {
2342 lz->stopped = 1;
2343 Py_RETURN_NONE;
2344 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002345 /* clamp the index */
2346 if (index < 0)
2347 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002348 else if (index > poolsize-1)
2349 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002350 lz->indices[i] = index;
2351 }
2352
2353 result = PyTuple_New(n);
2354 if (!result)
2355 return NULL;
2356 for (i=0; i<n; i++) {
2357 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2358 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2359 Py_INCREF(element);
2360 PyTuple_SET_ITEM(result, i, element);
2361 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002362 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002363 Py_RETURN_NONE;
2364}
2365
2366static PyMethodDef product_methods[] = {
2367 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2368 reduce_doc},
2369 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2370 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002371 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2372 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002373 {NULL, NULL} /* sentinel */
2374};
2375
Christian Heimesc3f30c42008-02-22 16:37:40 +00002376PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002377"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002378\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002379Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002380For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2381The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2382cycle in a manner similar to an odometer (with the rightmost element changing\n\
2383on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002384To compute the product of an iterable with itself, specify the number\n\
2385of repetitions with the optional repeat keyword argument. For example,\n\
2386product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002387product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2388product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2389
2390static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002391 PyVarObject_HEAD_INIT(NULL, 0)
2392 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002393 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002394 0, /* tp_itemsize */
2395 /* methods */
2396 (destructor)product_dealloc, /* tp_dealloc */
2397 0, /* tp_print */
2398 0, /* tp_getattr */
2399 0, /* tp_setattr */
2400 0, /* tp_reserved */
2401 0, /* tp_repr */
2402 0, /* tp_as_number */
2403 0, /* tp_as_sequence */
2404 0, /* tp_as_mapping */
2405 0, /* tp_hash */
2406 0, /* tp_call */
2407 0, /* tp_str */
2408 PyObject_GenericGetAttr, /* tp_getattro */
2409 0, /* tp_setattro */
2410 0, /* tp_as_buffer */
2411 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2412 Py_TPFLAGS_BASETYPE, /* tp_flags */
2413 product_doc, /* tp_doc */
2414 (traverseproc)product_traverse, /* tp_traverse */
2415 0, /* tp_clear */
2416 0, /* tp_richcompare */
2417 0, /* tp_weaklistoffset */
2418 PyObject_SelfIter, /* tp_iter */
2419 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002420 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002421 0, /* tp_members */
2422 0, /* tp_getset */
2423 0, /* tp_base */
2424 0, /* tp_dict */
2425 0, /* tp_descr_get */
2426 0, /* tp_descr_set */
2427 0, /* tp_dictoffset */
2428 0, /* tp_init */
2429 0, /* tp_alloc */
2430 product_new, /* tp_new */
2431 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002432};
2433
2434
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002435/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002436
2437typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002438 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002439 PyObject *pool; /* input converted to a tuple */
2440 Py_ssize_t *indices; /* one index per result element */
2441 PyObject *result; /* most recently returned result tuple */
2442 Py_ssize_t r; /* size of result tuple */
2443 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002444} combinationsobject;
2445
2446static PyTypeObject combinations_type;
2447
Tal Einatc4bccd32018-09-12 00:49:13 +03002448
2449/*[clinic input]
2450@classmethod
2451itertools.combinations.__new__
2452 iterable: object
2453 r: Py_ssize_t
2454Return successive r-length combinations of elements in the iterable.
2455
2456combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)
2457[clinic start generated code]*/
2458
Christian Heimes380f7f22008-02-28 11:19:05 +00002459static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002460itertools_combinations_impl(PyTypeObject *type, PyObject *iterable,
2461 Py_ssize_t r)
2462/*[clinic end generated code: output=87a689b39c40039c input=06bede09e3da20f8]*/
Christian Heimes380f7f22008-02-28 11:19:05 +00002463{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002464 combinationsobject *co;
2465 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002466 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002467 Py_ssize_t *indices = NULL;
2468 Py_ssize_t i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002469
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002470 pool = PySequence_Tuple(iterable);
2471 if (pool == NULL)
2472 goto error;
2473 n = PyTuple_GET_SIZE(pool);
2474 if (r < 0) {
2475 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2476 goto error;
2477 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002478
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002479 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002480 if (indices == NULL) {
2481 PyErr_NoMemory();
2482 goto error;
2483 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002484
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002485 for (i=0 ; i<r ; i++)
2486 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002487
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002488 /* create combinationsobject structure */
2489 co = (combinationsobject *)type->tp_alloc(type, 0);
2490 if (co == NULL)
2491 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002492
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002493 co->pool = pool;
2494 co->indices = indices;
2495 co->result = NULL;
2496 co->r = r;
2497 co->stopped = r > n ? 1 : 0;
2498
2499 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002500
2501error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002502 if (indices != NULL)
2503 PyMem_Free(indices);
2504 Py_XDECREF(pool);
2505 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002506}
2507
2508static void
2509combinations_dealloc(combinationsobject *co)
2510{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002511 PyObject_GC_UnTrack(co);
2512 Py_XDECREF(co->pool);
2513 Py_XDECREF(co->result);
2514 if (co->indices != NULL)
2515 PyMem_Free(co->indices);
2516 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002517}
2518
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002519static PyObject *
2520combinations_sizeof(combinationsobject *co, void *unused)
2521{
2522 Py_ssize_t res;
2523
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002524 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002525 res += co->r * sizeof(Py_ssize_t);
2526 return PyLong_FromSsize_t(res);
2527}
2528
Christian Heimes380f7f22008-02-28 11:19:05 +00002529static int
2530combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2531{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002532 Py_VISIT(co->pool);
2533 Py_VISIT(co->result);
2534 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002535}
2536
2537static PyObject *
2538combinations_next(combinationsobject *co)
2539{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002540 PyObject *elem;
2541 PyObject *oldelem;
2542 PyObject *pool = co->pool;
2543 Py_ssize_t *indices = co->indices;
2544 PyObject *result = co->result;
2545 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2546 Py_ssize_t r = co->r;
2547 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002548
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002549 if (co->stopped)
2550 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002551
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002552 if (result == NULL) {
2553 /* On the first pass, initialize result tuple using the indices */
2554 result = PyTuple_New(r);
2555 if (result == NULL)
2556 goto empty;
2557 co->result = result;
2558 for (i=0; i<r ; i++) {
2559 index = indices[i];
2560 elem = PyTuple_GET_ITEM(pool, index);
2561 Py_INCREF(elem);
2562 PyTuple_SET_ITEM(result, i, elem);
2563 }
2564 } else {
2565 /* Copy the previous result tuple or re-use it if available */
2566 if (Py_REFCNT(result) > 1) {
2567 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002568 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002569 if (result == NULL)
2570 goto empty;
2571 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002572 Py_DECREF(old_result);
2573 }
2574 /* Now, we've got the only copy so we can update it in-place
2575 * CPython's empty tuple is a singleton and cached in
2576 * PyTuple's freelist.
2577 */
2578 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002579
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002580 /* Scan indices right-to-left until finding one that is not
2581 at its maximum (i + n - r). */
2582 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2583 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002584
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002585 /* If i is negative, then the indices are all at
2586 their maximum value and we're done. */
2587 if (i < 0)
2588 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002589
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002590 /* Increment the current index which we know is not at its
2591 maximum. Then move back to the right setting each index
2592 to its lowest possible value (one higher than the index
2593 to its left -- this maintains the sort order invariant). */
2594 indices[i]++;
2595 for (j=i+1 ; j<r ; j++)
2596 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002597
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002598 /* Update the result tuple for the new indices
2599 starting with i, the leftmost index that changed */
2600 for ( ; i<r ; i++) {
2601 index = indices[i];
2602 elem = PyTuple_GET_ITEM(pool, index);
2603 Py_INCREF(elem);
2604 oldelem = PyTuple_GET_ITEM(result, i);
2605 PyTuple_SET_ITEM(result, i, elem);
2606 Py_DECREF(oldelem);
2607 }
2608 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002609
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002610 Py_INCREF(result);
2611 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002612
2613empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002614 co->stopped = 1;
2615 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002616}
2617
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002618static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302619combinations_reduce(combinationsobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002620{
2621 if (lz->result == NULL) {
2622 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2623 } else if (lz->stopped) {
2624 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2625 } else {
2626 PyObject *indices;
2627 Py_ssize_t i;
2628
2629 /* we must pickle the indices and use them for setstate */
2630 indices = PyTuple_New(lz->r);
2631 if (!indices)
2632 return NULL;
2633 for (i=0; i<lz->r; i++)
2634 {
2635 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2636 if (!index) {
2637 Py_DECREF(indices);
2638 return NULL;
2639 }
2640 PyTuple_SET_ITEM(indices, i, index);
2641 }
2642
2643 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2644 }
2645}
2646
2647static PyObject *
2648combinations_setstate(combinationsobject *lz, PyObject *state)
2649{
2650 PyObject *result;
2651 Py_ssize_t i;
2652 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2653
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002654 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002655 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2656 return NULL;
2657 }
2658
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002659 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002660 Py_ssize_t max;
2661 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2662 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002663
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002664 if (index == -1 && PyErr_Occurred())
2665 return NULL; /* not an integer */
2666 max = i + n - lz->r;
2667 /* clamp the index (beware of negative max) */
2668 if (index > max)
2669 index = max;
2670 if (index < 0)
2671 index = 0;
2672 lz->indices[i] = index;
2673 }
2674
2675 result = PyTuple_New(lz->r);
2676 if (result == NULL)
2677 return NULL;
2678 for (i=0; i<lz->r; i++) {
2679 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2680 Py_INCREF(element);
2681 PyTuple_SET_ITEM(result, i, element);
2682 }
2683
Serhiy Storchaka48842712016-04-06 09:45:48 +03002684 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002685 Py_RETURN_NONE;
2686}
2687
2688static PyMethodDef combinations_methods[] = {
2689 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2690 reduce_doc},
2691 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2692 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002693 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2694 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002695 {NULL, NULL} /* sentinel */
2696};
2697
Christian Heimes380f7f22008-02-28 11:19:05 +00002698static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002699 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002700 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002701 sizeof(combinationsobject), /* tp_basicsize */
2702 0, /* tp_itemsize */
2703 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002704 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002705 0, /* tp_print */
2706 0, /* tp_getattr */
2707 0, /* tp_setattr */
2708 0, /* tp_reserved */
2709 0, /* tp_repr */
2710 0, /* tp_as_number */
2711 0, /* tp_as_sequence */
2712 0, /* tp_as_mapping */
2713 0, /* tp_hash */
2714 0, /* tp_call */
2715 0, /* tp_str */
2716 PyObject_GenericGetAttr, /* tp_getattro */
2717 0, /* tp_setattro */
2718 0, /* tp_as_buffer */
2719 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2720 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03002721 itertools_combinations__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002722 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002723 0, /* tp_clear */
2724 0, /* tp_richcompare */
2725 0, /* tp_weaklistoffset */
2726 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002727 (iternextfunc)combinations_next, /* tp_iternext */
2728 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002729 0, /* tp_members */
2730 0, /* tp_getset */
2731 0, /* tp_base */
2732 0, /* tp_dict */
2733 0, /* tp_descr_get */
2734 0, /* tp_descr_set */
2735 0, /* tp_dictoffset */
2736 0, /* tp_init */
2737 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03002738 itertools_combinations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002739 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002740};
2741
2742
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002743/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002744
2745/* Equivalent to:
2746
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002747 def combinations_with_replacement(iterable, r):
2748 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2749 # number items returned: (n+r-1)! / r! / (n-1)!
2750 pool = tuple(iterable)
2751 n = len(pool)
2752 indices = [0] * r
2753 yield tuple(pool[i] for i in indices)
2754 while 1:
2755 for i in reversed(range(r)):
2756 if indices[i] != n - 1:
2757 break
2758 else:
2759 return
2760 indices[i:] = [indices[i] + 1] * (r - i)
2761 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002762
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002763 def combinations_with_replacement2(iterable, r):
2764 'Alternate version that filters from product()'
2765 pool = tuple(iterable)
2766 n = len(pool)
2767 for indices in product(range(n), repeat=r):
2768 if sorted(indices) == list(indices):
2769 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002770*/
2771typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002772 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002773 PyObject *pool; /* input converted to a tuple */
2774 Py_ssize_t *indices; /* one index per result element */
2775 PyObject *result; /* most recently returned result tuple */
2776 Py_ssize_t r; /* size of result tuple */
2777 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002778} cwrobject;
2779
2780static PyTypeObject cwr_type;
2781
Tal Einatc4bccd32018-09-12 00:49:13 +03002782/*[clinic input]
2783@classmethod
2784itertools.combinations_with_replacement.__new__
2785 iterable: object
2786 r: Py_ssize_t
2787Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.
2788
2789combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2790[clinic start generated code]*/
2791
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002792static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002793itertools_combinations_with_replacement_impl(PyTypeObject *type,
2794 PyObject *iterable,
2795 Py_ssize_t r)
2796/*[clinic end generated code: output=48b26856d4e659ca input=dc2a8c7ba785fad7]*/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002797{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002798 cwrobject *co;
2799 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002800 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002801 Py_ssize_t *indices = NULL;
2802 Py_ssize_t i;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002803
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002804 pool = PySequence_Tuple(iterable);
2805 if (pool == NULL)
2806 goto error;
2807 n = PyTuple_GET_SIZE(pool);
2808 if (r < 0) {
2809 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2810 goto error;
2811 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002812
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002813 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002814 if (indices == NULL) {
2815 PyErr_NoMemory();
2816 goto error;
2817 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002818
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002819 for (i=0 ; i<r ; i++)
2820 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002821
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002822 /* create cwrobject structure */
2823 co = (cwrobject *)type->tp_alloc(type, 0);
2824 if (co == NULL)
2825 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002826
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002827 co->pool = pool;
2828 co->indices = indices;
2829 co->result = NULL;
2830 co->r = r;
2831 co->stopped = !n && r;
2832
2833 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002834
2835error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002836 if (indices != NULL)
2837 PyMem_Free(indices);
2838 Py_XDECREF(pool);
2839 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002840}
2841
2842static void
2843cwr_dealloc(cwrobject *co)
2844{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002845 PyObject_GC_UnTrack(co);
2846 Py_XDECREF(co->pool);
2847 Py_XDECREF(co->result);
2848 if (co->indices != NULL)
2849 PyMem_Free(co->indices);
2850 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002851}
2852
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002853static PyObject *
2854cwr_sizeof(cwrobject *co, void *unused)
2855{
2856 Py_ssize_t res;
2857
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002858 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002859 res += co->r * sizeof(Py_ssize_t);
2860 return PyLong_FromSsize_t(res);
2861}
2862
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002863static int
2864cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2865{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002866 Py_VISIT(co->pool);
2867 Py_VISIT(co->result);
2868 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002869}
2870
2871static PyObject *
2872cwr_next(cwrobject *co)
2873{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002874 PyObject *elem;
2875 PyObject *oldelem;
2876 PyObject *pool = co->pool;
2877 Py_ssize_t *indices = co->indices;
2878 PyObject *result = co->result;
2879 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2880 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002881 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002882
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002883 if (co->stopped)
2884 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002885
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002886 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002887 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002888 result = PyTuple_New(r);
2889 if (result == NULL)
2890 goto empty;
2891 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002892 if (n > 0) {
2893 elem = PyTuple_GET_ITEM(pool, 0);
2894 for (i=0; i<r ; i++) {
2895 assert(indices[i] == 0);
2896 Py_INCREF(elem);
2897 PyTuple_SET_ITEM(result, i, elem);
2898 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002899 }
2900 } else {
2901 /* Copy the previous result tuple or re-use it if available */
2902 if (Py_REFCNT(result) > 1) {
2903 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002904 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002905 if (result == NULL)
2906 goto empty;
2907 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002908 Py_DECREF(old_result);
2909 }
2910 /* Now, we've got the only copy so we can update it in-place CPython's
2911 empty tuple is a singleton and cached in PyTuple's freelist. */
2912 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002913
Tim Peters9edb1682013-09-03 11:49:31 -05002914 /* Scan indices right-to-left until finding one that is not
2915 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002916 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2917 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002918
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002919 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002920 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002921 if (i < 0)
2922 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002923
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002924 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002925 maximum. Then set all to the right to the same value. */
2926 index = indices[i] + 1;
2927 assert(index < n);
2928 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002929 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002930 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002931 Py_INCREF(elem);
2932 oldelem = PyTuple_GET_ITEM(result, i);
2933 PyTuple_SET_ITEM(result, i, elem);
2934 Py_DECREF(oldelem);
2935 }
2936 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002937
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002938 Py_INCREF(result);
2939 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002940
2941empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002942 co->stopped = 1;
2943 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002944}
2945
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002946static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302947cwr_reduce(cwrobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002948{
2949 if (lz->result == NULL) {
2950 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2951 } else if (lz->stopped) {
2952 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2953 } else {
2954 PyObject *indices;
2955 Py_ssize_t i;
2956
2957 /* we must pickle the indices and use them for setstate */
2958 indices = PyTuple_New(lz->r);
2959 if (!indices)
2960 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002961 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002962 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2963 if (!index) {
2964 Py_DECREF(indices);
2965 return NULL;
2966 }
2967 PyTuple_SET_ITEM(indices, i, index);
2968 }
2969
2970 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2971 }
2972}
2973
2974static PyObject *
2975cwr_setstate(cwrobject *lz, PyObject *state)
2976{
2977 PyObject *result;
2978 Py_ssize_t n, i;
2979
2980 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2981 {
2982 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2983 return NULL;
2984 }
2985
2986 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002987 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002988 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2989 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002990
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002991 if (index < 0 && PyErr_Occurred())
2992 return NULL; /* not an integer */
2993 /* clamp the index */
2994 if (index < 0)
2995 index = 0;
2996 else if (index > n-1)
2997 index = n-1;
2998 lz->indices[i] = index;
2999 }
3000 result = PyTuple_New(lz->r);
3001 if (result == NULL)
3002 return NULL;
3003 for (i=0; i<lz->r; i++) {
3004 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
3005 Py_INCREF(element);
3006 PyTuple_SET_ITEM(result, i, element);
3007 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003008 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003009 Py_RETURN_NONE;
3010}
3011
3012static PyMethodDef cwr_methods[] = {
3013 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
3014 reduce_doc},
3015 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
3016 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003017 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
3018 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003019 {NULL, NULL} /* sentinel */
3020};
3021
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003022static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003023 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003024 "itertools.combinations_with_replacement", /* tp_name */
3025 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003026 0, /* tp_itemsize */
3027 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003028 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003029 0, /* tp_print */
3030 0, /* tp_getattr */
3031 0, /* tp_setattr */
3032 0, /* tp_reserved */
3033 0, /* tp_repr */
3034 0, /* tp_as_number */
3035 0, /* tp_as_sequence */
3036 0, /* tp_as_mapping */
3037 0, /* tp_hash */
3038 0, /* tp_call */
3039 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003040 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003041 0, /* tp_setattro */
3042 0, /* tp_as_buffer */
3043 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003044 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003045 itertools_combinations_with_replacement__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003046 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003047 0, /* tp_clear */
3048 0, /* tp_richcompare */
3049 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003050 PyObject_SelfIter, /* tp_iter */
3051 (iternextfunc)cwr_next, /* tp_iternext */
3052 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003053 0, /* tp_members */
3054 0, /* tp_getset */
3055 0, /* tp_base */
3056 0, /* tp_dict */
3057 0, /* tp_descr_get */
3058 0, /* tp_descr_set */
3059 0, /* tp_dictoffset */
3060 0, /* tp_init */
3061 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003062 itertools_combinations_with_replacement, /* tp_new */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003063 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003064};
3065
3066
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003067/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003068
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003069def permutations(iterable, r=None):
3070 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
3071 pool = tuple(iterable)
3072 n = len(pool)
3073 r = n if r is None else r
3074 indices = range(n)
3075 cycles = range(n-r+1, n+1)[::-1]
3076 yield tuple(pool[i] for i in indices[:r])
3077 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003078 for i in reversed(range(r)):
3079 cycles[i] -= 1
3080 if cycles[i] == 0:
3081 indices[i:] = indices[i+1:] + indices[i:i+1]
3082 cycles[i] = n - i
3083 else:
3084 j = cycles[i]
3085 indices[i], indices[-j] = indices[-j], indices[i]
3086 yield tuple(pool[i] for i in indices[:r])
3087 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003088 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003089 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003090*/
3091
3092typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003093 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003094 PyObject *pool; /* input converted to a tuple */
3095 Py_ssize_t *indices; /* one index per element in the pool */
3096 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3097 PyObject *result; /* most recently returned result tuple */
3098 Py_ssize_t r; /* size of result tuple */
3099 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003100} permutationsobject;
3101
3102static PyTypeObject permutations_type;
3103
Tal Einatc4bccd32018-09-12 00:49:13 +03003104/*[clinic input]
3105@classmethod
3106itertools.permutations.__new__
3107 iterable: object
3108 r as robj: object = None
3109Return successive r-length permutations of elements in the iterable.
3110
3111permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
3112[clinic start generated code]*/
3113
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003114static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003115itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
3116 PyObject *robj)
3117/*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003118{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003119 permutationsobject *po;
3120 Py_ssize_t n;
3121 Py_ssize_t r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003122 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003123 Py_ssize_t *indices = NULL;
3124 Py_ssize_t *cycles = NULL;
3125 Py_ssize_t i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003126
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003127 pool = PySequence_Tuple(iterable);
3128 if (pool == NULL)
3129 goto error;
3130 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003131
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003132 r = n;
3133 if (robj != Py_None) {
3134 if (!PyLong_Check(robj)) {
3135 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3136 goto error;
3137 }
3138 r = PyLong_AsSsize_t(robj);
3139 if (r == -1 && PyErr_Occurred())
3140 goto error;
3141 }
3142 if (r < 0) {
3143 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3144 goto error;
3145 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003146
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003147 indices = PyMem_New(Py_ssize_t, n);
3148 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003149 if (indices == NULL || cycles == NULL) {
3150 PyErr_NoMemory();
3151 goto error;
3152 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003153
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003154 for (i=0 ; i<n ; i++)
3155 indices[i] = i;
3156 for (i=0 ; i<r ; i++)
3157 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003158
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003159 /* create permutationsobject structure */
3160 po = (permutationsobject *)type->tp_alloc(type, 0);
3161 if (po == NULL)
3162 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003163
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003164 po->pool = pool;
3165 po->indices = indices;
3166 po->cycles = cycles;
3167 po->result = NULL;
3168 po->r = r;
3169 po->stopped = r > n ? 1 : 0;
3170
3171 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003172
3173error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003174 if (indices != NULL)
3175 PyMem_Free(indices);
3176 if (cycles != NULL)
3177 PyMem_Free(cycles);
3178 Py_XDECREF(pool);
3179 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003180}
3181
3182static void
3183permutations_dealloc(permutationsobject *po)
3184{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003185 PyObject_GC_UnTrack(po);
3186 Py_XDECREF(po->pool);
3187 Py_XDECREF(po->result);
3188 PyMem_Free(po->indices);
3189 PyMem_Free(po->cycles);
3190 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003191}
3192
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003193static PyObject *
3194permutations_sizeof(permutationsobject *po, void *unused)
3195{
3196 Py_ssize_t res;
3197
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003198 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003199 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3200 res += po->r * sizeof(Py_ssize_t);
3201 return PyLong_FromSsize_t(res);
3202}
3203
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003204static int
3205permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3206{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003207 Py_VISIT(po->pool);
3208 Py_VISIT(po->result);
3209 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003210}
3211
3212static PyObject *
3213permutations_next(permutationsobject *po)
3214{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003215 PyObject *elem;
3216 PyObject *oldelem;
3217 PyObject *pool = po->pool;
3218 Py_ssize_t *indices = po->indices;
3219 Py_ssize_t *cycles = po->cycles;
3220 PyObject *result = po->result;
3221 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3222 Py_ssize_t r = po->r;
3223 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003224
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003225 if (po->stopped)
3226 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003227
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003228 if (result == NULL) {
3229 /* On the first pass, initialize result tuple using the indices */
3230 result = PyTuple_New(r);
3231 if (result == NULL)
3232 goto empty;
3233 po->result = result;
3234 for (i=0; i<r ; i++) {
3235 index = indices[i];
3236 elem = PyTuple_GET_ITEM(pool, index);
3237 Py_INCREF(elem);
3238 PyTuple_SET_ITEM(result, i, elem);
3239 }
3240 } else {
3241 if (n == 0)
3242 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003243
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003244 /* Copy the previous result tuple or re-use it if available */
3245 if (Py_REFCNT(result) > 1) {
3246 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05003247 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003248 if (result == NULL)
3249 goto empty;
3250 po->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003251 Py_DECREF(old_result);
3252 }
3253 /* Now, we've got the only copy so we can update it in-place */
3254 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003255
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003256 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3257 for (i=r-1 ; i>=0 ; i--) {
3258 cycles[i] -= 1;
3259 if (cycles[i] == 0) {
3260 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3261 index = indices[i];
3262 for (j=i ; j<n-1 ; j++)
3263 indices[j] = indices[j+1];
3264 indices[n-1] = index;
3265 cycles[i] = n - i;
3266 } else {
3267 j = cycles[i];
3268 index = indices[i];
3269 indices[i] = indices[n-j];
3270 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003271
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003272 for (k=i; k<r ; k++) {
3273 /* start with i, the leftmost element that changed */
3274 /* yield tuple(pool[k] for k in indices[:r]) */
3275 index = indices[k];
3276 elem = PyTuple_GET_ITEM(pool, index);
3277 Py_INCREF(elem);
3278 oldelem = PyTuple_GET_ITEM(result, k);
3279 PyTuple_SET_ITEM(result, k, elem);
3280 Py_DECREF(oldelem);
3281 }
3282 break;
3283 }
3284 }
3285 /* If i is negative, then the cycles have all
3286 rolled-over and we're done. */
3287 if (i < 0)
3288 goto empty;
3289 }
3290 Py_INCREF(result);
3291 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003292
3293empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003294 po->stopped = 1;
3295 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003296}
3297
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003298static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303299permutations_reduce(permutationsobject *po, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003300{
3301 if (po->result == NULL) {
3302 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3303 } else if (po->stopped) {
3304 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3305 } else {
3306 PyObject *indices=NULL, *cycles=NULL;
3307 Py_ssize_t n, i;
3308
3309 /* we must pickle the indices and cycles and use them for setstate */
3310 n = PyTuple_GET_SIZE(po->pool);
3311 indices = PyTuple_New(n);
3312 if (indices == NULL)
3313 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003314 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003315 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3316 if (!index)
3317 goto err;
3318 PyTuple_SET_ITEM(indices, i, index);
3319 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003320
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003321 cycles = PyTuple_New(po->r);
3322 if (cycles == NULL)
3323 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003324 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003325 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3326 if (!index)
3327 goto err;
3328 PyTuple_SET_ITEM(cycles, i, index);
3329 }
3330 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3331 po->pool, po->r,
3332 indices, cycles);
3333 err:
3334 Py_XDECREF(indices);
3335 Py_XDECREF(cycles);
3336 return NULL;
3337 }
3338}
3339
3340static PyObject *
3341permutations_setstate(permutationsobject *po, PyObject *state)
3342{
3343 PyObject *indices, *cycles, *result;
3344 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003345
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003346 if (!PyTuple_Check(state)) {
3347 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3348 return NULL;
3349 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003350 if (!PyArg_ParseTuple(state, "O!O!",
3351 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003352 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003353 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003354 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003355
3356 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003357 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003358 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3359 return NULL;
3360 }
3361
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003362 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003363 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3364 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3365 if (index < 0 && PyErr_Occurred())
3366 return NULL; /* not an integer */
3367 /* clamp the index */
3368 if (index < 0)
3369 index = 0;
3370 else if (index > n-1)
3371 index = n-1;
3372 po->indices[i] = index;
3373 }
3374
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003375 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003376 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3377 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3378 if (index < 0 && PyErr_Occurred())
3379 return NULL; /* not an integer */
3380 if (index < 1)
3381 index = 1;
3382 else if (index > n-i)
3383 index = n-i;
3384 po->cycles[i] = index;
3385 }
3386 result = PyTuple_New(po->r);
3387 if (result == NULL)
3388 return NULL;
3389 for (i=0; i<po->r; i++) {
3390 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3391 Py_INCREF(element);
3392 PyTuple_SET_ITEM(result, i, element);
3393 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003394 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003395 Py_RETURN_NONE;
3396}
3397
3398static PyMethodDef permuations_methods[] = {
3399 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3400 reduce_doc},
3401 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3402 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003403 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3404 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003405 {NULL, NULL} /* sentinel */
3406};
3407
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003408static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003409 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003410 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003411 sizeof(permutationsobject), /* tp_basicsize */
3412 0, /* tp_itemsize */
3413 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003414 (destructor)permutations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003415 0, /* tp_print */
3416 0, /* tp_getattr */
3417 0, /* tp_setattr */
3418 0, /* tp_reserved */
3419 0, /* tp_repr */
3420 0, /* tp_as_number */
3421 0, /* tp_as_sequence */
3422 0, /* tp_as_mapping */
3423 0, /* tp_hash */
3424 0, /* tp_call */
3425 0, /* tp_str */
3426 PyObject_GenericGetAttr, /* tp_getattro */
3427 0, /* tp_setattro */
3428 0, /* tp_as_buffer */
3429 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3430 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003431 itertools_permutations__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003432 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003433 0, /* tp_clear */
3434 0, /* tp_richcompare */
3435 0, /* tp_weaklistoffset */
3436 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003437 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003438 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003439 0, /* tp_members */
3440 0, /* tp_getset */
3441 0, /* tp_base */
3442 0, /* tp_dict */
3443 0, /* tp_descr_get */
3444 0, /* tp_descr_set */
3445 0, /* tp_dictoffset */
3446 0, /* tp_init */
3447 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003448 itertools_permutations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003449 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003450};
3451
Tal Einatc4bccd32018-09-12 00:49:13 +03003452
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003453/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003454
3455typedef struct {
3456 PyObject_HEAD
3457 PyObject *total;
3458 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003459 PyObject *binop;
Lisa Roach9718b592018-09-23 17:34:59 -07003460 PyObject *initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003461} accumulateobject;
3462
3463static PyTypeObject accumulate_type;
3464
Tal Einatc4bccd32018-09-12 00:49:13 +03003465/*[clinic input]
3466@classmethod
3467itertools.accumulate.__new__
3468 iterable: object
3469 func as binop: object = None
Lisa Roach9718b592018-09-23 17:34:59 -07003470 *
3471 initial: object = None
Tal Einatc4bccd32018-09-12 00:49:13 +03003472Return series of accumulated sums (or other binary function results).
3473[clinic start generated code]*/
3474
Raymond Hettinger482ba772010-12-01 22:48:00 +00003475static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003476itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
Lisa Roach9718b592018-09-23 17:34:59 -07003477 PyObject *binop, PyObject *initial)
3478/*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003479{
Raymond Hettinger482ba772010-12-01 22:48:00 +00003480 PyObject *it;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003481 accumulateobject *lz;
3482
Raymond Hettinger482ba772010-12-01 22:48:00 +00003483 /* Get iterator. */
3484 it = PyObject_GetIter(iterable);
3485 if (it == NULL)
3486 return NULL;
3487
Raymond Hettinger482ba772010-12-01 22:48:00 +00003488 /* create accumulateobject structure */
3489 lz = (accumulateobject *)type->tp_alloc(type, 0);
3490 if (lz == NULL) {
3491 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003492 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003493 }
3494
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003495 if (binop != Py_None) {
3496 Py_XINCREF(binop);
3497 lz->binop = binop;
3498 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003499 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003500 lz->it = it;
Lisa Roach9718b592018-09-23 17:34:59 -07003501 Py_XINCREF(initial);
3502 lz->initial = initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003503 return (PyObject *)lz;
3504}
3505
3506static void
3507accumulate_dealloc(accumulateobject *lz)
3508{
3509 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003510 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003511 Py_XDECREF(lz->total);
3512 Py_XDECREF(lz->it);
Lisa Roach9718b592018-09-23 17:34:59 -07003513 Py_XDECREF(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003514 Py_TYPE(lz)->tp_free(lz);
3515}
3516
3517static int
3518accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3519{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003520 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003521 Py_VISIT(lz->it);
3522 Py_VISIT(lz->total);
Lisa Roach9718b592018-09-23 17:34:59 -07003523 Py_VISIT(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003524 return 0;
3525}
3526
3527static PyObject *
3528accumulate_next(accumulateobject *lz)
3529{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003530 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003531
Lisa Roach9718b592018-09-23 17:34:59 -07003532 if (lz->initial != Py_None) {
3533 lz->total = lz->initial;
3534 Py_INCREF(Py_None);
3535 lz->initial = Py_None;
3536 Py_INCREF(lz->total);
3537 return lz->total;
3538 }
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003539 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003540 if (val == NULL)
3541 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003542
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003543 if (lz->total == NULL) {
3544 Py_INCREF(val);
3545 lz->total = val;
3546 return lz->total;
3547 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003548
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003549 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003550 newtotal = PyNumber_Add(lz->total, val);
3551 else
3552 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003553 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003554 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003555 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003556
Raymond Hettinger482ba772010-12-01 22:48:00 +00003557 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003558 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003559 return newtotal;
3560}
3561
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003562static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303563accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003564{
Lisa Roach9718b592018-09-23 17:34:59 -07003565 if (lz->initial != Py_None) {
3566 PyObject *it;
3567
3568 assert(lz->total == NULL);
3569 if (PyType_Ready(&chain_type) < 0)
3570 return NULL;
3571 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3572 lz->initial, lz->it);
3573 if (it == NULL)
3574 return NULL;
3575 return Py_BuildValue("O(NO)O", Py_TYPE(lz),
3576 it, lz->binop?lz->binop:Py_None, Py_None);
3577 }
Serhiy Storchakad5516252016-03-06 14:00:45 +02003578 if (lz->total == Py_None) {
3579 PyObject *it;
3580
3581 if (PyType_Ready(&chain_type) < 0)
3582 return NULL;
3583 if (PyType_Ready(&islice_type) < 0)
3584 return NULL;
3585 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3586 lz->total, lz->it);
3587 if (it == NULL)
3588 return NULL;
3589 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3590 it, lz->binop ? lz->binop : Py_None);
3591 if (it == NULL)
3592 return NULL;
3593 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3594 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003595 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3596 lz->it, lz->binop?lz->binop:Py_None,
3597 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003598}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003599
3600static PyObject *
3601accumulate_setstate(accumulateobject *lz, PyObject *state)
3602{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003603 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003604 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003605 Py_RETURN_NONE;
3606}
3607
3608static PyMethodDef accumulate_methods[] = {
3609 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3610 reduce_doc},
3611 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3612 setstate_doc},
3613 {NULL, NULL} /* sentinel */
3614};
3615
Raymond Hettinger482ba772010-12-01 22:48:00 +00003616static PyTypeObject accumulate_type = {
3617 PyVarObject_HEAD_INIT(NULL, 0)
3618 "itertools.accumulate", /* tp_name */
3619 sizeof(accumulateobject), /* tp_basicsize */
3620 0, /* tp_itemsize */
3621 /* methods */
3622 (destructor)accumulate_dealloc, /* tp_dealloc */
3623 0, /* tp_print */
3624 0, /* tp_getattr */
3625 0, /* tp_setattr */
3626 0, /* tp_reserved */
3627 0, /* tp_repr */
3628 0, /* tp_as_number */
3629 0, /* tp_as_sequence */
3630 0, /* tp_as_mapping */
3631 0, /* tp_hash */
3632 0, /* tp_call */
3633 0, /* tp_str */
3634 PyObject_GenericGetAttr, /* tp_getattro */
3635 0, /* tp_setattro */
3636 0, /* tp_as_buffer */
3637 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3638 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003639 itertools_accumulate__doc__, /* tp_doc */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003640 (traverseproc)accumulate_traverse, /* tp_traverse */
3641 0, /* tp_clear */
3642 0, /* tp_richcompare */
3643 0, /* tp_weaklistoffset */
3644 PyObject_SelfIter, /* tp_iter */
3645 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003646 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003647 0, /* tp_members */
3648 0, /* tp_getset */
3649 0, /* tp_base */
3650 0, /* tp_dict */
3651 0, /* tp_descr_get */
3652 0, /* tp_descr_set */
3653 0, /* tp_dictoffset */
3654 0, /* tp_init */
3655 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003656 itertools_accumulate, /* tp_new */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003657 PyObject_GC_Del, /* tp_free */
3658};
3659
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003660
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003661/* compress object ************************************************************/
3662
3663/* Equivalent to:
3664
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003665 def compress(data, selectors):
3666 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3667 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003668*/
3669
3670typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003671 PyObject_HEAD
3672 PyObject *data;
3673 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003674} compressobject;
3675
3676static PyTypeObject compress_type;
3677
Tal Einatc4bccd32018-09-12 00:49:13 +03003678/*[clinic input]
3679@classmethod
3680itertools.compress.__new__
3681 data as seq1: object
3682 selectors as seq2: object
3683Return data elements corresponding to true selector elements.
3684
3685Forms a shorter iterator from selected data elements using the selectors to
3686choose the data elements.
3687[clinic start generated code]*/
3688
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003689static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003690itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
3691/*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003692{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003693 PyObject *data=NULL, *selectors=NULL;
3694 compressobject *lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003695
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003696 data = PyObject_GetIter(seq1);
3697 if (data == NULL)
3698 goto fail;
3699 selectors = PyObject_GetIter(seq2);
3700 if (selectors == NULL)
3701 goto fail;
3702
3703 /* create compressobject structure */
3704 lz = (compressobject *)type->tp_alloc(type, 0);
3705 if (lz == NULL)
3706 goto fail;
3707 lz->data = data;
3708 lz->selectors = selectors;
3709 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003710
3711fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003712 Py_XDECREF(data);
3713 Py_XDECREF(selectors);
3714 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003715}
3716
3717static void
3718compress_dealloc(compressobject *lz)
3719{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003720 PyObject_GC_UnTrack(lz);
3721 Py_XDECREF(lz->data);
3722 Py_XDECREF(lz->selectors);
3723 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003724}
3725
3726static int
3727compress_traverse(compressobject *lz, visitproc visit, void *arg)
3728{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003729 Py_VISIT(lz->data);
3730 Py_VISIT(lz->selectors);
3731 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003732}
3733
3734static PyObject *
3735compress_next(compressobject *lz)
3736{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003737 PyObject *data = lz->data, *selectors = lz->selectors;
3738 PyObject *datum, *selector;
3739 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3740 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3741 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003742
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003743 while (1) {
3744 /* Steps: get datum, get selector, evaluate selector.
3745 Order is important (to match the pure python version
3746 in terms of which input gets a chance to raise an
3747 exception first).
3748 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003749
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003750 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003751 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003752 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003753
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003754 selector = selectornext(selectors);
3755 if (selector == NULL) {
3756 Py_DECREF(datum);
3757 return NULL;
3758 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003759
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003760 ok = PyObject_IsTrue(selector);
3761 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003762 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003763 return datum;
3764 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003765 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003766 return NULL;
3767 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003768}
3769
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003770static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303771compress_reduce(compressobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003772{
3773 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3774 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003775}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003776
3777static PyMethodDef compress_methods[] = {
3778 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3779 reduce_doc},
3780 {NULL, NULL} /* sentinel */
3781};
3782
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003783static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003784 PyVarObject_HEAD_INIT(NULL, 0)
3785 "itertools.compress", /* tp_name */
3786 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003787 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003788 /* methods */
3789 (destructor)compress_dealloc, /* tp_dealloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003790 0, /* tp_print */
3791 0, /* tp_getattr */
3792 0, /* tp_setattr */
3793 0, /* tp_reserved */
3794 0, /* tp_repr */
3795 0, /* tp_as_number */
3796 0, /* tp_as_sequence */
3797 0, /* tp_as_mapping */
3798 0, /* tp_hash */
3799 0, /* tp_call */
3800 0, /* tp_str */
3801 PyObject_GenericGetAttr, /* tp_getattro */
3802 0, /* tp_setattro */
3803 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003804 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003805 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003806 itertools_compress__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003807 (traverseproc)compress_traverse, /* tp_traverse */
3808 0, /* tp_clear */
3809 0, /* tp_richcompare */
3810 0, /* tp_weaklistoffset */
3811 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003812 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003813 compress_methods, /* tp_methods */
3814 0, /* tp_members */
3815 0, /* tp_getset */
3816 0, /* tp_base */
3817 0, /* tp_dict */
3818 0, /* tp_descr_get */
3819 0, /* tp_descr_set */
3820 0, /* tp_dictoffset */
3821 0, /* tp_init */
3822 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003823 itertools_compress, /* tp_new */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003824 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003825};
3826
3827
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003828/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003829
3830typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003831 PyObject_HEAD
3832 PyObject *func;
3833 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003834} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003835
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003836static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003837
Tal Einatc4bccd32018-09-12 00:49:13 +03003838/*[clinic input]
3839@classmethod
3840itertools.filterfalse.__new__
3841 function as func: object
3842 iterable as seq: object
3843 /
3844Return those items of iterable for which function(item) is false.
3845
3846If function is None, return the items that are false.
3847[clinic start generated code]*/
3848
Raymond Hettinger60eca932003-02-09 06:40:58 +00003849static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003850itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
3851/*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003852{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003853 PyObject *it;
3854 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003855
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003856 /* Get iterator. */
3857 it = PyObject_GetIter(seq);
3858 if (it == NULL)
3859 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003860
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003861 /* create filterfalseobject structure */
3862 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3863 if (lz == NULL) {
3864 Py_DECREF(it);
3865 return NULL;
3866 }
3867 Py_INCREF(func);
3868 lz->func = func;
3869 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003870
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003871 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003872}
3873
3874static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003875filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003876{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003877 PyObject_GC_UnTrack(lz);
3878 Py_XDECREF(lz->func);
3879 Py_XDECREF(lz->it);
3880 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003881}
3882
3883static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003884filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003885{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003886 Py_VISIT(lz->it);
3887 Py_VISIT(lz->func);
3888 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003889}
3890
3891static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003892filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003893{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003894 PyObject *item;
3895 PyObject *it = lz->it;
3896 long ok;
3897 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003898
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003899 iternext = *Py_TYPE(it)->tp_iternext;
3900 for (;;) {
3901 item = iternext(it);
3902 if (item == NULL)
3903 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003904
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003905 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3906 ok = PyObject_IsTrue(item);
3907 } else {
3908 PyObject *good;
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01003909 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003910 if (good == NULL) {
3911 Py_DECREF(item);
3912 return NULL;
3913 }
3914 ok = PyObject_IsTrue(good);
3915 Py_DECREF(good);
3916 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003917 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003918 return item;
3919 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003920 if (ok < 0)
3921 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003922 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003923}
3924
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003925static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303926filterfalse_reduce(filterfalseobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003927{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003928 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3929}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003930
3931static PyMethodDef filterfalse_methods[] = {
3932 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3933 reduce_doc},
3934 {NULL, NULL} /* sentinel */
3935};
3936
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003937static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003938 PyVarObject_HEAD_INIT(NULL, 0)
3939 "itertools.filterfalse", /* tp_name */
3940 sizeof(filterfalseobject), /* tp_basicsize */
3941 0, /* tp_itemsize */
3942 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003943 (destructor)filterfalse_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003944 0, /* tp_print */
3945 0, /* tp_getattr */
3946 0, /* tp_setattr */
3947 0, /* tp_reserved */
3948 0, /* tp_repr */
3949 0, /* tp_as_number */
3950 0, /* tp_as_sequence */
3951 0, /* tp_as_mapping */
3952 0, /* tp_hash */
3953 0, /* tp_call */
3954 0, /* tp_str */
3955 PyObject_GenericGetAttr, /* tp_getattro */
3956 0, /* tp_setattro */
3957 0, /* tp_as_buffer */
3958 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3959 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003960 itertools_filterfalse__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003961 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003962 0, /* tp_clear */
3963 0, /* tp_richcompare */
3964 0, /* tp_weaklistoffset */
3965 PyObject_SelfIter, /* tp_iter */
3966 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003967 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003968 0, /* tp_members */
3969 0, /* tp_getset */
3970 0, /* tp_base */
3971 0, /* tp_dict */
3972 0, /* tp_descr_get */
3973 0, /* tp_descr_set */
3974 0, /* tp_dictoffset */
3975 0, /* tp_init */
3976 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003977 itertools_filterfalse, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003978 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003979};
3980
3981
3982/* count object ************************************************************/
3983
3984typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003985 PyObject_HEAD
3986 Py_ssize_t cnt;
3987 PyObject *long_cnt;
3988 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003989} countobject;
3990
Raymond Hettinger30729212009-02-12 06:28:27 +00003991/* Counting logic and invariants:
3992
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003993fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003994
Serhiy Storchaka483405b2015-02-17 10:14:30 +02003995 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003996 Advances with: cnt += 1
3997 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003998
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003999slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00004000
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004001 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
4002 All counting is done with python objects (no overflows or underflows).
4003 Advances with: long_cnt += long_step
4004 Step may be zero -- effectively a slow version of repeat(cnt).
4005 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00004006*/
4007
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004008static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004009
Tal Einatc4bccd32018-09-12 00:49:13 +03004010/*[clinic input]
4011@classmethod
4012itertools.count.__new__
4013 start as long_cnt: object(c_default="NULL") = 0
4014 step as long_step: object(c_default="NULL") = 1
4015Return a count object whose .__next__() method returns consecutive values.
4016
4017Equivalent to:
4018 def count(firstval=0, step=1):
4019 x = firstval
4020 while 1:
4021 yield x
4022 x += step
4023[clinic start generated code]*/
4024
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004025static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03004026itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
4027 PyObject *long_step)
4028/*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004029{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004030 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004031 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004032 Py_ssize_t cnt = 0;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00004033 long step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004034
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004035 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
4036 (long_step != NULL && !PyNumber_Check(long_step))) {
4037 PyErr_SetString(PyExc_TypeError, "a number is required");
4038 return NULL;
4039 }
Raymond Hettinger30729212009-02-12 06:28:27 +00004040
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004041 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
4042 (long_step == NULL || PyLong_Check(long_step));
4043
4044 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004045 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004046 if (fast_mode) {
4047 assert(PyLong_Check(long_cnt));
4048 cnt = PyLong_AsSsize_t(long_cnt);
4049 if (cnt == -1 && PyErr_Occurred()) {
4050 PyErr_Clear();
4051 fast_mode = 0;
4052 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004053 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004054 } else {
4055 cnt = 0;
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004056 long_cnt = _PyLong_Zero;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004057 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004058 Py_INCREF(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004059
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004060 /* If not specified, step defaults to 1 */
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004061 if (long_step == NULL)
4062 long_step = _PyLong_One;
4063 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004064
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004065 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004066
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004067 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004068 if (fast_mode) {
4069 assert(PyLong_Check(long_step));
4070 step = PyLong_AsLong(long_step);
4071 if (step != 1) {
4072 fast_mode = 0;
4073 if (step == -1 && PyErr_Occurred())
4074 PyErr_Clear();
4075 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004076 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004077
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004078 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004079 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004080 else
4081 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004082
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004083 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4084 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4085 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004086 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004087
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004088 /* create countobject structure */
4089 lz = (countobject *)type->tp_alloc(type, 0);
4090 if (lz == NULL) {
4091 Py_XDECREF(long_cnt);
Zackery Spytz0523c392019-03-26 00:05:29 -06004092 Py_DECREF(long_step);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004093 return NULL;
4094 }
4095 lz->cnt = cnt;
4096 lz->long_cnt = long_cnt;
4097 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004098
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004099 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004100}
4101
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004102static void
4103count_dealloc(countobject *lz)
4104{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004105 PyObject_GC_UnTrack(lz);
4106 Py_XDECREF(lz->long_cnt);
4107 Py_XDECREF(lz->long_step);
4108 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004109}
4110
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004111static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004112count_traverse(countobject *lz, visitproc visit, void *arg)
4113{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004114 Py_VISIT(lz->long_cnt);
4115 Py_VISIT(lz->long_step);
4116 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004117}
4118
4119static PyObject *
4120count_nextlong(countobject *lz)
4121{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004122 PyObject *long_cnt;
4123 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004124
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004125 long_cnt = lz->long_cnt;
4126 if (long_cnt == NULL) {
4127 /* Switch to slow_mode */
4128 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4129 if (long_cnt == NULL)
4130 return NULL;
4131 }
4132 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004133
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004134 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4135 if (stepped_up == NULL)
4136 return NULL;
4137 lz->long_cnt = stepped_up;
4138 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004139}
4140
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004141static PyObject *
4142count_next(countobject *lz)
4143{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004144 if (lz->cnt == PY_SSIZE_T_MAX)
4145 return count_nextlong(lz);
4146 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004147}
4148
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004149static PyObject *
4150count_repr(countobject *lz)
4151{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004152 if (lz->cnt != PY_SSIZE_T_MAX)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004153 return PyUnicode_FromFormat("%s(%zd)",
4154 _PyType_Name(Py_TYPE(lz)), lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004155
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004156 if (PyLong_Check(lz->long_step)) {
4157 long step = PyLong_AsLong(lz->long_step);
4158 if (step == -1 && PyErr_Occurred()) {
4159 PyErr_Clear();
4160 }
4161 if (step == 1) {
4162 /* Don't display step when it is an integer equal to 1 */
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004163 return PyUnicode_FromFormat("%s(%R)",
4164 _PyType_Name(Py_TYPE(lz)),
4165 lz->long_cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004166 }
4167 }
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004168 return PyUnicode_FromFormat("%s(%R, %R)",
4169 _PyType_Name(Py_TYPE(lz)),
4170 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004171}
4172
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004173static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304174count_reduce(countobject *lz, PyObject *Py_UNUSED(ignored))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004175{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004176 if (lz->cnt == PY_SSIZE_T_MAX)
4177 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4178 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004179}
4180
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004181static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004182 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004183 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004184 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004185};
4186
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004187static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004188 PyVarObject_HEAD_INIT(NULL, 0)
4189 "itertools.count", /* tp_name */
4190 sizeof(countobject), /* tp_basicsize */
4191 0, /* tp_itemsize */
4192 /* methods */
4193 (destructor)count_dealloc, /* tp_dealloc */
4194 0, /* tp_print */
4195 0, /* tp_getattr */
4196 0, /* tp_setattr */
4197 0, /* tp_reserved */
4198 (reprfunc)count_repr, /* tp_repr */
4199 0, /* tp_as_number */
4200 0, /* tp_as_sequence */
4201 0, /* tp_as_mapping */
4202 0, /* tp_hash */
4203 0, /* tp_call */
4204 0, /* tp_str */
4205 PyObject_GenericGetAttr, /* tp_getattro */
4206 0, /* tp_setattro */
4207 0, /* tp_as_buffer */
4208 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004209 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03004210 itertools_count__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004211 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004212 0, /* tp_clear */
4213 0, /* tp_richcompare */
4214 0, /* tp_weaklistoffset */
4215 PyObject_SelfIter, /* tp_iter */
4216 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004217 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004218 0, /* tp_members */
4219 0, /* tp_getset */
4220 0, /* tp_base */
4221 0, /* tp_dict */
4222 0, /* tp_descr_get */
4223 0, /* tp_descr_set */
4224 0, /* tp_dictoffset */
4225 0, /* tp_init */
4226 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03004227 itertools_count, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004228 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004229};
4230
4231
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004232/* repeat object ************************************************************/
4233
4234typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004235 PyObject_HEAD
4236 PyObject *element;
4237 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004238} repeatobject;
4239
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004240static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004241
4242static PyObject *
4243repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4244{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004245 repeatobject *ro;
4246 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004247 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004248 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004249
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004250 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4251 &element, &cnt))
4252 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004253
Raymond Hettinger97d35552014-06-24 21:36:58 -07004254 if (kwds != NULL)
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004255 n_kwds = PyDict_GET_SIZE(kwds);
Raymond Hettinger97d35552014-06-24 21:36:58 -07004256 /* Does user supply times argument? */
4257 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004258 cnt = 0;
4259
4260 ro = (repeatobject *)type->tp_alloc(type, 0);
4261 if (ro == NULL)
4262 return NULL;
4263 Py_INCREF(element);
4264 ro->element = element;
4265 ro->cnt = cnt;
4266 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004267}
4268
4269static void
4270repeat_dealloc(repeatobject *ro)
4271{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004272 PyObject_GC_UnTrack(ro);
4273 Py_XDECREF(ro->element);
4274 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004275}
4276
4277static int
4278repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4279{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004280 Py_VISIT(ro->element);
4281 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004282}
4283
4284static PyObject *
4285repeat_next(repeatobject *ro)
4286{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004287 if (ro->cnt == 0)
4288 return NULL;
4289 if (ro->cnt > 0)
4290 ro->cnt--;
4291 Py_INCREF(ro->element);
4292 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004293}
4294
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004295static PyObject *
4296repeat_repr(repeatobject *ro)
4297{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004298 if (ro->cnt == -1)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004299 return PyUnicode_FromFormat("%s(%R)",
4300 _PyType_Name(Py_TYPE(ro)), ro->element);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004301 else
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004302 return PyUnicode_FromFormat("%s(%R, %zd)",
4303 _PyType_Name(Py_TYPE(ro)), ro->element,
4304 ro->cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004305}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004306
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004307static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304308repeat_len(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004309{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004310 if (ro->cnt == -1) {
4311 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4312 return NULL;
4313 }
4314 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004315}
4316
Armin Rigof5b3e362006-02-11 21:32:43 +00004317PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004318
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004319static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304320repeat_reduce(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004321{
4322 /* unpickle this so that a new repeat iterator is constructed with an
4323 * object, then call __setstate__ on it to set cnt
4324 */
4325 if (ro->cnt >= 0)
4326 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4327 else
4328 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4329}
4330
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004331static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004332 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004333 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004334 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004335};
4336
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004337PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004338"repeat(object [,times]) -> create an iterator which returns the object\n\
4339for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004340endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004341
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004342static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004343 PyVarObject_HEAD_INIT(NULL, 0)
4344 "itertools.repeat", /* tp_name */
4345 sizeof(repeatobject), /* tp_basicsize */
4346 0, /* tp_itemsize */
4347 /* methods */
4348 (destructor)repeat_dealloc, /* tp_dealloc */
4349 0, /* tp_print */
4350 0, /* tp_getattr */
4351 0, /* tp_setattr */
4352 0, /* tp_reserved */
4353 (reprfunc)repeat_repr, /* tp_repr */
4354 0, /* tp_as_number */
4355 0, /* tp_as_sequence */
4356 0, /* tp_as_mapping */
4357 0, /* tp_hash */
4358 0, /* tp_call */
4359 0, /* tp_str */
4360 PyObject_GenericGetAttr, /* tp_getattro */
4361 0, /* tp_setattro */
4362 0, /* tp_as_buffer */
4363 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4364 Py_TPFLAGS_BASETYPE, /* tp_flags */
4365 repeat_doc, /* tp_doc */
4366 (traverseproc)repeat_traverse, /* tp_traverse */
4367 0, /* tp_clear */
4368 0, /* tp_richcompare */
4369 0, /* tp_weaklistoffset */
4370 PyObject_SelfIter, /* tp_iter */
4371 (iternextfunc)repeat_next, /* tp_iternext */
4372 repeat_methods, /* tp_methods */
4373 0, /* tp_members */
4374 0, /* tp_getset */
4375 0, /* tp_base */
4376 0, /* tp_dict */
4377 0, /* tp_descr_get */
4378 0, /* tp_descr_set */
4379 0, /* tp_dictoffset */
4380 0, /* tp_init */
4381 0, /* tp_alloc */
4382 repeat_new, /* tp_new */
4383 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004384};
4385
Tal Einatc4bccd32018-09-12 00:49:13 +03004386
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004387/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004388
4389typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004390 PyObject_HEAD
4391 Py_ssize_t tuplesize;
4392 Py_ssize_t numactive;
4393 PyObject *ittuple; /* tuple of iterators */
4394 PyObject *result;
4395 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004396} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004397
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004398static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004399
4400static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004401zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004402{
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004403 _Py_IDENTIFIER(fillvalue);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004404 ziplongestobject *lz;
4405 Py_ssize_t i;
4406 PyObject *ittuple; /* tuple of iterators */
4407 PyObject *result;
4408 PyObject *fillvalue = Py_None;
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004409 Py_ssize_t tuplesize;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004410
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004411 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004412 fillvalue = NULL;
4413 if (PyDict_GET_SIZE(kwds) == 1) {
4414 fillvalue = _PyDict_GetItemIdWithError(kwds, &PyId_fillvalue);
4415 }
4416 if (fillvalue == NULL) {
4417 if (!PyErr_Occurred()) {
4418 PyErr_SetString(PyExc_TypeError,
4419 "zip_longest() got an unexpected keyword argument");
4420 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004421 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004422 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004423 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004424
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004425 /* args must be a tuple */
4426 assert(PyTuple_Check(args));
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004427 tuplesize = PyTuple_GET_SIZE(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004428
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004429 /* obtain iterators */
4430 ittuple = PyTuple_New(tuplesize);
4431 if (ittuple == NULL)
4432 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004433 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004434 PyObject *item = PyTuple_GET_ITEM(args, i);
4435 PyObject *it = PyObject_GetIter(item);
4436 if (it == NULL) {
4437 if (PyErr_ExceptionMatches(PyExc_TypeError))
4438 PyErr_Format(PyExc_TypeError,
4439 "zip_longest argument #%zd must support iteration",
4440 i+1);
4441 Py_DECREF(ittuple);
4442 return NULL;
4443 }
4444 PyTuple_SET_ITEM(ittuple, i, it);
4445 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004446
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004447 /* create a result holder */
4448 result = PyTuple_New(tuplesize);
4449 if (result == NULL) {
4450 Py_DECREF(ittuple);
4451 return NULL;
4452 }
4453 for (i=0 ; i < tuplesize ; i++) {
4454 Py_INCREF(Py_None);
4455 PyTuple_SET_ITEM(result, i, Py_None);
4456 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004457
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004458 /* create ziplongestobject structure */
4459 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4460 if (lz == NULL) {
4461 Py_DECREF(ittuple);
4462 Py_DECREF(result);
4463 return NULL;
4464 }
4465 lz->ittuple = ittuple;
4466 lz->tuplesize = tuplesize;
4467 lz->numactive = tuplesize;
4468 lz->result = result;
4469 Py_INCREF(fillvalue);
4470 lz->fillvalue = fillvalue;
4471 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004472}
4473
4474static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004475zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004476{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004477 PyObject_GC_UnTrack(lz);
4478 Py_XDECREF(lz->ittuple);
4479 Py_XDECREF(lz->result);
4480 Py_XDECREF(lz->fillvalue);
4481 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004482}
4483
4484static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004485zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004486{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004487 Py_VISIT(lz->ittuple);
4488 Py_VISIT(lz->result);
4489 Py_VISIT(lz->fillvalue);
4490 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004491}
4492
4493static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004494zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004495{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004496 Py_ssize_t i;
4497 Py_ssize_t tuplesize = lz->tuplesize;
4498 PyObject *result = lz->result;
4499 PyObject *it;
4500 PyObject *item;
4501 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004502
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004503 if (tuplesize == 0)
4504 return NULL;
4505 if (lz->numactive == 0)
4506 return NULL;
4507 if (Py_REFCNT(result) == 1) {
4508 Py_INCREF(result);
4509 for (i=0 ; i < tuplesize ; i++) {
4510 it = PyTuple_GET_ITEM(lz->ittuple, i);
4511 if (it == NULL) {
4512 Py_INCREF(lz->fillvalue);
4513 item = lz->fillvalue;
4514 } else {
4515 item = PyIter_Next(it);
4516 if (item == NULL) {
4517 lz->numactive -= 1;
4518 if (lz->numactive == 0 || PyErr_Occurred()) {
4519 lz->numactive = 0;
4520 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004521 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004522 } else {
4523 Py_INCREF(lz->fillvalue);
4524 item = lz->fillvalue;
4525 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4526 Py_DECREF(it);
4527 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004528 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004529 }
4530 olditem = PyTuple_GET_ITEM(result, i);
4531 PyTuple_SET_ITEM(result, i, item);
4532 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004533 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004534 } else {
4535 result = PyTuple_New(tuplesize);
4536 if (result == NULL)
4537 return NULL;
4538 for (i=0 ; i < tuplesize ; i++) {
4539 it = PyTuple_GET_ITEM(lz->ittuple, i);
4540 if (it == NULL) {
4541 Py_INCREF(lz->fillvalue);
4542 item = lz->fillvalue;
4543 } else {
4544 item = PyIter_Next(it);
4545 if (item == NULL) {
4546 lz->numactive -= 1;
4547 if (lz->numactive == 0 || PyErr_Occurred()) {
4548 lz->numactive = 0;
4549 Py_DECREF(result);
4550 return NULL;
4551 } else {
4552 Py_INCREF(lz->fillvalue);
4553 item = lz->fillvalue;
4554 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4555 Py_DECREF(it);
4556 }
4557 }
4558 }
4559 PyTuple_SET_ITEM(result, i, item);
4560 }
4561 }
4562 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004563}
4564
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004565static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304566zip_longest_reduce(ziplongestobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004567{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004568
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004569 /* Create a new tuple with empty sequences where appropriate to pickle.
4570 * Then use setstate to set the fillvalue
4571 */
4572 int i;
4573 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004574
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004575 if (args == NULL)
4576 return NULL;
4577 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4578 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4579 if (elem == NULL) {
4580 elem = PyTuple_New(0);
4581 if (elem == NULL) {
4582 Py_DECREF(args);
4583 return NULL;
4584 }
4585 } else
4586 Py_INCREF(elem);
4587 PyTuple_SET_ITEM(args, i, elem);
4588 }
4589 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4590}
4591
4592static PyObject *
4593zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4594{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004595 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004596 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004597 Py_RETURN_NONE;
4598}
4599
4600static PyMethodDef zip_longest_methods[] = {
4601 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4602 reduce_doc},
4603 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4604 setstate_doc},
4605 {NULL, NULL} /* sentinel */
4606};
4607
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004608PyDoc_STRVAR(zip_longest_doc,
4609"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004610\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004611Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004612the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004613method continues until the longest iterable in the argument sequence\n\
4614is exhausted and then it raises StopIteration. When the shorter iterables\n\
4615are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4616defaults to None or can be specified by a keyword argument.\n\
4617");
4618
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004619static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004620 PyVarObject_HEAD_INIT(NULL, 0)
4621 "itertools.zip_longest", /* tp_name */
4622 sizeof(ziplongestobject), /* tp_basicsize */
4623 0, /* tp_itemsize */
4624 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004625 (destructor)zip_longest_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004626 0, /* tp_print */
4627 0, /* tp_getattr */
4628 0, /* tp_setattr */
4629 0, /* tp_reserved */
4630 0, /* tp_repr */
4631 0, /* tp_as_number */
4632 0, /* tp_as_sequence */
4633 0, /* tp_as_mapping */
4634 0, /* tp_hash */
4635 0, /* tp_call */
4636 0, /* tp_str */
4637 PyObject_GenericGetAttr, /* tp_getattro */
4638 0, /* tp_setattro */
4639 0, /* tp_as_buffer */
4640 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4641 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004642 zip_longest_doc, /* tp_doc */
4643 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004644 0, /* tp_clear */
4645 0, /* tp_richcompare */
4646 0, /* tp_weaklistoffset */
4647 PyObject_SelfIter, /* tp_iter */
4648 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004649 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004650 0, /* tp_members */
4651 0, /* tp_getset */
4652 0, /* tp_base */
4653 0, /* tp_dict */
4654 0, /* tp_descr_get */
4655 0, /* tp_descr_set */
4656 0, /* tp_dictoffset */
4657 0, /* tp_init */
4658 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004659 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004660 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004661};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004662
Tal Einatc4bccd32018-09-12 00:49:13 +03004663
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004664/* module level code ********************************************************/
4665
4666PyDoc_STRVAR(module_doc,
4667"Functional tools for creating and using iterators.\n\
4668\n\
4669Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004670count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004671cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004672repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004673\n\
4674Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004675accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
oldkaa0735f2018-02-02 16:52:55 +08004676chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
4677chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004678compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4679dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4680groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004681filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004682islice(seq, [start,] stop [, step]) --> elements from\n\
4683 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004684starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004685tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004686takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
oldkaa0735f2018-02-02 16:52:55 +08004687zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004688\n\
4689Combinatoric generators:\n\
4690product(p, q, ... [repeat=1]) --> cartesian product\n\
4691permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004692combinations(p, r)\n\
4693combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004694");
4695
4696
Raymond Hettingerad983e72003-11-12 14:32:26 +00004697static PyMethodDef module_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03004698 ITERTOOLS_TEE_METHODDEF
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004699 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004700};
4701
Martin v. Löwis1a214512008-06-11 05:26:20 +00004702
4703static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004704 PyModuleDef_HEAD_INIT,
4705 "itertools",
4706 module_doc,
4707 -1,
4708 module_methods,
4709 NULL,
4710 NULL,
4711 NULL,
4712 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004713};
4714
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004715PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004716PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004717{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004718 int i;
4719 PyObject *m;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004720 const char *name;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004721 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004722 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004723 &combinations_type,
4724 &cwr_type,
4725 &cycle_type,
4726 &dropwhile_type,
4727 &takewhile_type,
4728 &islice_type,
4729 &starmap_type,
4730 &chain_type,
4731 &compress_type,
4732 &filterfalse_type,
4733 &count_type,
4734 &ziplongest_type,
4735 &permutations_type,
4736 &product_type,
4737 &repeat_type,
4738 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004739 &_grouper_type,
4740 &tee_type,
4741 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004742 NULL
4743 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004744
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004745 Py_TYPE(&teedataobject_type) = &PyType_Type;
4746 m = PyModule_Create(&itertoolsmodule);
4747 if (m == NULL)
4748 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004749
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004750 for (i=0 ; typelist[i] != NULL ; i++) {
4751 if (PyType_Ready(typelist[i]) < 0)
4752 return NULL;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004753 name = _PyType_Name(typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004754 Py_INCREF(typelist[i]);
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004755 PyModule_AddObject(m, name, (PyObject *)typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004756 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004757
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004758 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004759}