blob: c00c2745d3f0cb1eae884de84bffa494a5184d0c [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 {
Petr Viktorinffd97532020-02-11 17:46:57 +0100137 newkey = PyObject_CallOneArg(gbo->keyfunc, newvalue);
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300138 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 */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200244 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000245 0, /* tp_getattr */
246 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200247 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000248 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 */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200393 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000394 0, /* tp_getattr */
395 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200396 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000397 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 */
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300446 int running;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000447 PyObject *nextlink;
448 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000449} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000450
451typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000452 PyObject_HEAD
453 teedataobject *dataobj;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700454 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000455 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000456} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000457
Raymond Hettingerad983e72003-11-12 14:32:26 +0000458static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000459
460static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000461teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000462{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000463 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000464
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000465 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
466 if (tdo == NULL)
467 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000468
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300469 tdo->running = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000470 tdo->numread = 0;
471 tdo->nextlink = NULL;
472 Py_INCREF(it);
473 tdo->it = it;
474 PyObject_GC_Track(tdo);
475 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000476}
477
478static PyObject *
479teedataobject_jumplink(teedataobject *tdo)
480{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000481 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000482 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000483 Py_XINCREF(tdo->nextlink);
484 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000485}
486
487static PyObject *
488teedataobject_getitem(teedataobject *tdo, int i)
489{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000490 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000491
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000492 assert(i < LINKCELLS);
493 if (i < tdo->numread)
494 value = tdo->values[i];
495 else {
496 /* this is the lead iterator, so fetch more data */
497 assert(i == tdo->numread);
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300498 if (tdo->running) {
499 PyErr_SetString(PyExc_RuntimeError,
500 "cannot re-enter the tee iterator");
501 return NULL;
502 }
503 tdo->running = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000504 value = PyIter_Next(tdo->it);
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300505 tdo->running = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000506 if (value == NULL)
507 return NULL;
508 tdo->numread++;
509 tdo->values[i] = value;
510 }
511 Py_INCREF(value);
512 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000513}
514
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000515static int
516teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
517{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000518 int i;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700519
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000520 Py_VISIT(tdo->it);
521 for (i = 0; i < tdo->numread; i++)
522 Py_VISIT(tdo->values[i]);
523 Py_VISIT(tdo->nextlink);
524 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000525}
526
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200527static void
528teedataobject_safe_decref(PyObject *obj)
529{
530 while (obj && Py_TYPE(obj) == &teedataobject_type &&
531 Py_REFCNT(obj) == 1) {
532 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
533 ((teedataobject *)obj)->nextlink = NULL;
534 Py_DECREF(obj);
535 obj = nextlink;
536 }
537 Py_XDECREF(obj);
538}
539
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000540static int
541teedataobject_clear(teedataobject *tdo)
542{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000543 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200544 PyObject *tmp;
545
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000546 Py_CLEAR(tdo->it);
547 for (i=0 ; i<tdo->numread ; i++)
548 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200549 tmp = tdo->nextlink;
550 tdo->nextlink = NULL;
551 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000552 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000553}
554
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000555static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000556teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000557{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000558 PyObject_GC_UnTrack(tdo);
559 teedataobject_clear(tdo);
560 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000561}
562
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000563static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530564teedataobject_reduce(teedataobject *tdo, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000565{
566 int i;
567 /* create a temporary list of already iterated values */
568 PyObject *values = PyList_New(tdo->numread);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700569
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000570 if (!values)
571 return NULL;
572 for (i=0 ; i<tdo->numread ; i++) {
573 Py_INCREF(tdo->values[i]);
574 PyList_SET_ITEM(values, i, tdo->values[i]);
575 }
576 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
577 values,
578 tdo->nextlink ? tdo->nextlink : Py_None);
579}
580
Tal Einatc4bccd32018-09-12 00:49:13 +0300581/*[clinic input]
582@classmethod
583itertools.teedataobject.__new__
584 iterable as it: object
585 values: object(subclass_of='&PyList_Type')
586 next: object
587 /
588Data container common to multiple tee objects.
589[clinic start generated code]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000590
591static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300592itertools_teedataobject_impl(PyTypeObject *type, PyObject *it,
593 PyObject *values, PyObject *next)
594/*[clinic end generated code: output=3343ceb07e08df5e input=be60f2fabd2b72ba]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000595{
596 teedataobject *tdo;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000597 Py_ssize_t i, len;
598
599 assert(type == &teedataobject_type);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000600
601 tdo = (teedataobject *)teedataobject_newinternal(it);
602 if (!tdo)
603 return NULL;
604
605 len = PyList_GET_SIZE(values);
606 if (len > LINKCELLS)
607 goto err;
608 for (i=0; i<len; i++) {
609 tdo->values[i] = PyList_GET_ITEM(values, i);
610 Py_INCREF(tdo->values[i]);
611 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200612 /* len <= LINKCELLS < INT_MAX */
613 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000614
615 if (len == LINKCELLS) {
616 if (next != Py_None) {
617 if (Py_TYPE(next) != &teedataobject_type)
618 goto err;
619 assert(tdo->nextlink == NULL);
620 Py_INCREF(next);
621 tdo->nextlink = next;
622 }
623 } else {
624 if (next != Py_None)
625 goto err; /* shouldn't have a next if we are not full */
626 }
627 return (PyObject*)tdo;
628
629err:
630 Py_XDECREF(tdo);
631 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
632 return NULL;
633}
634
635static PyMethodDef teedataobject_methods[] = {
636 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
637 reduce_doc},
638 {NULL, NULL} /* sentinel */
639};
640
Raymond Hettingerad983e72003-11-12 14:32:26 +0000641static PyTypeObject teedataobject_type = {
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700642 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000643 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000644 sizeof(teedataobject), /* tp_basicsize */
645 0, /* tp_itemsize */
646 /* methods */
647 (destructor)teedataobject_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200648 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000649 0, /* tp_getattr */
650 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200651 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000652 0, /* tp_repr */
653 0, /* tp_as_number */
654 0, /* tp_as_sequence */
655 0, /* tp_as_mapping */
656 0, /* tp_hash */
657 0, /* tp_call */
658 0, /* tp_str */
659 PyObject_GenericGetAttr, /* tp_getattro */
660 0, /* tp_setattro */
661 0, /* tp_as_buffer */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700662 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300663 itertools_teedataobject__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000664 (traverseproc)teedataobject_traverse, /* tp_traverse */
665 (inquiry)teedataobject_clear, /* tp_clear */
666 0, /* tp_richcompare */
667 0, /* tp_weaklistoffset */
668 0, /* tp_iter */
669 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000670 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000671 0, /* tp_members */
672 0, /* tp_getset */
673 0, /* tp_base */
674 0, /* tp_dict */
675 0, /* tp_descr_get */
676 0, /* tp_descr_set */
677 0, /* tp_dictoffset */
678 0, /* tp_init */
679 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300680 itertools_teedataobject, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000681 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000682};
683
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000684
685static PyTypeObject tee_type;
686
687static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000688tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000689{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000690 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000691
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000692 if (to->index >= LINKCELLS) {
693 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200694 if (link == NULL)
695 return NULL;
Serhiy Storchaka57a01d32016-04-10 18:05:40 +0300696 Py_SETREF(to->dataobj, (teedataobject *)link);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000697 to->index = 0;
698 }
699 value = teedataobject_getitem(to->dataobj, to->index);
700 if (value == NULL)
701 return NULL;
702 to->index++;
703 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000704}
705
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000706static int
707tee_traverse(teeobject *to, visitproc visit, void *arg)
708{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000709 Py_VISIT((PyObject *)to->dataobj);
710 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000711}
712
Raymond Hettingerad983e72003-11-12 14:32:26 +0000713static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530714tee_copy(teeobject *to, PyObject *Py_UNUSED(ignored))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000715{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000716 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000717
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000718 newto = PyObject_GC_New(teeobject, &tee_type);
719 if (newto == NULL)
720 return NULL;
721 Py_INCREF(to->dataobj);
722 newto->dataobj = to->dataobj;
723 newto->index = to->index;
724 newto->weakreflist = NULL;
725 PyObject_GC_Track(newto);
726 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000727}
728
729PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
730
731static PyObject *
732tee_fromiterable(PyObject *iterable)
733{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000734 teeobject *to;
735 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000736
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000737 it = PyObject_GetIter(iterable);
738 if (it == NULL)
739 return NULL;
740 if (PyObject_TypeCheck(it, &tee_type)) {
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530741 to = (teeobject *)tee_copy((teeobject *)it, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000742 goto done;
743 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000744
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000745 to = PyObject_GC_New(teeobject, &tee_type);
746 if (to == NULL)
747 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000748 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000749 if (!to->dataobj) {
750 PyObject_GC_Del(to);
751 to = NULL;
752 goto done;
753 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000754
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000755 to->index = 0;
756 to->weakreflist = NULL;
757 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000758done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000759 Py_XDECREF(it);
760 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000761}
762
Tal Einatc4bccd32018-09-12 00:49:13 +0300763/*[clinic input]
764@classmethod
765itertools._tee.__new__
766 iterable: object
767 /
768Iterator wrapped to make it copyable.
769[clinic start generated code]*/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000770
Tal Einatc4bccd32018-09-12 00:49:13 +0300771static PyObject *
772itertools__tee_impl(PyTypeObject *type, PyObject *iterable)
773/*[clinic end generated code: output=b02d3fd26c810c3f input=adc0779d2afe37a2]*/
774{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000775 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000776}
777
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000778static int
779tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000780{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000781 if (to->weakreflist != NULL)
782 PyObject_ClearWeakRefs((PyObject *) to);
783 Py_CLEAR(to->dataobj);
784 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000785}
786
787static void
788tee_dealloc(teeobject *to)
789{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000790 PyObject_GC_UnTrack(to);
791 tee_clear(to);
792 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000793}
794
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000795static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530796tee_reduce(teeobject *to, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000797{
798 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
799}
800
801static PyObject *
802tee_setstate(teeobject *to, PyObject *state)
803{
804 teedataobject *tdo;
805 int index;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300806 if (!PyTuple_Check(state)) {
807 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000808 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300809 }
810 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
811 return NULL;
812 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000813 if (index < 0 || index > LINKCELLS) {
814 PyErr_SetString(PyExc_ValueError, "Index out of range");
815 return NULL;
816 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200817 Py_INCREF(tdo);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300818 Py_XSETREF(to->dataobj, tdo);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000819 to->index = index;
820 Py_RETURN_NONE;
821}
822
Raymond Hettingerad983e72003-11-12 14:32:26 +0000823static PyMethodDef tee_methods[] = {
Raymond Hettingera6a2d442015-08-16 14:38:07 -0700824 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
825 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
826 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000827 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000828};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000829
830static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000831 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000832 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000833 sizeof(teeobject), /* tp_basicsize */
834 0, /* tp_itemsize */
835 /* methods */
836 (destructor)tee_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200837 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000838 0, /* tp_getattr */
839 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200840 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000841 0, /* tp_repr */
842 0, /* tp_as_number */
843 0, /* tp_as_sequence */
844 0, /* tp_as_mapping */
845 0, /* tp_hash */
846 0, /* tp_call */
847 0, /* tp_str */
848 0, /* tp_getattro */
849 0, /* tp_setattro */
850 0, /* tp_as_buffer */
851 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300852 itertools__tee__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000853 (traverseproc)tee_traverse, /* tp_traverse */
854 (inquiry)tee_clear, /* tp_clear */
855 0, /* tp_richcompare */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700856 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000857 PyObject_SelfIter, /* tp_iter */
858 (iternextfunc)tee_next, /* tp_iternext */
859 tee_methods, /* tp_methods */
860 0, /* tp_members */
861 0, /* tp_getset */
862 0, /* tp_base */
863 0, /* tp_dict */
864 0, /* tp_descr_get */
865 0, /* tp_descr_set */
866 0, /* tp_dictoffset */
867 0, /* tp_init */
868 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300869 itertools__tee, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000870 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000871};
872
Tal Einatc4bccd32018-09-12 00:49:13 +0300873/*[clinic input]
874itertools.tee
875 iterable: object
876 n: Py_ssize_t = 2
877 /
878Returns a tuple of n independent iterators.
879[clinic start generated code]*/
880
Raymond Hettingerad983e72003-11-12 14:32:26 +0000881static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300882itertools_tee_impl(PyObject *module, PyObject *iterable, Py_ssize_t n)
883/*[clinic end generated code: output=1c64519cd859c2f0 input=c99a1472c425d66d]*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000884{
Tal Einatc4bccd32018-09-12 00:49:13 +0300885 Py_ssize_t i;
886 PyObject *it, *copyable, *copyfunc, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200887 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000888
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000889 if (n < 0) {
890 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
891 return NULL;
892 }
893 result = PyTuple_New(n);
894 if (result == NULL)
895 return NULL;
896 if (n == 0)
897 return result;
898 it = PyObject_GetIter(iterable);
899 if (it == NULL) {
900 Py_DECREF(result);
901 return NULL;
902 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200903
Serhiy Storchakaf320be72018-01-25 10:49:40 +0200904 if (_PyObject_LookupAttrId(it, &PyId___copy__, &copyfunc) < 0) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200905 Py_DECREF(it);
906 Py_DECREF(result);
907 return NULL;
908 }
Serhiy Storchakaf320be72018-01-25 10:49:40 +0200909 if (copyfunc != NULL) {
910 copyable = it;
911 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200912 else {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000913 copyable = tee_fromiterable(it);
914 Py_DECREF(it);
915 if (copyable == NULL) {
916 Py_DECREF(result);
917 return NULL;
918 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200919 copyfunc = _PyObject_GetAttrId(copyable, &PyId___copy__);
920 if (copyfunc == NULL) {
921 Py_DECREF(copyable);
922 Py_DECREF(result);
923 return NULL;
924 }
925 }
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200926
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200927 PyTuple_SET_ITEM(result, 0, copyable);
928 for (i = 1; i < n; i++) {
929 copyable = _PyObject_CallNoArg(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000930 if (copyable == NULL) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200931 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000932 Py_DECREF(result);
933 return NULL;
934 }
935 PyTuple_SET_ITEM(result, i, copyable);
936 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200937 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000938 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000939}
940
Raymond Hettingerad983e72003-11-12 14:32:26 +0000941
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700942/* cycle object **************************************************************/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000943
944typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000945 PyObject_HEAD
946 PyObject *it;
947 PyObject *saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700948 Py_ssize_t index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000949 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000950} cycleobject;
951
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000952static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000953
Tal Einatc4bccd32018-09-12 00:49:13 +0300954/*[clinic input]
955@classmethod
956itertools.cycle.__new__
957 iterable: object
958 /
959Return elements from the iterable until it is exhausted. Then repeat the sequence indefinitely.
960[clinic start generated code]*/
961
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000962static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300963itertools_cycle_impl(PyTypeObject *type, PyObject *iterable)
964/*[clinic end generated code: output=f60e5ec17a45b35c input=9d1d84bcf66e908b]*/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000965{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000966 PyObject *it;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000967 PyObject *saved;
968 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000969
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000970 /* Get iterator. */
971 it = PyObject_GetIter(iterable);
972 if (it == NULL)
973 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000974
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000975 saved = PyList_New(0);
976 if (saved == NULL) {
977 Py_DECREF(it);
978 return NULL;
979 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000980
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000981 /* create cycleobject structure */
982 lz = (cycleobject *)type->tp_alloc(type, 0);
983 if (lz == NULL) {
984 Py_DECREF(it);
985 Py_DECREF(saved);
986 return NULL;
987 }
988 lz->it = it;
989 lz->saved = saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700990 lz->index = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000991 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000992
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000993 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000994}
995
996static void
997cycle_dealloc(cycleobject *lz)
998{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000999 PyObject_GC_UnTrack(lz);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001000 Py_XDECREF(lz->it);
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001001 Py_XDECREF(lz->saved);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001002 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001003}
1004
1005static int
1006cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
1007{
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001008 if (lz->it)
1009 Py_VISIT(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001010 Py_VISIT(lz->saved);
1011 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001012}
1013
1014static PyObject *
1015cycle_next(cycleobject *lz)
1016{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001017 PyObject *item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001018
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001019 if (lz->it != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001020 item = PyIter_Next(lz->it);
1021 if (item != NULL) {
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001022 if (lz->firstpass)
1023 return item;
1024 if (PyList_Append(lz->saved, item)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001025 Py_DECREF(item);
1026 return NULL;
1027 }
1028 return item;
1029 }
Raymond Hettinger98958fe2015-08-15 15:09:30 -07001030 /* Note: StopIteration is already cleared by PyIter_Next() */
1031 if (PyErr_Occurred())
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001032 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001033 Py_CLEAR(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001034 }
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001035 if (PyList_GET_SIZE(lz->saved) == 0)
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001036 return NULL;
1037 item = PyList_GET_ITEM(lz->saved, lz->index);
1038 lz->index++;
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001039 if (lz->index >= PyList_GET_SIZE(lz->saved))
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001040 lz->index = 0;
1041 Py_INCREF(item);
1042 return item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001043}
1044
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001045static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301046cycle_reduce(cycleobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001047{
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001048 /* Create a new cycle with the iterator tuple, then set the saved state */
1049 if (lz->it == NULL) {
1050 PyObject *it = PyObject_GetIter(lz->saved);
1051 if (it == NULL)
1052 return NULL;
1053 if (lz->index != 0) {
1054 _Py_IDENTIFIER(__setstate__);
1055 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
1056 "n", lz->index);
1057 if (res == NULL) {
1058 Py_DECREF(it);
1059 return NULL;
1060 }
1061 Py_DECREF(res);
1062 }
Serhiy Storchaka1f21eaa2019-09-01 12:16:51 +03001063 return Py_BuildValue("O(N)(OO)", Py_TYPE(lz), it, lz->saved, Py_True);
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001064 }
Serhiy Storchaka1f21eaa2019-09-01 12:16:51 +03001065 return Py_BuildValue("O(O)(OO)", Py_TYPE(lz), lz->it, lz->saved,
1066 lz->firstpass ? Py_True : Py_False);
Raymond Hettingerc39786d2015-08-15 15:16:12 -07001067}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001068
1069static PyObject *
1070cycle_setstate(cycleobject *lz, PyObject *state)
1071{
1072 PyObject *saved=NULL;
1073 int firstpass;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001074 if (!PyTuple_Check(state)) {
1075 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001076 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001077 }
1078 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
1079 return NULL;
1080 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001081 Py_INCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001082 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001083 lz->firstpass = firstpass != 0;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001084 lz->index = 0;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001085 Py_RETURN_NONE;
1086}
1087
1088static PyMethodDef cycle_methods[] = {
1089 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1090 reduce_doc},
1091 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1092 setstate_doc},
1093 {NULL, NULL} /* sentinel */
1094};
1095
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001096static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001097 PyVarObject_HEAD_INIT(NULL, 0)
1098 "itertools.cycle", /* tp_name */
1099 sizeof(cycleobject), /* tp_basicsize */
1100 0, /* tp_itemsize */
1101 /* methods */
1102 (destructor)cycle_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001103 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001104 0, /* tp_getattr */
1105 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001106 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001107 0, /* tp_repr */
1108 0, /* tp_as_number */
1109 0, /* tp_as_sequence */
1110 0, /* tp_as_mapping */
1111 0, /* tp_hash */
1112 0, /* tp_call */
1113 0, /* tp_str */
1114 PyObject_GenericGetAttr, /* tp_getattro */
1115 0, /* tp_setattro */
1116 0, /* tp_as_buffer */
1117 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1118 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001119 itertools_cycle__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001120 (traverseproc)cycle_traverse, /* tp_traverse */
1121 0, /* tp_clear */
1122 0, /* tp_richcompare */
1123 0, /* tp_weaklistoffset */
1124 PyObject_SelfIter, /* tp_iter */
1125 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001126 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001127 0, /* tp_members */
1128 0, /* tp_getset */
1129 0, /* tp_base */
1130 0, /* tp_dict */
1131 0, /* tp_descr_get */
1132 0, /* tp_descr_set */
1133 0, /* tp_dictoffset */
1134 0, /* tp_init */
1135 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001136 itertools_cycle, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001137 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001138};
1139
1140
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001141/* dropwhile object **********************************************************/
1142
1143typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001144 PyObject_HEAD
1145 PyObject *func;
1146 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001147 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001148} dropwhileobject;
1149
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001150static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001151
Tal Einatc4bccd32018-09-12 00:49:13 +03001152/*[clinic input]
1153@classmethod
1154itertools.dropwhile.__new__
1155 predicate as func: object
1156 iterable as seq: object
1157 /
1158Drop items from the iterable while predicate(item) is true.
1159
1160Afterwards, return every element until the iterable is exhausted.
1161[clinic start generated code]*/
1162
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001163static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001164itertools_dropwhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1165/*[clinic end generated code: output=92f9d0d89af149e4 input=d39737147c9f0a26]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001166{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001167 PyObject *it;
1168 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001169
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001170 /* Get iterator. */
1171 it = PyObject_GetIter(seq);
1172 if (it == NULL)
1173 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001174
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001175 /* create dropwhileobject structure */
1176 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1177 if (lz == NULL) {
1178 Py_DECREF(it);
1179 return NULL;
1180 }
1181 Py_INCREF(func);
1182 lz->func = func;
1183 lz->it = it;
1184 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001185
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001186 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001187}
1188
1189static void
1190dropwhile_dealloc(dropwhileobject *lz)
1191{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001192 PyObject_GC_UnTrack(lz);
1193 Py_XDECREF(lz->func);
1194 Py_XDECREF(lz->it);
1195 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196}
1197
1198static int
1199dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1200{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001201 Py_VISIT(lz->it);
1202 Py_VISIT(lz->func);
1203 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204}
1205
1206static PyObject *
1207dropwhile_next(dropwhileobject *lz)
1208{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001209 PyObject *item, *good;
1210 PyObject *it = lz->it;
1211 long ok;
1212 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001213
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001214 iternext = *Py_TYPE(it)->tp_iternext;
1215 for (;;) {
1216 item = iternext(it);
1217 if (item == NULL)
1218 return NULL;
1219 if (lz->start == 1)
1220 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001221
Petr Viktorinffd97532020-02-11 17:46:57 +01001222 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001223 if (good == NULL) {
1224 Py_DECREF(item);
1225 return NULL;
1226 }
1227 ok = PyObject_IsTrue(good);
1228 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001229 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001230 lz->start = 1;
1231 return item;
1232 }
1233 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001234 if (ok < 0)
1235 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001236 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001237}
1238
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001239static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301240dropwhile_reduce(dropwhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001241{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001242 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 +00001243}
1244
1245static PyObject *
1246dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1247{
1248 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001249 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001250 return NULL;
1251 lz->start = start;
1252 Py_RETURN_NONE;
1253}
1254
1255static PyMethodDef dropwhile_methods[] = {
1256 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1257 reduce_doc},
1258 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1259 setstate_doc},
1260 {NULL, NULL} /* sentinel */
1261};
1262
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001263static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001264 PyVarObject_HEAD_INIT(NULL, 0)
1265 "itertools.dropwhile", /* tp_name */
1266 sizeof(dropwhileobject), /* tp_basicsize */
1267 0, /* tp_itemsize */
1268 /* methods */
1269 (destructor)dropwhile_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001270 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001271 0, /* tp_getattr */
1272 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001273 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001274 0, /* tp_repr */
1275 0, /* tp_as_number */
1276 0, /* tp_as_sequence */
1277 0, /* tp_as_mapping */
1278 0, /* tp_hash */
1279 0, /* tp_call */
1280 0, /* tp_str */
1281 PyObject_GenericGetAttr, /* tp_getattro */
1282 0, /* tp_setattro */
1283 0, /* tp_as_buffer */
1284 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1285 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001286 itertools_dropwhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001287 (traverseproc)dropwhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001288 0, /* tp_clear */
1289 0, /* tp_richcompare */
1290 0, /* tp_weaklistoffset */
1291 PyObject_SelfIter, /* tp_iter */
1292 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001293 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001294 0, /* tp_members */
1295 0, /* tp_getset */
1296 0, /* tp_base */
1297 0, /* tp_dict */
1298 0, /* tp_descr_get */
1299 0, /* tp_descr_set */
1300 0, /* tp_dictoffset */
1301 0, /* tp_init */
1302 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001303 itertools_dropwhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001304 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001305};
1306
1307
1308/* takewhile object **********************************************************/
1309
1310typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001311 PyObject_HEAD
1312 PyObject *func;
1313 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001314 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001315} takewhileobject;
1316
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001317static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001318
Tal Einatc4bccd32018-09-12 00:49:13 +03001319/*[clinic input]
1320@classmethod
1321itertools.takewhile.__new__
1322 predicate as func: object
1323 iterable as seq: object
1324 /
1325Return successive entries from an iterable as long as the predicate evaluates to true for each entry.
1326[clinic start generated code]*/
1327
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001328static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001329itertools_takewhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1330/*[clinic end generated code: output=bb179ea7864e2ef6 input=ba5255f7519aa119]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001331{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001332 PyObject *it;
1333 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001334
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001335 /* Get iterator. */
1336 it = PyObject_GetIter(seq);
1337 if (it == NULL)
1338 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001339
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001340 /* create takewhileobject structure */
1341 lz = (takewhileobject *)type->tp_alloc(type, 0);
1342 if (lz == NULL) {
1343 Py_DECREF(it);
1344 return NULL;
1345 }
1346 Py_INCREF(func);
1347 lz->func = func;
1348 lz->it = it;
1349 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001350
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001351 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001352}
1353
1354static void
1355takewhile_dealloc(takewhileobject *lz)
1356{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001357 PyObject_GC_UnTrack(lz);
1358 Py_XDECREF(lz->func);
1359 Py_XDECREF(lz->it);
1360 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001361}
1362
1363static int
1364takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1365{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001366 Py_VISIT(lz->it);
1367 Py_VISIT(lz->func);
1368 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001369}
1370
1371static PyObject *
1372takewhile_next(takewhileobject *lz)
1373{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001374 PyObject *item, *good;
1375 PyObject *it = lz->it;
1376 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001377
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001378 if (lz->stop == 1)
1379 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001380
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001381 item = (*Py_TYPE(it)->tp_iternext)(it);
1382 if (item == NULL)
1383 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001384
Petr Viktorinffd97532020-02-11 17:46:57 +01001385 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001386 if (good == NULL) {
1387 Py_DECREF(item);
1388 return NULL;
1389 }
1390 ok = PyObject_IsTrue(good);
1391 Py_DECREF(good);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001392 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001393 return item;
1394 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001395 if (ok == 0)
1396 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001397 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001398}
1399
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001400static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301401takewhile_reduce(takewhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001402{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001403 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 +00001404}
1405
1406static PyObject *
1407takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1408{
1409 int stop = PyObject_IsTrue(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001410
Antoine Pitrou721738f2012-08-15 23:20:39 +02001411 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001412 return NULL;
1413 lz->stop = stop;
1414 Py_RETURN_NONE;
1415}
1416
1417static PyMethodDef takewhile_reduce_methods[] = {
1418 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1419 reduce_doc},
1420 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1421 setstate_doc},
1422 {NULL, NULL} /* sentinel */
1423};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001424
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001425static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001426 PyVarObject_HEAD_INIT(NULL, 0)
1427 "itertools.takewhile", /* tp_name */
1428 sizeof(takewhileobject), /* tp_basicsize */
1429 0, /* tp_itemsize */
1430 /* methods */
1431 (destructor)takewhile_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001432 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001433 0, /* tp_getattr */
1434 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001435 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001436 0, /* tp_repr */
1437 0, /* tp_as_number */
1438 0, /* tp_as_sequence */
1439 0, /* tp_as_mapping */
1440 0, /* tp_hash */
1441 0, /* tp_call */
1442 0, /* tp_str */
1443 PyObject_GenericGetAttr, /* tp_getattro */
1444 0, /* tp_setattro */
1445 0, /* tp_as_buffer */
1446 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1447 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001448 itertools_takewhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001449 (traverseproc)takewhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001450 0, /* tp_clear */
1451 0, /* tp_richcompare */
1452 0, /* tp_weaklistoffset */
1453 PyObject_SelfIter, /* tp_iter */
1454 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001455 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001456 0, /* tp_members */
1457 0, /* tp_getset */
1458 0, /* tp_base */
1459 0, /* tp_dict */
1460 0, /* tp_descr_get */
1461 0, /* tp_descr_set */
1462 0, /* tp_dictoffset */
1463 0, /* tp_init */
1464 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001465 itertools_takewhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001466 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001467};
1468
1469
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001470/* islice object *************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001471
1472typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001473 PyObject_HEAD
1474 PyObject *it;
1475 Py_ssize_t next;
1476 Py_ssize_t stop;
1477 Py_ssize_t step;
1478 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001479} isliceobject;
1480
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001481static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001482
1483static PyObject *
1484islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1485{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001486 PyObject *seq;
1487 Py_ssize_t start=0, stop=-1, step=1;
1488 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1489 Py_ssize_t numargs;
1490 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001491
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001492 if (type == &islice_type && !_PyArg_NoKeywords("islice", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001493 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001494
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001495 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1496 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001497
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001498 numargs = PyTuple_Size(args);
1499 if (numargs == 2) {
1500 if (a1 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001501 stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001502 if (stop == -1) {
1503 if (PyErr_Occurred())
1504 PyErr_Clear();
1505 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001506 "Stop argument for islice() must be None or "
1507 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001508 return NULL;
1509 }
1510 }
1511 } else {
1512 if (a1 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001513 start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001514 if (start == -1 && PyErr_Occurred())
1515 PyErr_Clear();
1516 if (a2 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001517 stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001518 if (stop == -1) {
1519 if (PyErr_Occurred())
1520 PyErr_Clear();
1521 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001522 "Stop argument for islice() must be None or "
1523 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001524 return NULL;
1525 }
1526 }
1527 }
1528 if (start<0 || stop<-1) {
1529 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001530 "Indices for islice() must be None or "
1531 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001532 return NULL;
1533 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001534
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001535 if (a3 != NULL) {
1536 if (a3 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001537 step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001538 if (step == -1 && PyErr_Occurred())
1539 PyErr_Clear();
1540 }
1541 if (step<1) {
1542 PyErr_SetString(PyExc_ValueError,
1543 "Step for islice() must be a positive integer or None.");
1544 return NULL;
1545 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001546
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001547 /* Get iterator. */
1548 it = PyObject_GetIter(seq);
1549 if (it == NULL)
1550 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001551
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001552 /* create isliceobject structure */
1553 lz = (isliceobject *)type->tp_alloc(type, 0);
1554 if (lz == NULL) {
1555 Py_DECREF(it);
1556 return NULL;
1557 }
1558 lz->it = it;
1559 lz->next = start;
1560 lz->stop = stop;
1561 lz->step = step;
1562 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001563
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001564 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001565}
1566
1567static void
1568islice_dealloc(isliceobject *lz)
1569{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001570 PyObject_GC_UnTrack(lz);
1571 Py_XDECREF(lz->it);
1572 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001573}
1574
1575static int
1576islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1577{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001578 Py_VISIT(lz->it);
1579 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001580}
1581
1582static PyObject *
1583islice_next(isliceobject *lz)
1584{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001585 PyObject *item;
1586 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001587 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001588 Py_ssize_t oldnext;
1589 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001590
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001591 if (it == NULL)
1592 return NULL;
1593
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001594 iternext = *Py_TYPE(it)->tp_iternext;
1595 while (lz->cnt < lz->next) {
1596 item = iternext(it);
1597 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001598 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001599 Py_DECREF(item);
1600 lz->cnt++;
1601 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001602 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001603 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001604 item = iternext(it);
1605 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001606 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001607 lz->cnt++;
1608 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001609 /* The (size_t) cast below avoids the danger of undefined
1610 behaviour from signed integer overflow. */
1611 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001612 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1613 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001614 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001615
1616empty:
1617 Py_CLEAR(lz->it);
1618 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001619}
1620
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001621static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301622islice_reduce(isliceobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001623{
1624 /* When unpickled, generate a new object with the same bounds,
1625 * then 'setstate' with the next and count
1626 */
1627 PyObject *stop;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001628
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001629 if (lz->it == NULL) {
1630 PyObject *empty_list;
1631 PyObject *empty_it;
1632 empty_list = PyList_New(0);
1633 if (empty_list == NULL)
1634 return NULL;
1635 empty_it = PyObject_GetIter(empty_list);
1636 Py_DECREF(empty_list);
1637 if (empty_it == NULL)
1638 return NULL;
1639 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1640 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001641 if (lz->stop == -1) {
1642 stop = Py_None;
1643 Py_INCREF(stop);
1644 } else {
1645 stop = PyLong_FromSsize_t(lz->stop);
1646 if (stop == NULL)
1647 return NULL;
1648 }
1649 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1650 lz->it, lz->next, stop, lz->step,
1651 lz->cnt);
1652}
1653
1654static PyObject *
1655islice_setstate(isliceobject *lz, PyObject *state)
1656{
1657 Py_ssize_t cnt = PyLong_AsSsize_t(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001658
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001659 if (cnt == -1 && PyErr_Occurred())
1660 return NULL;
1661 lz->cnt = cnt;
1662 Py_RETURN_NONE;
1663}
1664
1665static PyMethodDef islice_methods[] = {
1666 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1667 reduce_doc},
1668 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1669 setstate_doc},
1670 {NULL, NULL} /* sentinel */
1671};
1672
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001673PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001674"islice(iterable, stop) --> islice object\n\
1675islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001676\n\
1677Return an iterator whose next() method returns selected values from an\n\
1678iterable. If start is specified, will skip all preceding elements;\n\
1679otherwise, start defaults to zero. Step defaults to one. If\n\
oldkaa0735f2018-02-02 16:52:55 +08001680specified as another value, step determines how many values are\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001681skipped between successive calls. Works like a slice() on a list\n\
1682but returns an iterator.");
1683
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001684static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001685 PyVarObject_HEAD_INIT(NULL, 0)
1686 "itertools.islice", /* tp_name */
1687 sizeof(isliceobject), /* tp_basicsize */
1688 0, /* tp_itemsize */
1689 /* methods */
1690 (destructor)islice_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001691 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001692 0, /* tp_getattr */
1693 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001694 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001695 0, /* tp_repr */
1696 0, /* tp_as_number */
1697 0, /* tp_as_sequence */
1698 0, /* tp_as_mapping */
1699 0, /* tp_hash */
1700 0, /* tp_call */
1701 0, /* tp_str */
1702 PyObject_GenericGetAttr, /* tp_getattro */
1703 0, /* tp_setattro */
1704 0, /* tp_as_buffer */
1705 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1706 Py_TPFLAGS_BASETYPE, /* tp_flags */
1707 islice_doc, /* tp_doc */
1708 (traverseproc)islice_traverse, /* tp_traverse */
1709 0, /* tp_clear */
1710 0, /* tp_richcompare */
1711 0, /* tp_weaklistoffset */
1712 PyObject_SelfIter, /* tp_iter */
1713 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001714 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001715 0, /* tp_members */
1716 0, /* tp_getset */
1717 0, /* tp_base */
1718 0, /* tp_dict */
1719 0, /* tp_descr_get */
1720 0, /* tp_descr_set */
1721 0, /* tp_dictoffset */
1722 0, /* tp_init */
1723 0, /* tp_alloc */
1724 islice_new, /* tp_new */
1725 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001726};
1727
1728
1729/* starmap object ************************************************************/
1730
1731typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001732 PyObject_HEAD
1733 PyObject *func;
1734 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001735} starmapobject;
1736
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001737static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001738
Tal Einatc4bccd32018-09-12 00:49:13 +03001739/*[clinic input]
1740@classmethod
1741itertools.starmap.__new__
1742 function as func: object
1743 iterable as seq: object
1744 /
1745Return an iterator whose values are returned from the function evaluated with an argument tuple taken from the given sequence.
1746[clinic start generated code]*/
1747
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001748static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001749itertools_starmap_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1750/*[clinic end generated code: output=79eeb81d452c6e8d input=844766df6a0d4dad]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001751{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001752 PyObject *it;
1753 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001754
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001755 /* Get iterator. */
1756 it = PyObject_GetIter(seq);
1757 if (it == NULL)
1758 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001759
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001760 /* create starmapobject structure */
1761 lz = (starmapobject *)type->tp_alloc(type, 0);
1762 if (lz == NULL) {
1763 Py_DECREF(it);
1764 return NULL;
1765 }
1766 Py_INCREF(func);
1767 lz->func = func;
1768 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001769
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001770 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001771}
1772
1773static void
1774starmap_dealloc(starmapobject *lz)
1775{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001776 PyObject_GC_UnTrack(lz);
1777 Py_XDECREF(lz->func);
1778 Py_XDECREF(lz->it);
1779 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001780}
1781
1782static int
1783starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1784{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001785 Py_VISIT(lz->it);
1786 Py_VISIT(lz->func);
1787 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001788}
1789
1790static PyObject *
1791starmap_next(starmapobject *lz)
1792{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001793 PyObject *args;
1794 PyObject *result;
1795 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001796
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001797 args = (*Py_TYPE(it)->tp_iternext)(it);
1798 if (args == NULL)
1799 return NULL;
1800 if (!PyTuple_CheckExact(args)) {
1801 PyObject *newargs = PySequence_Tuple(args);
1802 Py_DECREF(args);
1803 if (newargs == NULL)
1804 return NULL;
1805 args = newargs;
1806 }
1807 result = PyObject_Call(lz->func, args, NULL);
1808 Py_DECREF(args);
1809 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001810}
1811
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001812static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301813starmap_reduce(starmapobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001814{
1815 /* Just pickle the iterator */
1816 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1817}
1818
1819static PyMethodDef starmap_methods[] = {
1820 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1821 reduce_doc},
1822 {NULL, NULL} /* sentinel */
1823};
1824
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001825static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001826 PyVarObject_HEAD_INIT(NULL, 0)
1827 "itertools.starmap", /* tp_name */
1828 sizeof(starmapobject), /* tp_basicsize */
1829 0, /* tp_itemsize */
1830 /* methods */
1831 (destructor)starmap_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001832 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001833 0, /* tp_getattr */
1834 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001835 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001836 0, /* tp_repr */
1837 0, /* tp_as_number */
1838 0, /* tp_as_sequence */
1839 0, /* tp_as_mapping */
1840 0, /* tp_hash */
1841 0, /* tp_call */
1842 0, /* tp_str */
1843 PyObject_GenericGetAttr, /* tp_getattro */
1844 0, /* tp_setattro */
1845 0, /* tp_as_buffer */
1846 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1847 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001848 itertools_starmap__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001849 (traverseproc)starmap_traverse, /* tp_traverse */
1850 0, /* tp_clear */
1851 0, /* tp_richcompare */
1852 0, /* tp_weaklistoffset */
1853 PyObject_SelfIter, /* tp_iter */
1854 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001855 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001856 0, /* tp_members */
1857 0, /* tp_getset */
1858 0, /* tp_base */
1859 0, /* tp_dict */
1860 0, /* tp_descr_get */
1861 0, /* tp_descr_set */
1862 0, /* tp_dictoffset */
1863 0, /* tp_init */
1864 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001865 itertools_starmap, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001866 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001867};
1868
1869
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001870/* chain object **************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001871
1872typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001873 PyObject_HEAD
1874 PyObject *source; /* Iterator over input iterables */
1875 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001876} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001877
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001878static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001879
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001880static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001881chain_new_internal(PyTypeObject *type, PyObject *source)
1882{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001883 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001884
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001885 lz = (chainobject *)type->tp_alloc(type, 0);
1886 if (lz == NULL) {
1887 Py_DECREF(source);
1888 return NULL;
1889 }
1890
1891 lz->source = source;
1892 lz->active = NULL;
1893 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001894}
1895
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001896static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001897chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001898{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001899 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001900
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001901 if (type == &chain_type && !_PyArg_NoKeywords("chain", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001902 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001903
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001904 source = PyObject_GetIter(args);
1905 if (source == NULL)
1906 return NULL;
1907
1908 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001909}
1910
Tal Einatc4bccd32018-09-12 00:49:13 +03001911/*[clinic input]
1912@classmethod
1913itertools.chain.from_iterable
1914 iterable as arg: object
1915 /
1916Alternative chain() constructor taking a single iterable argument that evaluates lazily.
1917[clinic start generated code]*/
1918
Christian Heimesf16baeb2008-02-29 14:57:44 +00001919static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001920itertools_chain_from_iterable(PyTypeObject *type, PyObject *arg)
1921/*[clinic end generated code: output=667ae7a7f7b68654 input=72c39e3a2ca3be85]*/
Christian Heimesf16baeb2008-02-29 14:57:44 +00001922{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001923 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001924
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001925 source = PyObject_GetIter(arg);
1926 if (source == NULL)
1927 return NULL;
1928
1929 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001930}
1931
1932static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001933chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001934{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001935 PyObject_GC_UnTrack(lz);
1936 Py_XDECREF(lz->active);
1937 Py_XDECREF(lz->source);
1938 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001939}
1940
Raymond Hettinger2012f172003-02-07 05:32:58 +00001941static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001942chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001943{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001944 Py_VISIT(lz->source);
1945 Py_VISIT(lz->active);
1946 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001947}
1948
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001949static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001950chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001951{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001952 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001953
T. Wouters5466d4a2017-03-30 09:58:35 -07001954 /* lz->source is the iterator of iterables. If it's NULL, we've already
1955 * consumed them all. lz->active is the current iterator. If it's NULL,
1956 * we should grab a new one from lz->source. */
1957 while (lz->source != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001958 if (lz->active == NULL) {
T. Wouters5466d4a2017-03-30 09:58:35 -07001959 PyObject *iterable = PyIter_Next(lz->source);
1960 if (iterable == NULL) {
1961 Py_CLEAR(lz->source);
1962 return NULL; /* no more input sources */
1963 }
1964 lz->active = PyObject_GetIter(iterable);
1965 Py_DECREF(iterable);
1966 if (lz->active == NULL) {
1967 Py_CLEAR(lz->source);
1968 return NULL; /* input not iterable */
1969 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001970 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001971 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1972 if (item != NULL)
1973 return item;
1974 if (PyErr_Occurred()) {
1975 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1976 PyErr_Clear();
1977 else
1978 return NULL; /* input raised an exception */
1979 }
1980 /* lz->active is consumed, try with the next iterable. */
1981 Py_CLEAR(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001982 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001983 /* Everything had been consumed already. */
1984 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001985}
1986
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001987static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301988chain_reduce(chainobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001989{
1990 if (lz->source) {
1991 /* we can't pickle function objects (itertools.from_iterable) so
1992 * we must use setstate to replace the iterable. One day we
1993 * will fix pickling of functions
1994 */
1995 if (lz->active) {
1996 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1997 } else {
1998 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1999 }
2000 } else {
2001 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
2002 }
2003 return NULL;
2004}
2005
2006static PyObject *
2007chain_setstate(chainobject *lz, PyObject *state)
2008{
2009 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002010
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002011 if (!PyTuple_Check(state)) {
2012 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002013 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002014 }
2015 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
2016 return NULL;
2017 }
2018 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
2019 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
2020 return NULL;
2021 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002022
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002023 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002024 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002025 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002026 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002027 Py_RETURN_NONE;
2028}
2029
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002030PyDoc_STRVAR(chain_doc,
2031"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002032\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002033Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002034first iterable until it is exhausted, then elements from the next\n\
2035iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002036
Christian Heimesf16baeb2008-02-29 14:57:44 +00002037static PyMethodDef chain_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03002038 ITERTOOLS_CHAIN_FROM_ITERABLE_METHODDEF
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002039 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
2040 reduce_doc},
2041 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
2042 setstate_doc},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07002043 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00002044};
2045
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002046static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002047 PyVarObject_HEAD_INIT(NULL, 0)
2048 "itertools.chain", /* tp_name */
2049 sizeof(chainobject), /* tp_basicsize */
2050 0, /* tp_itemsize */
2051 /* methods */
2052 (destructor)chain_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002053 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002054 0, /* tp_getattr */
2055 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002056 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002057 0, /* tp_repr */
2058 0, /* tp_as_number */
2059 0, /* tp_as_sequence */
2060 0, /* tp_as_mapping */
2061 0, /* tp_hash */
2062 0, /* tp_call */
2063 0, /* tp_str */
2064 PyObject_GenericGetAttr, /* tp_getattro */
2065 0, /* tp_setattro */
2066 0, /* tp_as_buffer */
2067 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2068 Py_TPFLAGS_BASETYPE, /* tp_flags */
2069 chain_doc, /* tp_doc */
2070 (traverseproc)chain_traverse, /* tp_traverse */
2071 0, /* tp_clear */
2072 0, /* tp_richcompare */
2073 0, /* tp_weaklistoffset */
2074 PyObject_SelfIter, /* tp_iter */
2075 (iternextfunc)chain_next, /* tp_iternext */
2076 chain_methods, /* tp_methods */
2077 0, /* tp_members */
2078 0, /* tp_getset */
2079 0, /* tp_base */
2080 0, /* tp_dict */
2081 0, /* tp_descr_get */
2082 0, /* tp_descr_set */
2083 0, /* tp_dictoffset */
2084 0, /* tp_init */
2085 0, /* tp_alloc */
2086 chain_new, /* tp_new */
2087 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002088};
2089
2090
Christian Heimesc3f30c42008-02-22 16:37:40 +00002091/* product object ************************************************************/
2092
2093typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002094 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002095 PyObject *pools; /* tuple of pool tuples */
2096 Py_ssize_t *indices; /* one index per pool */
2097 PyObject *result; /* most recently returned result tuple */
2098 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002099} productobject;
2100
2101static PyTypeObject product_type;
2102
2103static PyObject *
2104product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2105{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002106 productobject *lz;
2107 Py_ssize_t nargs, npools, repeat=1;
2108 PyObject *pools = NULL;
2109 Py_ssize_t *indices = NULL;
2110 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002111
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002112 if (kwds != NULL) {
2113 char *kwlist[] = {"repeat", 0};
2114 PyObject *tmpargs = PyTuple_New(0);
2115 if (tmpargs == NULL)
2116 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002117 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2118 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002119 Py_DECREF(tmpargs);
2120 return NULL;
2121 }
2122 Py_DECREF(tmpargs);
2123 if (repeat < 0) {
2124 PyErr_SetString(PyExc_ValueError,
2125 "repeat argument cannot be negative");
2126 return NULL;
2127 }
2128 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002129
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002130 assert(PyTuple_CheckExact(args));
2131 if (repeat == 0) {
2132 nargs = 0;
2133 } else {
2134 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002135 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002136 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2137 return NULL;
2138 }
2139 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002140 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002141
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002142 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002143 if (indices == NULL) {
2144 PyErr_NoMemory();
2145 goto error;
2146 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002147
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002148 pools = PyTuple_New(npools);
2149 if (pools == NULL)
2150 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002151
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002152 for (i=0; i < nargs ; ++i) {
2153 PyObject *item = PyTuple_GET_ITEM(args, i);
2154 PyObject *pool = PySequence_Tuple(item);
2155 if (pool == NULL)
2156 goto error;
2157 PyTuple_SET_ITEM(pools, i, pool);
2158 indices[i] = 0;
2159 }
2160 for ( ; i < npools; ++i) {
2161 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2162 Py_INCREF(pool);
2163 PyTuple_SET_ITEM(pools, i, pool);
2164 indices[i] = 0;
2165 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002166
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002167 /* create productobject structure */
2168 lz = (productobject *)type->tp_alloc(type, 0);
2169 if (lz == NULL)
2170 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002171
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002172 lz->pools = pools;
2173 lz->indices = indices;
2174 lz->result = NULL;
2175 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002176
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002177 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002178
2179error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002180 if (indices != NULL)
2181 PyMem_Free(indices);
2182 Py_XDECREF(pools);
2183 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002184}
2185
2186static void
2187product_dealloc(productobject *lz)
2188{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002189 PyObject_GC_UnTrack(lz);
2190 Py_XDECREF(lz->pools);
2191 Py_XDECREF(lz->result);
2192 if (lz->indices != NULL)
2193 PyMem_Free(lz->indices);
2194 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002195}
2196
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002197static PyObject *
2198product_sizeof(productobject *lz, void *unused)
2199{
2200 Py_ssize_t res;
2201
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002202 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002203 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2204 return PyLong_FromSsize_t(res);
2205}
2206
2207PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2208
Christian Heimesc3f30c42008-02-22 16:37:40 +00002209static int
2210product_traverse(productobject *lz, visitproc visit, void *arg)
2211{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002212 Py_VISIT(lz->pools);
2213 Py_VISIT(lz->result);
2214 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002215}
2216
2217static PyObject *
2218product_next(productobject *lz)
2219{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002220 PyObject *pool;
2221 PyObject *elem;
2222 PyObject *oldelem;
2223 PyObject *pools = lz->pools;
2224 PyObject *result = lz->result;
2225 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2226 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002227
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002228 if (lz->stopped)
2229 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002230
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002231 if (result == NULL) {
2232 /* On the first pass, return an initial tuple filled with the
2233 first element from each pool. */
2234 result = PyTuple_New(npools);
2235 if (result == NULL)
2236 goto empty;
2237 lz->result = result;
2238 for (i=0; i < npools; i++) {
2239 pool = PyTuple_GET_ITEM(pools, i);
2240 if (PyTuple_GET_SIZE(pool) == 0)
2241 goto empty;
2242 elem = PyTuple_GET_ITEM(pool, 0);
2243 Py_INCREF(elem);
2244 PyTuple_SET_ITEM(result, i, elem);
2245 }
2246 } else {
2247 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002248
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002249 /* Copy the previous result tuple or re-use it if available */
2250 if (Py_REFCNT(result) > 1) {
2251 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002252 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002253 if (result == NULL)
2254 goto empty;
2255 lz->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002256 Py_DECREF(old_result);
2257 }
2258 /* Now, we've got the only copy so we can update it in-place */
2259 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002260
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002261 /* Update the pool indices right-to-left. Only advance to the
2262 next pool when the previous one rolls-over */
2263 for (i=npools-1 ; i >= 0 ; i--) {
2264 pool = PyTuple_GET_ITEM(pools, i);
2265 indices[i]++;
2266 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2267 /* Roll-over and advance to next pool */
2268 indices[i] = 0;
2269 elem = PyTuple_GET_ITEM(pool, 0);
2270 Py_INCREF(elem);
2271 oldelem = PyTuple_GET_ITEM(result, i);
2272 PyTuple_SET_ITEM(result, i, elem);
2273 Py_DECREF(oldelem);
2274 } else {
2275 /* No rollover. Just increment and stop here. */
2276 elem = PyTuple_GET_ITEM(pool, indices[i]);
2277 Py_INCREF(elem);
2278 oldelem = PyTuple_GET_ITEM(result, i);
2279 PyTuple_SET_ITEM(result, i, elem);
2280 Py_DECREF(oldelem);
2281 break;
2282 }
2283 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002284
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002285 /* If i is negative, then the indices have all rolled-over
2286 and we're done. */
2287 if (i < 0)
2288 goto empty;
2289 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002290
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002291 Py_INCREF(result);
2292 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002293
2294empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002295 lz->stopped = 1;
2296 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002297}
2298
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002299static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302300product_reduce(productobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002301{
2302 if (lz->stopped) {
2303 return Py_BuildValue("O(())", Py_TYPE(lz));
2304 } else if (lz->result == NULL) {
2305 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2306 } else {
2307 PyObject *indices;
2308 Py_ssize_t n, i;
2309
2310 /* we must pickle the indices use them for setstate, and
2311 * additionally indicate that the iterator has started
2312 */
2313 n = PyTuple_GET_SIZE(lz->pools);
2314 indices = PyTuple_New(n);
2315 if (indices == NULL)
2316 return NULL;
2317 for (i=0; i<n; i++){
2318 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2319 if (!index) {
2320 Py_DECREF(indices);
2321 return NULL;
2322 }
2323 PyTuple_SET_ITEM(indices, i, index);
2324 }
2325 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2326 }
2327}
2328
2329static PyObject *
2330product_setstate(productobject *lz, PyObject *state)
2331{
2332 PyObject *result;
2333 Py_ssize_t n, i;
2334
2335 n = PyTuple_GET_SIZE(lz->pools);
2336 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2337 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2338 return NULL;
2339 }
2340 for (i=0; i<n; i++)
2341 {
2342 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2343 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002344 PyObject* pool;
2345 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002346 if (index < 0 && PyErr_Occurred())
2347 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002348 pool = PyTuple_GET_ITEM(lz->pools, i);
2349 poolsize = PyTuple_GET_SIZE(pool);
2350 if (poolsize == 0) {
2351 lz->stopped = 1;
2352 Py_RETURN_NONE;
2353 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002354 /* clamp the index */
2355 if (index < 0)
2356 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002357 else if (index > poolsize-1)
2358 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002359 lz->indices[i] = index;
2360 }
2361
2362 result = PyTuple_New(n);
2363 if (!result)
2364 return NULL;
2365 for (i=0; i<n; i++) {
2366 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2367 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2368 Py_INCREF(element);
2369 PyTuple_SET_ITEM(result, i, element);
2370 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002371 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002372 Py_RETURN_NONE;
2373}
2374
2375static PyMethodDef product_methods[] = {
2376 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2377 reduce_doc},
2378 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2379 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002380 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2381 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002382 {NULL, NULL} /* sentinel */
2383};
2384
Christian Heimesc3f30c42008-02-22 16:37:40 +00002385PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002386"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002387\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002388Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002389For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2390The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2391cycle in a manner similar to an odometer (with the rightmost element changing\n\
2392on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002393To compute the product of an iterable with itself, specify the number\n\
2394of repetitions with the optional repeat keyword argument. For example,\n\
2395product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002396product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2397product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2398
2399static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002400 PyVarObject_HEAD_INIT(NULL, 0)
2401 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002402 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002403 0, /* tp_itemsize */
2404 /* methods */
2405 (destructor)product_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002406 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002407 0, /* tp_getattr */
2408 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002409 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002410 0, /* tp_repr */
2411 0, /* tp_as_number */
2412 0, /* tp_as_sequence */
2413 0, /* tp_as_mapping */
2414 0, /* tp_hash */
2415 0, /* tp_call */
2416 0, /* tp_str */
2417 PyObject_GenericGetAttr, /* tp_getattro */
2418 0, /* tp_setattro */
2419 0, /* tp_as_buffer */
2420 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2421 Py_TPFLAGS_BASETYPE, /* tp_flags */
2422 product_doc, /* tp_doc */
2423 (traverseproc)product_traverse, /* tp_traverse */
2424 0, /* tp_clear */
2425 0, /* tp_richcompare */
2426 0, /* tp_weaklistoffset */
2427 PyObject_SelfIter, /* tp_iter */
2428 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002429 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002430 0, /* tp_members */
2431 0, /* tp_getset */
2432 0, /* tp_base */
2433 0, /* tp_dict */
2434 0, /* tp_descr_get */
2435 0, /* tp_descr_set */
2436 0, /* tp_dictoffset */
2437 0, /* tp_init */
2438 0, /* tp_alloc */
2439 product_new, /* tp_new */
2440 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002441};
2442
2443
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002444/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002445
2446typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002447 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002448 PyObject *pool; /* input converted to a tuple */
2449 Py_ssize_t *indices; /* one index per result element */
2450 PyObject *result; /* most recently returned result tuple */
2451 Py_ssize_t r; /* size of result tuple */
2452 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002453} combinationsobject;
2454
2455static PyTypeObject combinations_type;
2456
Tal Einatc4bccd32018-09-12 00:49:13 +03002457
2458/*[clinic input]
2459@classmethod
2460itertools.combinations.__new__
2461 iterable: object
2462 r: Py_ssize_t
2463Return successive r-length combinations of elements in the iterable.
2464
2465combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)
2466[clinic start generated code]*/
2467
Christian Heimes380f7f22008-02-28 11:19:05 +00002468static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002469itertools_combinations_impl(PyTypeObject *type, PyObject *iterable,
2470 Py_ssize_t r)
2471/*[clinic end generated code: output=87a689b39c40039c input=06bede09e3da20f8]*/
Christian Heimes380f7f22008-02-28 11:19:05 +00002472{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002473 combinationsobject *co;
2474 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002475 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002476 Py_ssize_t *indices = NULL;
2477 Py_ssize_t i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002478
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002479 pool = PySequence_Tuple(iterable);
2480 if (pool == NULL)
2481 goto error;
2482 n = PyTuple_GET_SIZE(pool);
2483 if (r < 0) {
2484 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2485 goto error;
2486 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002487
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002488 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002489 if (indices == NULL) {
2490 PyErr_NoMemory();
2491 goto error;
2492 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002493
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002494 for (i=0 ; i<r ; i++)
2495 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002496
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002497 /* create combinationsobject structure */
2498 co = (combinationsobject *)type->tp_alloc(type, 0);
2499 if (co == NULL)
2500 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002501
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002502 co->pool = pool;
2503 co->indices = indices;
2504 co->result = NULL;
2505 co->r = r;
2506 co->stopped = r > n ? 1 : 0;
2507
2508 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002509
2510error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002511 if (indices != NULL)
2512 PyMem_Free(indices);
2513 Py_XDECREF(pool);
2514 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002515}
2516
2517static void
2518combinations_dealloc(combinationsobject *co)
2519{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002520 PyObject_GC_UnTrack(co);
2521 Py_XDECREF(co->pool);
2522 Py_XDECREF(co->result);
2523 if (co->indices != NULL)
2524 PyMem_Free(co->indices);
2525 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002526}
2527
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002528static PyObject *
2529combinations_sizeof(combinationsobject *co, void *unused)
2530{
2531 Py_ssize_t res;
2532
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002533 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002534 res += co->r * sizeof(Py_ssize_t);
2535 return PyLong_FromSsize_t(res);
2536}
2537
Christian Heimes380f7f22008-02-28 11:19:05 +00002538static int
2539combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2540{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002541 Py_VISIT(co->pool);
2542 Py_VISIT(co->result);
2543 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002544}
2545
2546static PyObject *
2547combinations_next(combinationsobject *co)
2548{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002549 PyObject *elem;
2550 PyObject *oldelem;
2551 PyObject *pool = co->pool;
2552 Py_ssize_t *indices = co->indices;
2553 PyObject *result = co->result;
2554 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2555 Py_ssize_t r = co->r;
2556 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002557
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002558 if (co->stopped)
2559 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002560
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002561 if (result == NULL) {
2562 /* On the first pass, initialize result tuple using the indices */
2563 result = PyTuple_New(r);
2564 if (result == NULL)
2565 goto empty;
2566 co->result = result;
2567 for (i=0; i<r ; i++) {
2568 index = indices[i];
2569 elem = PyTuple_GET_ITEM(pool, index);
2570 Py_INCREF(elem);
2571 PyTuple_SET_ITEM(result, i, elem);
2572 }
2573 } else {
2574 /* Copy the previous result tuple or re-use it if available */
2575 if (Py_REFCNT(result) > 1) {
2576 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002577 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002578 if (result == NULL)
2579 goto empty;
2580 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002581 Py_DECREF(old_result);
2582 }
2583 /* Now, we've got the only copy so we can update it in-place
2584 * CPython's empty tuple is a singleton and cached in
2585 * PyTuple's freelist.
2586 */
2587 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002588
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002589 /* Scan indices right-to-left until finding one that is not
2590 at its maximum (i + n - r). */
2591 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2592 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002593
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002594 /* If i is negative, then the indices are all at
2595 their maximum value and we're done. */
2596 if (i < 0)
2597 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002598
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002599 /* Increment the current index which we know is not at its
2600 maximum. Then move back to the right setting each index
2601 to its lowest possible value (one higher than the index
2602 to its left -- this maintains the sort order invariant). */
2603 indices[i]++;
2604 for (j=i+1 ; j<r ; j++)
2605 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002606
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002607 /* Update the result tuple for the new indices
2608 starting with i, the leftmost index that changed */
2609 for ( ; i<r ; i++) {
2610 index = indices[i];
2611 elem = PyTuple_GET_ITEM(pool, index);
2612 Py_INCREF(elem);
2613 oldelem = PyTuple_GET_ITEM(result, i);
2614 PyTuple_SET_ITEM(result, i, elem);
2615 Py_DECREF(oldelem);
2616 }
2617 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002618
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002619 Py_INCREF(result);
2620 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002621
2622empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002623 co->stopped = 1;
2624 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002625}
2626
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002627static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302628combinations_reduce(combinationsobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002629{
2630 if (lz->result == NULL) {
2631 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2632 } else if (lz->stopped) {
2633 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2634 } else {
2635 PyObject *indices;
2636 Py_ssize_t i;
2637
2638 /* we must pickle the indices and use them for setstate */
2639 indices = PyTuple_New(lz->r);
2640 if (!indices)
2641 return NULL;
2642 for (i=0; i<lz->r; i++)
2643 {
2644 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2645 if (!index) {
2646 Py_DECREF(indices);
2647 return NULL;
2648 }
2649 PyTuple_SET_ITEM(indices, i, index);
2650 }
2651
2652 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2653 }
2654}
2655
2656static PyObject *
2657combinations_setstate(combinationsobject *lz, PyObject *state)
2658{
2659 PyObject *result;
2660 Py_ssize_t i;
2661 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2662
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002663 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002664 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2665 return NULL;
2666 }
2667
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002668 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002669 Py_ssize_t max;
2670 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2671 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002672
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002673 if (index == -1 && PyErr_Occurred())
2674 return NULL; /* not an integer */
2675 max = i + n - lz->r;
2676 /* clamp the index (beware of negative max) */
2677 if (index > max)
2678 index = max;
2679 if (index < 0)
2680 index = 0;
2681 lz->indices[i] = index;
2682 }
2683
2684 result = PyTuple_New(lz->r);
2685 if (result == NULL)
2686 return NULL;
2687 for (i=0; i<lz->r; i++) {
2688 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2689 Py_INCREF(element);
2690 PyTuple_SET_ITEM(result, i, element);
2691 }
2692
Serhiy Storchaka48842712016-04-06 09:45:48 +03002693 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002694 Py_RETURN_NONE;
2695}
2696
2697static PyMethodDef combinations_methods[] = {
2698 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2699 reduce_doc},
2700 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2701 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002702 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2703 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002704 {NULL, NULL} /* sentinel */
2705};
2706
Christian Heimes380f7f22008-02-28 11:19:05 +00002707static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002708 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002709 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002710 sizeof(combinationsobject), /* tp_basicsize */
2711 0, /* tp_itemsize */
2712 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002713 (destructor)combinations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002714 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002715 0, /* tp_getattr */
2716 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002717 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002718 0, /* tp_repr */
2719 0, /* tp_as_number */
2720 0, /* tp_as_sequence */
2721 0, /* tp_as_mapping */
2722 0, /* tp_hash */
2723 0, /* tp_call */
2724 0, /* tp_str */
2725 PyObject_GenericGetAttr, /* tp_getattro */
2726 0, /* tp_setattro */
2727 0, /* tp_as_buffer */
2728 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2729 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03002730 itertools_combinations__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002731 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002732 0, /* tp_clear */
2733 0, /* tp_richcompare */
2734 0, /* tp_weaklistoffset */
2735 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002736 (iternextfunc)combinations_next, /* tp_iternext */
2737 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002738 0, /* tp_members */
2739 0, /* tp_getset */
2740 0, /* tp_base */
2741 0, /* tp_dict */
2742 0, /* tp_descr_get */
2743 0, /* tp_descr_set */
2744 0, /* tp_dictoffset */
2745 0, /* tp_init */
2746 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03002747 itertools_combinations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002748 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002749};
2750
2751
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002752/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002753
2754/* Equivalent to:
2755
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002756 def combinations_with_replacement(iterable, r):
2757 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2758 # number items returned: (n+r-1)! / r! / (n-1)!
2759 pool = tuple(iterable)
2760 n = len(pool)
2761 indices = [0] * r
2762 yield tuple(pool[i] for i in indices)
2763 while 1:
2764 for i in reversed(range(r)):
2765 if indices[i] != n - 1:
2766 break
2767 else:
2768 return
2769 indices[i:] = [indices[i] + 1] * (r - i)
2770 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002771
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002772 def combinations_with_replacement2(iterable, r):
2773 'Alternate version that filters from product()'
2774 pool = tuple(iterable)
2775 n = len(pool)
2776 for indices in product(range(n), repeat=r):
2777 if sorted(indices) == list(indices):
2778 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002779*/
2780typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002781 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002782 PyObject *pool; /* input converted to a tuple */
2783 Py_ssize_t *indices; /* one index per result element */
2784 PyObject *result; /* most recently returned result tuple */
2785 Py_ssize_t r; /* size of result tuple */
2786 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002787} cwrobject;
2788
2789static PyTypeObject cwr_type;
2790
Tal Einatc4bccd32018-09-12 00:49:13 +03002791/*[clinic input]
2792@classmethod
2793itertools.combinations_with_replacement.__new__
2794 iterable: object
2795 r: Py_ssize_t
2796Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.
2797
2798combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2799[clinic start generated code]*/
2800
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002801static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002802itertools_combinations_with_replacement_impl(PyTypeObject *type,
2803 PyObject *iterable,
2804 Py_ssize_t r)
2805/*[clinic end generated code: output=48b26856d4e659ca input=dc2a8c7ba785fad7]*/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002806{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002807 cwrobject *co;
2808 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002809 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002810 Py_ssize_t *indices = NULL;
2811 Py_ssize_t i;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002812
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002813 pool = PySequence_Tuple(iterable);
2814 if (pool == NULL)
2815 goto error;
2816 n = PyTuple_GET_SIZE(pool);
2817 if (r < 0) {
2818 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2819 goto error;
2820 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002821
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002822 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002823 if (indices == NULL) {
2824 PyErr_NoMemory();
2825 goto error;
2826 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002827
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002828 for (i=0 ; i<r ; i++)
2829 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002830
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002831 /* create cwrobject structure */
2832 co = (cwrobject *)type->tp_alloc(type, 0);
2833 if (co == NULL)
2834 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002835
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002836 co->pool = pool;
2837 co->indices = indices;
2838 co->result = NULL;
2839 co->r = r;
2840 co->stopped = !n && r;
2841
2842 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002843
2844error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002845 if (indices != NULL)
2846 PyMem_Free(indices);
2847 Py_XDECREF(pool);
2848 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002849}
2850
2851static void
2852cwr_dealloc(cwrobject *co)
2853{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002854 PyObject_GC_UnTrack(co);
2855 Py_XDECREF(co->pool);
2856 Py_XDECREF(co->result);
2857 if (co->indices != NULL)
2858 PyMem_Free(co->indices);
2859 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002860}
2861
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002862static PyObject *
2863cwr_sizeof(cwrobject *co, void *unused)
2864{
2865 Py_ssize_t res;
2866
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002867 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002868 res += co->r * sizeof(Py_ssize_t);
2869 return PyLong_FromSsize_t(res);
2870}
2871
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002872static int
2873cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2874{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002875 Py_VISIT(co->pool);
2876 Py_VISIT(co->result);
2877 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002878}
2879
2880static PyObject *
2881cwr_next(cwrobject *co)
2882{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002883 PyObject *elem;
2884 PyObject *oldelem;
2885 PyObject *pool = co->pool;
2886 Py_ssize_t *indices = co->indices;
2887 PyObject *result = co->result;
2888 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2889 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002890 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002891
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002892 if (co->stopped)
2893 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002894
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002895 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002896 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002897 result = PyTuple_New(r);
2898 if (result == NULL)
2899 goto empty;
2900 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002901 if (n > 0) {
2902 elem = PyTuple_GET_ITEM(pool, 0);
2903 for (i=0; i<r ; i++) {
2904 assert(indices[i] == 0);
2905 Py_INCREF(elem);
2906 PyTuple_SET_ITEM(result, i, elem);
2907 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002908 }
2909 } else {
2910 /* Copy the previous result tuple or re-use it if available */
2911 if (Py_REFCNT(result) > 1) {
2912 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002913 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002914 if (result == NULL)
2915 goto empty;
2916 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002917 Py_DECREF(old_result);
2918 }
2919 /* Now, we've got the only copy so we can update it in-place CPython's
2920 empty tuple is a singleton and cached in PyTuple's freelist. */
2921 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002922
Tim Peters9edb1682013-09-03 11:49:31 -05002923 /* Scan indices right-to-left until finding one that is not
2924 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002925 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2926 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002927
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002928 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002929 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002930 if (i < 0)
2931 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002932
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002933 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002934 maximum. Then set all to the right to the same value. */
2935 index = indices[i] + 1;
2936 assert(index < n);
2937 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002938 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002939 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002940 Py_INCREF(elem);
2941 oldelem = PyTuple_GET_ITEM(result, i);
2942 PyTuple_SET_ITEM(result, i, elem);
2943 Py_DECREF(oldelem);
2944 }
2945 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002946
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002947 Py_INCREF(result);
2948 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002949
2950empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002951 co->stopped = 1;
2952 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002953}
2954
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002955static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302956cwr_reduce(cwrobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002957{
2958 if (lz->result == NULL) {
2959 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2960 } else if (lz->stopped) {
2961 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2962 } else {
2963 PyObject *indices;
2964 Py_ssize_t i;
2965
2966 /* we must pickle the indices and use them for setstate */
2967 indices = PyTuple_New(lz->r);
2968 if (!indices)
2969 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002970 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002971 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2972 if (!index) {
2973 Py_DECREF(indices);
2974 return NULL;
2975 }
2976 PyTuple_SET_ITEM(indices, i, index);
2977 }
2978
2979 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2980 }
2981}
2982
2983static PyObject *
2984cwr_setstate(cwrobject *lz, PyObject *state)
2985{
2986 PyObject *result;
2987 Py_ssize_t n, i;
2988
2989 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2990 {
2991 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2992 return NULL;
2993 }
2994
2995 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002996 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002997 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2998 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002999
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003000 if (index < 0 && PyErr_Occurred())
3001 return NULL; /* not an integer */
3002 /* clamp the index */
3003 if (index < 0)
3004 index = 0;
3005 else if (index > n-1)
3006 index = n-1;
3007 lz->indices[i] = index;
3008 }
3009 result = PyTuple_New(lz->r);
3010 if (result == NULL)
3011 return NULL;
3012 for (i=0; i<lz->r; i++) {
3013 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
3014 Py_INCREF(element);
3015 PyTuple_SET_ITEM(result, i, element);
3016 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003017 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003018 Py_RETURN_NONE;
3019}
3020
3021static PyMethodDef cwr_methods[] = {
3022 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
3023 reduce_doc},
3024 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
3025 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003026 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
3027 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003028 {NULL, NULL} /* sentinel */
3029};
3030
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003031static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003032 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003033 "itertools.combinations_with_replacement", /* tp_name */
3034 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003035 0, /* tp_itemsize */
3036 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003037 (destructor)cwr_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003038 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003039 0, /* tp_getattr */
3040 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003041 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003042 0, /* tp_repr */
3043 0, /* tp_as_number */
3044 0, /* tp_as_sequence */
3045 0, /* tp_as_mapping */
3046 0, /* tp_hash */
3047 0, /* tp_call */
3048 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003049 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003050 0, /* tp_setattro */
3051 0, /* tp_as_buffer */
3052 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003053 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003054 itertools_combinations_with_replacement__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003055 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003056 0, /* tp_clear */
3057 0, /* tp_richcompare */
3058 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003059 PyObject_SelfIter, /* tp_iter */
3060 (iternextfunc)cwr_next, /* tp_iternext */
3061 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003062 0, /* tp_members */
3063 0, /* tp_getset */
3064 0, /* tp_base */
3065 0, /* tp_dict */
3066 0, /* tp_descr_get */
3067 0, /* tp_descr_set */
3068 0, /* tp_dictoffset */
3069 0, /* tp_init */
3070 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003071 itertools_combinations_with_replacement, /* tp_new */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003072 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003073};
3074
3075
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003076/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003077
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003078def permutations(iterable, r=None):
Sergey0078a0c2019-10-29 08:10:24 +03003079 # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
3080 # permutations(range(3)) --> 012 021 102 120 201 210
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003081 pool = tuple(iterable)
3082 n = len(pool)
3083 r = n if r is None else r
Sergey0078a0c2019-10-29 08:10:24 +03003084 if r > n:
3085 return
3086 indices = list(range(n))
3087 cycles = list(range(n, n-r, -1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003088 yield tuple(pool[i] for i in indices[:r])
3089 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003090 for i in reversed(range(r)):
3091 cycles[i] -= 1
3092 if cycles[i] == 0:
3093 indices[i:] = indices[i+1:] + indices[i:i+1]
3094 cycles[i] = n - i
3095 else:
3096 j = cycles[i]
3097 indices[i], indices[-j] = indices[-j], indices[i]
3098 yield tuple(pool[i] for i in indices[:r])
3099 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003100 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003101 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003102*/
3103
3104typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003105 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003106 PyObject *pool; /* input converted to a tuple */
3107 Py_ssize_t *indices; /* one index per element in the pool */
3108 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3109 PyObject *result; /* most recently returned result tuple */
3110 Py_ssize_t r; /* size of result tuple */
3111 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003112} permutationsobject;
3113
3114static PyTypeObject permutations_type;
3115
Tal Einatc4bccd32018-09-12 00:49:13 +03003116/*[clinic input]
3117@classmethod
3118itertools.permutations.__new__
3119 iterable: object
3120 r as robj: object = None
3121Return successive r-length permutations of elements in the iterable.
3122
3123permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
3124[clinic start generated code]*/
3125
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003126static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003127itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
3128 PyObject *robj)
3129/*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003130{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003131 permutationsobject *po;
3132 Py_ssize_t n;
3133 Py_ssize_t r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003134 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003135 Py_ssize_t *indices = NULL;
3136 Py_ssize_t *cycles = NULL;
3137 Py_ssize_t i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003138
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003139 pool = PySequence_Tuple(iterable);
3140 if (pool == NULL)
3141 goto error;
3142 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003143
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003144 r = n;
3145 if (robj != Py_None) {
3146 if (!PyLong_Check(robj)) {
3147 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3148 goto error;
3149 }
3150 r = PyLong_AsSsize_t(robj);
3151 if (r == -1 && PyErr_Occurred())
3152 goto error;
3153 }
3154 if (r < 0) {
3155 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3156 goto error;
3157 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003158
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003159 indices = PyMem_New(Py_ssize_t, n);
3160 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003161 if (indices == NULL || cycles == NULL) {
3162 PyErr_NoMemory();
3163 goto error;
3164 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003165
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003166 for (i=0 ; i<n ; i++)
3167 indices[i] = i;
3168 for (i=0 ; i<r ; i++)
3169 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003170
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003171 /* create permutationsobject structure */
3172 po = (permutationsobject *)type->tp_alloc(type, 0);
3173 if (po == NULL)
3174 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003175
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003176 po->pool = pool;
3177 po->indices = indices;
3178 po->cycles = cycles;
3179 po->result = NULL;
3180 po->r = r;
3181 po->stopped = r > n ? 1 : 0;
3182
3183 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003184
3185error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003186 if (indices != NULL)
3187 PyMem_Free(indices);
3188 if (cycles != NULL)
3189 PyMem_Free(cycles);
3190 Py_XDECREF(pool);
3191 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003192}
3193
3194static void
3195permutations_dealloc(permutationsobject *po)
3196{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003197 PyObject_GC_UnTrack(po);
3198 Py_XDECREF(po->pool);
3199 Py_XDECREF(po->result);
3200 PyMem_Free(po->indices);
3201 PyMem_Free(po->cycles);
3202 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003203}
3204
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003205static PyObject *
3206permutations_sizeof(permutationsobject *po, void *unused)
3207{
3208 Py_ssize_t res;
3209
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003210 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003211 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3212 res += po->r * sizeof(Py_ssize_t);
3213 return PyLong_FromSsize_t(res);
3214}
3215
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003216static int
3217permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3218{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003219 Py_VISIT(po->pool);
3220 Py_VISIT(po->result);
3221 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003222}
3223
3224static PyObject *
3225permutations_next(permutationsobject *po)
3226{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003227 PyObject *elem;
3228 PyObject *oldelem;
3229 PyObject *pool = po->pool;
3230 Py_ssize_t *indices = po->indices;
3231 Py_ssize_t *cycles = po->cycles;
3232 PyObject *result = po->result;
3233 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3234 Py_ssize_t r = po->r;
3235 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003236
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003237 if (po->stopped)
3238 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003239
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003240 if (result == NULL) {
3241 /* On the first pass, initialize result tuple using the indices */
3242 result = PyTuple_New(r);
3243 if (result == NULL)
3244 goto empty;
3245 po->result = result;
3246 for (i=0; i<r ; i++) {
3247 index = indices[i];
3248 elem = PyTuple_GET_ITEM(pool, index);
3249 Py_INCREF(elem);
3250 PyTuple_SET_ITEM(result, i, elem);
3251 }
3252 } else {
3253 if (n == 0)
3254 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003255
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003256 /* Copy the previous result tuple or re-use it if available */
3257 if (Py_REFCNT(result) > 1) {
3258 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05003259 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003260 if (result == NULL)
3261 goto empty;
3262 po->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003263 Py_DECREF(old_result);
3264 }
3265 /* Now, we've got the only copy so we can update it in-place */
3266 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003267
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003268 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3269 for (i=r-1 ; i>=0 ; i--) {
3270 cycles[i] -= 1;
3271 if (cycles[i] == 0) {
3272 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3273 index = indices[i];
3274 for (j=i ; j<n-1 ; j++)
3275 indices[j] = indices[j+1];
3276 indices[n-1] = index;
3277 cycles[i] = n - i;
3278 } else {
3279 j = cycles[i];
3280 index = indices[i];
3281 indices[i] = indices[n-j];
3282 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003283
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003284 for (k=i; k<r ; k++) {
3285 /* start with i, the leftmost element that changed */
3286 /* yield tuple(pool[k] for k in indices[:r]) */
3287 index = indices[k];
3288 elem = PyTuple_GET_ITEM(pool, index);
3289 Py_INCREF(elem);
3290 oldelem = PyTuple_GET_ITEM(result, k);
3291 PyTuple_SET_ITEM(result, k, elem);
3292 Py_DECREF(oldelem);
3293 }
3294 break;
3295 }
3296 }
3297 /* If i is negative, then the cycles have all
3298 rolled-over and we're done. */
3299 if (i < 0)
3300 goto empty;
3301 }
3302 Py_INCREF(result);
3303 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003304
3305empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003306 po->stopped = 1;
3307 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003308}
3309
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003310static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303311permutations_reduce(permutationsobject *po, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003312{
3313 if (po->result == NULL) {
3314 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3315 } else if (po->stopped) {
3316 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3317 } else {
3318 PyObject *indices=NULL, *cycles=NULL;
3319 Py_ssize_t n, i;
3320
3321 /* we must pickle the indices and cycles and use them for setstate */
3322 n = PyTuple_GET_SIZE(po->pool);
3323 indices = PyTuple_New(n);
3324 if (indices == NULL)
3325 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003326 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003327 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3328 if (!index)
3329 goto err;
3330 PyTuple_SET_ITEM(indices, i, index);
3331 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003332
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003333 cycles = PyTuple_New(po->r);
3334 if (cycles == NULL)
3335 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003336 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003337 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3338 if (!index)
3339 goto err;
3340 PyTuple_SET_ITEM(cycles, i, index);
3341 }
3342 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3343 po->pool, po->r,
3344 indices, cycles);
3345 err:
3346 Py_XDECREF(indices);
3347 Py_XDECREF(cycles);
3348 return NULL;
3349 }
3350}
3351
3352static PyObject *
3353permutations_setstate(permutationsobject *po, PyObject *state)
3354{
3355 PyObject *indices, *cycles, *result;
3356 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003357
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003358 if (!PyTuple_Check(state)) {
3359 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3360 return NULL;
3361 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003362 if (!PyArg_ParseTuple(state, "O!O!",
3363 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003364 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003365 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003366 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003367
3368 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003369 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003370 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3371 return NULL;
3372 }
3373
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003374 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003375 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3376 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3377 if (index < 0 && PyErr_Occurred())
3378 return NULL; /* not an integer */
3379 /* clamp the index */
3380 if (index < 0)
3381 index = 0;
3382 else if (index > n-1)
3383 index = n-1;
3384 po->indices[i] = index;
3385 }
3386
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003387 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003388 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3389 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3390 if (index < 0 && PyErr_Occurred())
3391 return NULL; /* not an integer */
3392 if (index < 1)
3393 index = 1;
3394 else if (index > n-i)
3395 index = n-i;
3396 po->cycles[i] = index;
3397 }
3398 result = PyTuple_New(po->r);
3399 if (result == NULL)
3400 return NULL;
3401 for (i=0; i<po->r; i++) {
3402 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3403 Py_INCREF(element);
3404 PyTuple_SET_ITEM(result, i, element);
3405 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003406 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003407 Py_RETURN_NONE;
3408}
3409
3410static PyMethodDef permuations_methods[] = {
3411 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3412 reduce_doc},
3413 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3414 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003415 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3416 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003417 {NULL, NULL} /* sentinel */
3418};
3419
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003420static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003421 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003422 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003423 sizeof(permutationsobject), /* tp_basicsize */
3424 0, /* tp_itemsize */
3425 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003426 (destructor)permutations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003427 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003428 0, /* tp_getattr */
3429 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003430 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003431 0, /* tp_repr */
3432 0, /* tp_as_number */
3433 0, /* tp_as_sequence */
3434 0, /* tp_as_mapping */
3435 0, /* tp_hash */
3436 0, /* tp_call */
3437 0, /* tp_str */
3438 PyObject_GenericGetAttr, /* tp_getattro */
3439 0, /* tp_setattro */
3440 0, /* tp_as_buffer */
3441 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3442 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003443 itertools_permutations__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003444 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003445 0, /* tp_clear */
3446 0, /* tp_richcompare */
3447 0, /* tp_weaklistoffset */
3448 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003449 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003450 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003451 0, /* tp_members */
3452 0, /* tp_getset */
3453 0, /* tp_base */
3454 0, /* tp_dict */
3455 0, /* tp_descr_get */
3456 0, /* tp_descr_set */
3457 0, /* tp_dictoffset */
3458 0, /* tp_init */
3459 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003460 itertools_permutations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003461 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003462};
3463
Tal Einatc4bccd32018-09-12 00:49:13 +03003464
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003465/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003466
3467typedef struct {
3468 PyObject_HEAD
3469 PyObject *total;
3470 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003471 PyObject *binop;
Lisa Roach9718b592018-09-23 17:34:59 -07003472 PyObject *initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003473} accumulateobject;
3474
3475static PyTypeObject accumulate_type;
3476
Tal Einatc4bccd32018-09-12 00:49:13 +03003477/*[clinic input]
3478@classmethod
3479itertools.accumulate.__new__
3480 iterable: object
3481 func as binop: object = None
Lisa Roach9718b592018-09-23 17:34:59 -07003482 *
3483 initial: object = None
Tal Einatc4bccd32018-09-12 00:49:13 +03003484Return series of accumulated sums (or other binary function results).
3485[clinic start generated code]*/
3486
Raymond Hettinger482ba772010-12-01 22:48:00 +00003487static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003488itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
Lisa Roach9718b592018-09-23 17:34:59 -07003489 PyObject *binop, PyObject *initial)
3490/*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003491{
Raymond Hettinger482ba772010-12-01 22:48:00 +00003492 PyObject *it;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003493 accumulateobject *lz;
3494
Raymond Hettinger482ba772010-12-01 22:48:00 +00003495 /* Get iterator. */
3496 it = PyObject_GetIter(iterable);
3497 if (it == NULL)
3498 return NULL;
3499
Raymond Hettinger482ba772010-12-01 22:48:00 +00003500 /* create accumulateobject structure */
3501 lz = (accumulateobject *)type->tp_alloc(type, 0);
3502 if (lz == NULL) {
3503 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003504 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003505 }
3506
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003507 if (binop != Py_None) {
3508 Py_XINCREF(binop);
3509 lz->binop = binop;
3510 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003511 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003512 lz->it = it;
Lisa Roach9718b592018-09-23 17:34:59 -07003513 Py_XINCREF(initial);
3514 lz->initial = initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003515 return (PyObject *)lz;
3516}
3517
3518static void
3519accumulate_dealloc(accumulateobject *lz)
3520{
3521 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003522 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003523 Py_XDECREF(lz->total);
3524 Py_XDECREF(lz->it);
Lisa Roach9718b592018-09-23 17:34:59 -07003525 Py_XDECREF(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003526 Py_TYPE(lz)->tp_free(lz);
3527}
3528
3529static int
3530accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3531{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003532 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003533 Py_VISIT(lz->it);
3534 Py_VISIT(lz->total);
Lisa Roach9718b592018-09-23 17:34:59 -07003535 Py_VISIT(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003536 return 0;
3537}
3538
3539static PyObject *
3540accumulate_next(accumulateobject *lz)
3541{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003542 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003543
Lisa Roach9718b592018-09-23 17:34:59 -07003544 if (lz->initial != Py_None) {
3545 lz->total = lz->initial;
3546 Py_INCREF(Py_None);
3547 lz->initial = Py_None;
3548 Py_INCREF(lz->total);
3549 return lz->total;
3550 }
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003551 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003552 if (val == NULL)
3553 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003554
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003555 if (lz->total == NULL) {
3556 Py_INCREF(val);
3557 lz->total = val;
3558 return lz->total;
3559 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003560
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003561 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003562 newtotal = PyNumber_Add(lz->total, val);
3563 else
3564 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003565 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003566 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003567 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003568
Raymond Hettinger482ba772010-12-01 22:48:00 +00003569 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003570 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003571 return newtotal;
3572}
3573
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003574static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303575accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003576{
Lisa Roach9718b592018-09-23 17:34:59 -07003577 if (lz->initial != Py_None) {
3578 PyObject *it;
3579
3580 assert(lz->total == NULL);
3581 if (PyType_Ready(&chain_type) < 0)
3582 return NULL;
3583 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3584 lz->initial, lz->it);
3585 if (it == NULL)
3586 return NULL;
3587 return Py_BuildValue("O(NO)O", Py_TYPE(lz),
3588 it, lz->binop?lz->binop:Py_None, Py_None);
3589 }
Serhiy Storchakad5516252016-03-06 14:00:45 +02003590 if (lz->total == Py_None) {
3591 PyObject *it;
3592
3593 if (PyType_Ready(&chain_type) < 0)
3594 return NULL;
3595 if (PyType_Ready(&islice_type) < 0)
3596 return NULL;
3597 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3598 lz->total, lz->it);
3599 if (it == NULL)
3600 return NULL;
3601 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3602 it, lz->binop ? lz->binop : Py_None);
3603 if (it == NULL)
3604 return NULL;
3605 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3606 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003607 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3608 lz->it, lz->binop?lz->binop:Py_None,
3609 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003610}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003611
3612static PyObject *
3613accumulate_setstate(accumulateobject *lz, PyObject *state)
3614{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003615 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003616 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003617 Py_RETURN_NONE;
3618}
3619
3620static PyMethodDef accumulate_methods[] = {
3621 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3622 reduce_doc},
3623 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3624 setstate_doc},
3625 {NULL, NULL} /* sentinel */
3626};
3627
Raymond Hettinger482ba772010-12-01 22:48:00 +00003628static PyTypeObject accumulate_type = {
3629 PyVarObject_HEAD_INIT(NULL, 0)
3630 "itertools.accumulate", /* tp_name */
3631 sizeof(accumulateobject), /* tp_basicsize */
3632 0, /* tp_itemsize */
3633 /* methods */
3634 (destructor)accumulate_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003635 0, /* tp_vectorcall_offset */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003636 0, /* tp_getattr */
3637 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003638 0, /* tp_as_async */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003639 0, /* tp_repr */
3640 0, /* tp_as_number */
3641 0, /* tp_as_sequence */
3642 0, /* tp_as_mapping */
3643 0, /* tp_hash */
3644 0, /* tp_call */
3645 0, /* tp_str */
3646 PyObject_GenericGetAttr, /* tp_getattro */
3647 0, /* tp_setattro */
3648 0, /* tp_as_buffer */
3649 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3650 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003651 itertools_accumulate__doc__, /* tp_doc */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003652 (traverseproc)accumulate_traverse, /* tp_traverse */
3653 0, /* tp_clear */
3654 0, /* tp_richcompare */
3655 0, /* tp_weaklistoffset */
3656 PyObject_SelfIter, /* tp_iter */
3657 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003658 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003659 0, /* tp_members */
3660 0, /* tp_getset */
3661 0, /* tp_base */
3662 0, /* tp_dict */
3663 0, /* tp_descr_get */
3664 0, /* tp_descr_set */
3665 0, /* tp_dictoffset */
3666 0, /* tp_init */
3667 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003668 itertools_accumulate, /* tp_new */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003669 PyObject_GC_Del, /* tp_free */
3670};
3671
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003672
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003673/* compress object ************************************************************/
3674
3675/* Equivalent to:
3676
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003677 def compress(data, selectors):
3678 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3679 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003680*/
3681
3682typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003683 PyObject_HEAD
3684 PyObject *data;
3685 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003686} compressobject;
3687
3688static PyTypeObject compress_type;
3689
Tal Einatc4bccd32018-09-12 00:49:13 +03003690/*[clinic input]
3691@classmethod
3692itertools.compress.__new__
3693 data as seq1: object
3694 selectors as seq2: object
3695Return data elements corresponding to true selector elements.
3696
3697Forms a shorter iterator from selected data elements using the selectors to
3698choose the data elements.
3699[clinic start generated code]*/
3700
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003701static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003702itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
3703/*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003704{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003705 PyObject *data=NULL, *selectors=NULL;
3706 compressobject *lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003707
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003708 data = PyObject_GetIter(seq1);
3709 if (data == NULL)
3710 goto fail;
3711 selectors = PyObject_GetIter(seq2);
3712 if (selectors == NULL)
3713 goto fail;
3714
3715 /* create compressobject structure */
3716 lz = (compressobject *)type->tp_alloc(type, 0);
3717 if (lz == NULL)
3718 goto fail;
3719 lz->data = data;
3720 lz->selectors = selectors;
3721 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003722
3723fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003724 Py_XDECREF(data);
3725 Py_XDECREF(selectors);
3726 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003727}
3728
3729static void
3730compress_dealloc(compressobject *lz)
3731{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003732 PyObject_GC_UnTrack(lz);
3733 Py_XDECREF(lz->data);
3734 Py_XDECREF(lz->selectors);
3735 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003736}
3737
3738static int
3739compress_traverse(compressobject *lz, visitproc visit, void *arg)
3740{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003741 Py_VISIT(lz->data);
3742 Py_VISIT(lz->selectors);
3743 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003744}
3745
3746static PyObject *
3747compress_next(compressobject *lz)
3748{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003749 PyObject *data = lz->data, *selectors = lz->selectors;
3750 PyObject *datum, *selector;
3751 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3752 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3753 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003754
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003755 while (1) {
3756 /* Steps: get datum, get selector, evaluate selector.
3757 Order is important (to match the pure python version
3758 in terms of which input gets a chance to raise an
3759 exception first).
3760 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003761
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003762 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003763 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003764 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003765
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003766 selector = selectornext(selectors);
3767 if (selector == NULL) {
3768 Py_DECREF(datum);
3769 return NULL;
3770 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003771
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003772 ok = PyObject_IsTrue(selector);
3773 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003774 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003775 return datum;
3776 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003777 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003778 return NULL;
3779 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003780}
3781
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003782static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303783compress_reduce(compressobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003784{
3785 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3786 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003787}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003788
3789static PyMethodDef compress_methods[] = {
3790 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3791 reduce_doc},
3792 {NULL, NULL} /* sentinel */
3793};
3794
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003795static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003796 PyVarObject_HEAD_INIT(NULL, 0)
3797 "itertools.compress", /* tp_name */
3798 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003799 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003800 /* methods */
3801 (destructor)compress_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003802 0, /* tp_vectorcall_offset */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003803 0, /* tp_getattr */
3804 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003805 0, /* tp_as_async */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003806 0, /* tp_repr */
3807 0, /* tp_as_number */
3808 0, /* tp_as_sequence */
3809 0, /* tp_as_mapping */
3810 0, /* tp_hash */
3811 0, /* tp_call */
3812 0, /* tp_str */
3813 PyObject_GenericGetAttr, /* tp_getattro */
3814 0, /* tp_setattro */
3815 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003816 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003817 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003818 itertools_compress__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003819 (traverseproc)compress_traverse, /* tp_traverse */
3820 0, /* tp_clear */
3821 0, /* tp_richcompare */
3822 0, /* tp_weaklistoffset */
3823 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003824 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003825 compress_methods, /* tp_methods */
3826 0, /* tp_members */
3827 0, /* tp_getset */
3828 0, /* tp_base */
3829 0, /* tp_dict */
3830 0, /* tp_descr_get */
3831 0, /* tp_descr_set */
3832 0, /* tp_dictoffset */
3833 0, /* tp_init */
3834 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003835 itertools_compress, /* tp_new */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003836 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003837};
3838
3839
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003840/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003841
3842typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003843 PyObject_HEAD
3844 PyObject *func;
3845 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003846} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003847
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003848static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003849
Tal Einatc4bccd32018-09-12 00:49:13 +03003850/*[clinic input]
3851@classmethod
3852itertools.filterfalse.__new__
3853 function as func: object
3854 iterable as seq: object
3855 /
3856Return those items of iterable for which function(item) is false.
3857
3858If function is None, return the items that are false.
3859[clinic start generated code]*/
3860
Raymond Hettinger60eca932003-02-09 06:40:58 +00003861static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003862itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
3863/*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003864{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003865 PyObject *it;
3866 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003867
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003868 /* Get iterator. */
3869 it = PyObject_GetIter(seq);
3870 if (it == NULL)
3871 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003872
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003873 /* create filterfalseobject structure */
3874 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3875 if (lz == NULL) {
3876 Py_DECREF(it);
3877 return NULL;
3878 }
3879 Py_INCREF(func);
3880 lz->func = func;
3881 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003882
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003883 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003884}
3885
3886static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003887filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003888{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003889 PyObject_GC_UnTrack(lz);
3890 Py_XDECREF(lz->func);
3891 Py_XDECREF(lz->it);
3892 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003893}
3894
3895static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003896filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003897{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003898 Py_VISIT(lz->it);
3899 Py_VISIT(lz->func);
3900 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003901}
3902
3903static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003904filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003905{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003906 PyObject *item;
3907 PyObject *it = lz->it;
3908 long ok;
3909 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003910
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003911 iternext = *Py_TYPE(it)->tp_iternext;
3912 for (;;) {
3913 item = iternext(it);
3914 if (item == NULL)
3915 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003916
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003917 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3918 ok = PyObject_IsTrue(item);
3919 } else {
3920 PyObject *good;
Petr Viktorinffd97532020-02-11 17:46:57 +01003921 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003922 if (good == NULL) {
3923 Py_DECREF(item);
3924 return NULL;
3925 }
3926 ok = PyObject_IsTrue(good);
3927 Py_DECREF(good);
3928 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003929 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003930 return item;
3931 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003932 if (ok < 0)
3933 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003934 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003935}
3936
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003937static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303938filterfalse_reduce(filterfalseobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003939{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003940 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3941}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003942
3943static PyMethodDef filterfalse_methods[] = {
3944 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3945 reduce_doc},
3946 {NULL, NULL} /* sentinel */
3947};
3948
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003949static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003950 PyVarObject_HEAD_INIT(NULL, 0)
3951 "itertools.filterfalse", /* tp_name */
3952 sizeof(filterfalseobject), /* tp_basicsize */
3953 0, /* tp_itemsize */
3954 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003955 (destructor)filterfalse_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003956 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003957 0, /* tp_getattr */
3958 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003959 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003960 0, /* tp_repr */
3961 0, /* tp_as_number */
3962 0, /* tp_as_sequence */
3963 0, /* tp_as_mapping */
3964 0, /* tp_hash */
3965 0, /* tp_call */
3966 0, /* tp_str */
3967 PyObject_GenericGetAttr, /* tp_getattro */
3968 0, /* tp_setattro */
3969 0, /* tp_as_buffer */
3970 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3971 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003972 itertools_filterfalse__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003973 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003974 0, /* tp_clear */
3975 0, /* tp_richcompare */
3976 0, /* tp_weaklistoffset */
3977 PyObject_SelfIter, /* tp_iter */
3978 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003979 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003980 0, /* tp_members */
3981 0, /* tp_getset */
3982 0, /* tp_base */
3983 0, /* tp_dict */
3984 0, /* tp_descr_get */
3985 0, /* tp_descr_set */
3986 0, /* tp_dictoffset */
3987 0, /* tp_init */
3988 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003989 itertools_filterfalse, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003990 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003991};
3992
3993
3994/* count object ************************************************************/
3995
3996typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003997 PyObject_HEAD
3998 Py_ssize_t cnt;
3999 PyObject *long_cnt;
4000 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004001} countobject;
4002
Raymond Hettinger30729212009-02-12 06:28:27 +00004003/* Counting logic and invariants:
4004
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004005fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00004006
Serhiy Storchaka483405b2015-02-17 10:14:30 +02004007 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004008 Advances with: cnt += 1
4009 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00004010
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004011slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00004012
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004013 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
4014 All counting is done with python objects (no overflows or underflows).
4015 Advances with: long_cnt += long_step
4016 Step may be zero -- effectively a slow version of repeat(cnt).
4017 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00004018*/
4019
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004020static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004021
Tal Einatc4bccd32018-09-12 00:49:13 +03004022/*[clinic input]
4023@classmethod
4024itertools.count.__new__
4025 start as long_cnt: object(c_default="NULL") = 0
4026 step as long_step: object(c_default="NULL") = 1
4027Return a count object whose .__next__() method returns consecutive values.
4028
4029Equivalent to:
4030 def count(firstval=0, step=1):
4031 x = firstval
4032 while 1:
4033 yield x
4034 x += step
4035[clinic start generated code]*/
4036
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004037static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03004038itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
4039 PyObject *long_step)
4040/*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004041{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004042 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004043 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004044 Py_ssize_t cnt = 0;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00004045 long step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004046
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004047 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
4048 (long_step != NULL && !PyNumber_Check(long_step))) {
4049 PyErr_SetString(PyExc_TypeError, "a number is required");
4050 return NULL;
4051 }
Raymond Hettinger30729212009-02-12 06:28:27 +00004052
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004053 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
4054 (long_step == NULL || PyLong_Check(long_step));
4055
4056 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004057 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004058 if (fast_mode) {
4059 assert(PyLong_Check(long_cnt));
4060 cnt = PyLong_AsSsize_t(long_cnt);
4061 if (cnt == -1 && PyErr_Occurred()) {
4062 PyErr_Clear();
4063 fast_mode = 0;
4064 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004065 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004066 } else {
4067 cnt = 0;
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004068 long_cnt = _PyLong_Zero;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004069 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004070 Py_INCREF(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004071
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004072 /* If not specified, step defaults to 1 */
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004073 if (long_step == NULL)
4074 long_step = _PyLong_One;
4075 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004076
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004077 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004078
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004079 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004080 if (fast_mode) {
4081 assert(PyLong_Check(long_step));
4082 step = PyLong_AsLong(long_step);
4083 if (step != 1) {
4084 fast_mode = 0;
4085 if (step == -1 && PyErr_Occurred())
4086 PyErr_Clear();
4087 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004088 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004089
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004090 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004091 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004092 else
4093 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004094
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004095 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4096 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4097 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004098 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004099
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004100 /* create countobject structure */
4101 lz = (countobject *)type->tp_alloc(type, 0);
4102 if (lz == NULL) {
4103 Py_XDECREF(long_cnt);
Zackery Spytz0523c392019-03-26 00:05:29 -06004104 Py_DECREF(long_step);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004105 return NULL;
4106 }
4107 lz->cnt = cnt;
4108 lz->long_cnt = long_cnt;
4109 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004110
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004111 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004112}
4113
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004114static void
4115count_dealloc(countobject *lz)
4116{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004117 PyObject_GC_UnTrack(lz);
4118 Py_XDECREF(lz->long_cnt);
4119 Py_XDECREF(lz->long_step);
4120 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004121}
4122
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004123static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004124count_traverse(countobject *lz, visitproc visit, void *arg)
4125{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004126 Py_VISIT(lz->long_cnt);
4127 Py_VISIT(lz->long_step);
4128 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004129}
4130
4131static PyObject *
4132count_nextlong(countobject *lz)
4133{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004134 PyObject *long_cnt;
4135 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004136
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004137 long_cnt = lz->long_cnt;
4138 if (long_cnt == NULL) {
4139 /* Switch to slow_mode */
4140 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4141 if (long_cnt == NULL)
4142 return NULL;
4143 }
4144 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004145
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004146 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4147 if (stepped_up == NULL)
4148 return NULL;
4149 lz->long_cnt = stepped_up;
4150 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004151}
4152
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004153static PyObject *
4154count_next(countobject *lz)
4155{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004156 if (lz->cnt == PY_SSIZE_T_MAX)
4157 return count_nextlong(lz);
4158 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004159}
4160
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004161static PyObject *
4162count_repr(countobject *lz)
4163{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004164 if (lz->cnt != PY_SSIZE_T_MAX)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004165 return PyUnicode_FromFormat("%s(%zd)",
4166 _PyType_Name(Py_TYPE(lz)), lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004167
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004168 if (PyLong_Check(lz->long_step)) {
4169 long step = PyLong_AsLong(lz->long_step);
4170 if (step == -1 && PyErr_Occurred()) {
4171 PyErr_Clear();
4172 }
4173 if (step == 1) {
4174 /* Don't display step when it is an integer equal to 1 */
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004175 return PyUnicode_FromFormat("%s(%R)",
4176 _PyType_Name(Py_TYPE(lz)),
4177 lz->long_cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004178 }
4179 }
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004180 return PyUnicode_FromFormat("%s(%R, %R)",
4181 _PyType_Name(Py_TYPE(lz)),
4182 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004183}
4184
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004185static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304186count_reduce(countobject *lz, PyObject *Py_UNUSED(ignored))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004187{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004188 if (lz->cnt == PY_SSIZE_T_MAX)
4189 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4190 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004191}
4192
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004193static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004194 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004195 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004196 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004197};
4198
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004199static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004200 PyVarObject_HEAD_INIT(NULL, 0)
4201 "itertools.count", /* tp_name */
4202 sizeof(countobject), /* tp_basicsize */
4203 0, /* tp_itemsize */
4204 /* methods */
4205 (destructor)count_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004206 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004207 0, /* tp_getattr */
4208 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004209 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004210 (reprfunc)count_repr, /* tp_repr */
4211 0, /* tp_as_number */
4212 0, /* tp_as_sequence */
4213 0, /* tp_as_mapping */
4214 0, /* tp_hash */
4215 0, /* tp_call */
4216 0, /* tp_str */
4217 PyObject_GenericGetAttr, /* tp_getattro */
4218 0, /* tp_setattro */
4219 0, /* tp_as_buffer */
4220 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004221 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03004222 itertools_count__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004223 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004224 0, /* tp_clear */
4225 0, /* tp_richcompare */
4226 0, /* tp_weaklistoffset */
4227 PyObject_SelfIter, /* tp_iter */
4228 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004229 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004230 0, /* tp_members */
4231 0, /* tp_getset */
4232 0, /* tp_base */
4233 0, /* tp_dict */
4234 0, /* tp_descr_get */
4235 0, /* tp_descr_set */
4236 0, /* tp_dictoffset */
4237 0, /* tp_init */
4238 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03004239 itertools_count, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004240 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004241};
4242
4243
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004244/* repeat object ************************************************************/
4245
4246typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004247 PyObject_HEAD
4248 PyObject *element;
4249 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004250} repeatobject;
4251
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004252static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004253
4254static PyObject *
4255repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4256{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004257 repeatobject *ro;
4258 PyObject *element;
Raymond Hettingeraf464502019-11-09 20:28:31 -08004259 Py_ssize_t cnt = -1, n_args;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004260 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004261
Raymond Hettingeraf464502019-11-09 20:28:31 -08004262 n_args = PyTuple_GET_SIZE(args);
4263 if (kwds != NULL)
4264 n_args += PyDict_GET_SIZE(kwds);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004265 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4266 &element, &cnt))
4267 return NULL;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004268 /* Does user supply times argument? */
Raymond Hettingeraf464502019-11-09 20:28:31 -08004269 if (n_args == 2 && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004270 cnt = 0;
4271
4272 ro = (repeatobject *)type->tp_alloc(type, 0);
4273 if (ro == NULL)
4274 return NULL;
4275 Py_INCREF(element);
4276 ro->element = element;
4277 ro->cnt = cnt;
4278 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004279}
4280
4281static void
4282repeat_dealloc(repeatobject *ro)
4283{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004284 PyObject_GC_UnTrack(ro);
4285 Py_XDECREF(ro->element);
4286 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004287}
4288
4289static int
4290repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4291{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004292 Py_VISIT(ro->element);
4293 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004294}
4295
4296static PyObject *
4297repeat_next(repeatobject *ro)
4298{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004299 if (ro->cnt == 0)
4300 return NULL;
4301 if (ro->cnt > 0)
4302 ro->cnt--;
4303 Py_INCREF(ro->element);
4304 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004305}
4306
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004307static PyObject *
4308repeat_repr(repeatobject *ro)
4309{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004310 if (ro->cnt == -1)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004311 return PyUnicode_FromFormat("%s(%R)",
4312 _PyType_Name(Py_TYPE(ro)), ro->element);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004313 else
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004314 return PyUnicode_FromFormat("%s(%R, %zd)",
4315 _PyType_Name(Py_TYPE(ro)), ro->element,
4316 ro->cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004317}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004318
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004319static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304320repeat_len(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004321{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004322 if (ro->cnt == -1) {
4323 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4324 return NULL;
4325 }
4326 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004327}
4328
Armin Rigof5b3e362006-02-11 21:32:43 +00004329PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004330
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004331static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304332repeat_reduce(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004333{
4334 /* unpickle this so that a new repeat iterator is constructed with an
4335 * object, then call __setstate__ on it to set cnt
4336 */
4337 if (ro->cnt >= 0)
4338 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4339 else
4340 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4341}
4342
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004343static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004344 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004345 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004346 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004347};
4348
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004349PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004350"repeat(object [,times]) -> create an iterator which returns the object\n\
4351for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004352endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004353
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004354static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004355 PyVarObject_HEAD_INIT(NULL, 0)
4356 "itertools.repeat", /* tp_name */
4357 sizeof(repeatobject), /* tp_basicsize */
4358 0, /* tp_itemsize */
4359 /* methods */
4360 (destructor)repeat_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004361 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004362 0, /* tp_getattr */
4363 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004364 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004365 (reprfunc)repeat_repr, /* tp_repr */
4366 0, /* tp_as_number */
4367 0, /* tp_as_sequence */
4368 0, /* tp_as_mapping */
4369 0, /* tp_hash */
4370 0, /* tp_call */
4371 0, /* tp_str */
4372 PyObject_GenericGetAttr, /* tp_getattro */
4373 0, /* tp_setattro */
4374 0, /* tp_as_buffer */
4375 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4376 Py_TPFLAGS_BASETYPE, /* tp_flags */
4377 repeat_doc, /* tp_doc */
4378 (traverseproc)repeat_traverse, /* tp_traverse */
4379 0, /* tp_clear */
4380 0, /* tp_richcompare */
4381 0, /* tp_weaklistoffset */
4382 PyObject_SelfIter, /* tp_iter */
4383 (iternextfunc)repeat_next, /* tp_iternext */
4384 repeat_methods, /* tp_methods */
4385 0, /* tp_members */
4386 0, /* tp_getset */
4387 0, /* tp_base */
4388 0, /* tp_dict */
4389 0, /* tp_descr_get */
4390 0, /* tp_descr_set */
4391 0, /* tp_dictoffset */
4392 0, /* tp_init */
4393 0, /* tp_alloc */
4394 repeat_new, /* tp_new */
4395 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004396};
4397
Tal Einatc4bccd32018-09-12 00:49:13 +03004398
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004399/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004400
4401typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004402 PyObject_HEAD
4403 Py_ssize_t tuplesize;
4404 Py_ssize_t numactive;
4405 PyObject *ittuple; /* tuple of iterators */
4406 PyObject *result;
4407 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004408} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004409
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004410static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004411
4412static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004413zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004414{
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004415 _Py_IDENTIFIER(fillvalue);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004416 ziplongestobject *lz;
4417 Py_ssize_t i;
4418 PyObject *ittuple; /* tuple of iterators */
4419 PyObject *result;
4420 PyObject *fillvalue = Py_None;
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004421 Py_ssize_t tuplesize;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004422
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004423 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004424 fillvalue = NULL;
4425 if (PyDict_GET_SIZE(kwds) == 1) {
4426 fillvalue = _PyDict_GetItemIdWithError(kwds, &PyId_fillvalue);
4427 }
4428 if (fillvalue == NULL) {
4429 if (!PyErr_Occurred()) {
4430 PyErr_SetString(PyExc_TypeError,
4431 "zip_longest() got an unexpected keyword argument");
4432 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004433 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004434 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004435 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004436
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004437 /* args must be a tuple */
4438 assert(PyTuple_Check(args));
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004439 tuplesize = PyTuple_GET_SIZE(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004440
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004441 /* obtain iterators */
4442 ittuple = PyTuple_New(tuplesize);
4443 if (ittuple == NULL)
4444 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004445 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004446 PyObject *item = PyTuple_GET_ITEM(args, i);
4447 PyObject *it = PyObject_GetIter(item);
4448 if (it == NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004449 Py_DECREF(ittuple);
4450 return NULL;
4451 }
4452 PyTuple_SET_ITEM(ittuple, i, it);
4453 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004454
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004455 /* create a result holder */
4456 result = PyTuple_New(tuplesize);
4457 if (result == NULL) {
4458 Py_DECREF(ittuple);
4459 return NULL;
4460 }
4461 for (i=0 ; i < tuplesize ; i++) {
4462 Py_INCREF(Py_None);
4463 PyTuple_SET_ITEM(result, i, Py_None);
4464 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004465
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004466 /* create ziplongestobject structure */
4467 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4468 if (lz == NULL) {
4469 Py_DECREF(ittuple);
4470 Py_DECREF(result);
4471 return NULL;
4472 }
4473 lz->ittuple = ittuple;
4474 lz->tuplesize = tuplesize;
4475 lz->numactive = tuplesize;
4476 lz->result = result;
4477 Py_INCREF(fillvalue);
4478 lz->fillvalue = fillvalue;
4479 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004480}
4481
4482static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004483zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004484{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004485 PyObject_GC_UnTrack(lz);
4486 Py_XDECREF(lz->ittuple);
4487 Py_XDECREF(lz->result);
4488 Py_XDECREF(lz->fillvalue);
4489 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004490}
4491
4492static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004493zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004494{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004495 Py_VISIT(lz->ittuple);
4496 Py_VISIT(lz->result);
4497 Py_VISIT(lz->fillvalue);
4498 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004499}
4500
4501static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004502zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004503{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004504 Py_ssize_t i;
4505 Py_ssize_t tuplesize = lz->tuplesize;
4506 PyObject *result = lz->result;
4507 PyObject *it;
4508 PyObject *item;
4509 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004510
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004511 if (tuplesize == 0)
4512 return NULL;
4513 if (lz->numactive == 0)
4514 return NULL;
4515 if (Py_REFCNT(result) == 1) {
4516 Py_INCREF(result);
4517 for (i=0 ; i < tuplesize ; i++) {
4518 it = PyTuple_GET_ITEM(lz->ittuple, i);
4519 if (it == NULL) {
4520 Py_INCREF(lz->fillvalue);
4521 item = lz->fillvalue;
4522 } else {
4523 item = PyIter_Next(it);
4524 if (item == NULL) {
4525 lz->numactive -= 1;
4526 if (lz->numactive == 0 || PyErr_Occurred()) {
4527 lz->numactive = 0;
4528 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004529 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004530 } else {
4531 Py_INCREF(lz->fillvalue);
4532 item = lz->fillvalue;
4533 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4534 Py_DECREF(it);
4535 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004536 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004537 }
4538 olditem = PyTuple_GET_ITEM(result, i);
4539 PyTuple_SET_ITEM(result, i, item);
4540 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004541 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004542 } else {
4543 result = PyTuple_New(tuplesize);
4544 if (result == NULL)
4545 return NULL;
4546 for (i=0 ; i < tuplesize ; i++) {
4547 it = PyTuple_GET_ITEM(lz->ittuple, i);
4548 if (it == NULL) {
4549 Py_INCREF(lz->fillvalue);
4550 item = lz->fillvalue;
4551 } else {
4552 item = PyIter_Next(it);
4553 if (item == NULL) {
4554 lz->numactive -= 1;
4555 if (lz->numactive == 0 || PyErr_Occurred()) {
4556 lz->numactive = 0;
4557 Py_DECREF(result);
4558 return NULL;
4559 } else {
4560 Py_INCREF(lz->fillvalue);
4561 item = lz->fillvalue;
4562 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4563 Py_DECREF(it);
4564 }
4565 }
4566 }
4567 PyTuple_SET_ITEM(result, i, item);
4568 }
4569 }
4570 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004571}
4572
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004573static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304574zip_longest_reduce(ziplongestobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004575{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004576
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004577 /* Create a new tuple with empty sequences where appropriate to pickle.
4578 * Then use setstate to set the fillvalue
4579 */
4580 int i;
4581 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004582
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004583 if (args == NULL)
4584 return NULL;
4585 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4586 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4587 if (elem == NULL) {
4588 elem = PyTuple_New(0);
4589 if (elem == NULL) {
4590 Py_DECREF(args);
4591 return NULL;
4592 }
4593 } else
4594 Py_INCREF(elem);
4595 PyTuple_SET_ITEM(args, i, elem);
4596 }
4597 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4598}
4599
4600static PyObject *
4601zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4602{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004603 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004604 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004605 Py_RETURN_NONE;
4606}
4607
4608static PyMethodDef zip_longest_methods[] = {
4609 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4610 reduce_doc},
4611 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4612 setstate_doc},
4613 {NULL, NULL} /* sentinel */
4614};
4615
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004616PyDoc_STRVAR(zip_longest_doc,
4617"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004618\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004619Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004620the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004621method continues until the longest iterable in the argument sequence\n\
4622is exhausted and then it raises StopIteration. When the shorter iterables\n\
4623are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4624defaults to None or can be specified by a keyword argument.\n\
4625");
4626
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004627static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004628 PyVarObject_HEAD_INIT(NULL, 0)
4629 "itertools.zip_longest", /* tp_name */
4630 sizeof(ziplongestobject), /* tp_basicsize */
4631 0, /* tp_itemsize */
4632 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004633 (destructor)zip_longest_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004634 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004635 0, /* tp_getattr */
4636 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004637 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004638 0, /* tp_repr */
4639 0, /* tp_as_number */
4640 0, /* tp_as_sequence */
4641 0, /* tp_as_mapping */
4642 0, /* tp_hash */
4643 0, /* tp_call */
4644 0, /* tp_str */
4645 PyObject_GenericGetAttr, /* tp_getattro */
4646 0, /* tp_setattro */
4647 0, /* tp_as_buffer */
4648 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4649 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004650 zip_longest_doc, /* tp_doc */
4651 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004652 0, /* tp_clear */
4653 0, /* tp_richcompare */
4654 0, /* tp_weaklistoffset */
4655 PyObject_SelfIter, /* tp_iter */
4656 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004657 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004658 0, /* tp_members */
4659 0, /* tp_getset */
4660 0, /* tp_base */
4661 0, /* tp_dict */
4662 0, /* tp_descr_get */
4663 0, /* tp_descr_set */
4664 0, /* tp_dictoffset */
4665 0, /* tp_init */
4666 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004667 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004668 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004669};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004670
Tal Einatc4bccd32018-09-12 00:49:13 +03004671
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004672/* module level code ********************************************************/
4673
4674PyDoc_STRVAR(module_doc,
4675"Functional tools for creating and using iterators.\n\
4676\n\
4677Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004678count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004679cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004680repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004681\n\
4682Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004683accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
oldkaa0735f2018-02-02 16:52:55 +08004684chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
4685chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004686compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4687dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4688groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004689filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004690islice(seq, [start,] stop [, step]) --> elements from\n\
4691 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004692starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004693tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004694takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
oldkaa0735f2018-02-02 16:52:55 +08004695zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004696\n\
4697Combinatoric generators:\n\
4698product(p, q, ... [repeat=1]) --> cartesian product\n\
4699permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004700combinations(p, r)\n\
4701combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004702");
4703
4704
Raymond Hettingerad983e72003-11-12 14:32:26 +00004705static PyMethodDef module_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03004706 ITERTOOLS_TEE_METHODDEF
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004707 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004708};
4709
Martin v. Löwis1a214512008-06-11 05:26:20 +00004710
4711static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004712 PyModuleDef_HEAD_INIT,
4713 "itertools",
4714 module_doc,
4715 -1,
4716 module_methods,
4717 NULL,
4718 NULL,
4719 NULL,
4720 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004721};
4722
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004723PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004724PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004725{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004726 int i;
4727 PyObject *m;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004728 const char *name;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004729 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004730 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004731 &combinations_type,
4732 &cwr_type,
4733 &cycle_type,
4734 &dropwhile_type,
4735 &takewhile_type,
4736 &islice_type,
4737 &starmap_type,
4738 &chain_type,
4739 &compress_type,
4740 &filterfalse_type,
4741 &count_type,
4742 &ziplongest_type,
4743 &permutations_type,
4744 &product_type,
4745 &repeat_type,
4746 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004747 &_grouper_type,
4748 &tee_type,
4749 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004750 NULL
4751 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004752
Victor Stinnerd2ec81a2020-02-07 09:17:07 +01004753 Py_SET_TYPE(&teedataobject_type, &PyType_Type);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004754 m = PyModule_Create(&itertoolsmodule);
Victor Stinnerd2ec81a2020-02-07 09:17:07 +01004755 if (m == NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004756 return NULL;
Victor Stinnerd2ec81a2020-02-07 09:17:07 +01004757 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004758
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004759 for (i=0 ; typelist[i] != NULL ; i++) {
Victor Stinnerd2ec81a2020-02-07 09:17:07 +01004760 if (PyType_Ready(typelist[i]) < 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004761 return NULL;
Victor Stinnerd2ec81a2020-02-07 09:17:07 +01004762 }
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004763 name = _PyType_Name(typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004764 Py_INCREF(typelist[i]);
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004765 PyModule_AddObject(m, name, (PyObject *)typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004766 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004767
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004768 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004769}