blob: 4ebeb741b6b5c88049a61ae4dc6d190b7627df49 [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 Hettinger6a5b0272003-10-24 08:45:23 +0000458static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000459teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000460{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000461 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000462
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000463 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
464 if (tdo == NULL)
465 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000466
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300467 tdo->running = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000468 tdo->numread = 0;
469 tdo->nextlink = NULL;
470 Py_INCREF(it);
471 tdo->it = it;
472 PyObject_GC_Track(tdo);
473 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000474}
475
476static PyObject *
477teedataobject_jumplink(teedataobject *tdo)
478{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000479 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000480 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000481 Py_XINCREF(tdo->nextlink);
482 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000483}
484
485static PyObject *
486teedataobject_getitem(teedataobject *tdo, int i)
487{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000488 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000489
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000490 assert(i < LINKCELLS);
491 if (i < tdo->numread)
492 value = tdo->values[i];
493 else {
494 /* this is the lead iterator, so fetch more data */
495 assert(i == tdo->numread);
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300496 if (tdo->running) {
497 PyErr_SetString(PyExc_RuntimeError,
498 "cannot re-enter the tee iterator");
499 return NULL;
500 }
501 tdo->running = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000502 value = PyIter_Next(tdo->it);
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300503 tdo->running = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000504 if (value == NULL)
505 return NULL;
506 tdo->numread++;
507 tdo->values[i] = value;
508 }
509 Py_INCREF(value);
510 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000511}
512
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000513static int
514teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
515{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000516 int i;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700517
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000518 Py_VISIT(tdo->it);
519 for (i = 0; i < tdo->numread; i++)
520 Py_VISIT(tdo->values[i]);
521 Py_VISIT(tdo->nextlink);
522 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000523}
524
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200525static void
526teedataobject_safe_decref(PyObject *obj)
527{
Andy Lesterdffe4c02020-03-04 07:15:20 -0600528 while (obj && Py_IS_TYPE(obj, &teedataobject_type) &&
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200529 Py_REFCNT(obj) == 1) {
530 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
531 ((teedataobject *)obj)->nextlink = NULL;
532 Py_DECREF(obj);
533 obj = nextlink;
534 }
535 Py_XDECREF(obj);
536}
537
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000538static int
539teedataobject_clear(teedataobject *tdo)
540{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000541 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200542 PyObject *tmp;
543
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000544 Py_CLEAR(tdo->it);
545 for (i=0 ; i<tdo->numread ; i++)
546 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200547 tmp = tdo->nextlink;
548 tdo->nextlink = NULL;
549 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000550 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000551}
552
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000553static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000554teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000555{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000556 PyObject_GC_UnTrack(tdo);
557 teedataobject_clear(tdo);
558 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000559}
560
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000561static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530562teedataobject_reduce(teedataobject *tdo, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000563{
564 int i;
565 /* create a temporary list of already iterated values */
566 PyObject *values = PyList_New(tdo->numread);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700567
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000568 if (!values)
569 return NULL;
570 for (i=0 ; i<tdo->numread ; i++) {
571 Py_INCREF(tdo->values[i]);
572 PyList_SET_ITEM(values, i, tdo->values[i]);
573 }
574 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
575 values,
576 tdo->nextlink ? tdo->nextlink : Py_None);
577}
578
Tal Einatc4bccd32018-09-12 00:49:13 +0300579/*[clinic input]
580@classmethod
581itertools.teedataobject.__new__
582 iterable as it: object
583 values: object(subclass_of='&PyList_Type')
584 next: object
585 /
586Data container common to multiple tee objects.
587[clinic start generated code]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000588
589static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300590itertools_teedataobject_impl(PyTypeObject *type, PyObject *it,
591 PyObject *values, PyObject *next)
592/*[clinic end generated code: output=3343ceb07e08df5e input=be60f2fabd2b72ba]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000593{
594 teedataobject *tdo;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000595 Py_ssize_t i, len;
596
597 assert(type == &teedataobject_type);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000598
599 tdo = (teedataobject *)teedataobject_newinternal(it);
600 if (!tdo)
601 return NULL;
602
603 len = PyList_GET_SIZE(values);
604 if (len > LINKCELLS)
605 goto err;
606 for (i=0; i<len; i++) {
607 tdo->values[i] = PyList_GET_ITEM(values, i);
608 Py_INCREF(tdo->values[i]);
609 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200610 /* len <= LINKCELLS < INT_MAX */
611 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000612
613 if (len == LINKCELLS) {
614 if (next != Py_None) {
Andy Lester55728702020-03-06 16:53:17 -0600615 if (!Py_IS_TYPE(next, &teedataobject_type))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000616 goto err;
617 assert(tdo->nextlink == NULL);
618 Py_INCREF(next);
619 tdo->nextlink = next;
620 }
621 } else {
622 if (next != Py_None)
623 goto err; /* shouldn't have a next if we are not full */
624 }
625 return (PyObject*)tdo;
626
627err:
628 Py_XDECREF(tdo);
629 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
630 return NULL;
631}
632
633static PyMethodDef teedataobject_methods[] = {
634 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
635 reduce_doc},
636 {NULL, NULL} /* sentinel */
637};
638
Raymond Hettingerad983e72003-11-12 14:32:26 +0000639static PyTypeObject teedataobject_type = {
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700640 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000641 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000642 sizeof(teedataobject), /* tp_basicsize */
643 0, /* tp_itemsize */
644 /* methods */
645 (destructor)teedataobject_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200646 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000647 0, /* tp_getattr */
648 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200649 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000650 0, /* tp_repr */
651 0, /* tp_as_number */
652 0, /* tp_as_sequence */
653 0, /* tp_as_mapping */
654 0, /* tp_hash */
655 0, /* tp_call */
656 0, /* tp_str */
657 PyObject_GenericGetAttr, /* tp_getattro */
658 0, /* tp_setattro */
659 0, /* tp_as_buffer */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700660 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300661 itertools_teedataobject__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000662 (traverseproc)teedataobject_traverse, /* tp_traverse */
663 (inquiry)teedataobject_clear, /* tp_clear */
664 0, /* tp_richcompare */
665 0, /* tp_weaklistoffset */
666 0, /* tp_iter */
667 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000668 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000669 0, /* tp_members */
670 0, /* tp_getset */
671 0, /* tp_base */
672 0, /* tp_dict */
673 0, /* tp_descr_get */
674 0, /* tp_descr_set */
675 0, /* tp_dictoffset */
676 0, /* tp_init */
677 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300678 itertools_teedataobject, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000679 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000680};
681
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000682
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000683static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000684tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000685{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000686 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000687
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000688 if (to->index >= LINKCELLS) {
689 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200690 if (link == NULL)
691 return NULL;
Serhiy Storchaka57a01d32016-04-10 18:05:40 +0300692 Py_SETREF(to->dataobj, (teedataobject *)link);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000693 to->index = 0;
694 }
695 value = teedataobject_getitem(to->dataobj, to->index);
696 if (value == NULL)
697 return NULL;
698 to->index++;
699 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000700}
701
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000702static int
703tee_traverse(teeobject *to, visitproc visit, void *arg)
704{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000705 Py_VISIT((PyObject *)to->dataobj);
706 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000707}
708
Raymond Hettingerad983e72003-11-12 14:32:26 +0000709static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530710tee_copy(teeobject *to, PyObject *Py_UNUSED(ignored))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000711{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000712 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000713
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000714 newto = PyObject_GC_New(teeobject, &tee_type);
715 if (newto == NULL)
716 return NULL;
717 Py_INCREF(to->dataobj);
718 newto->dataobj = to->dataobj;
719 newto->index = to->index;
720 newto->weakreflist = NULL;
721 PyObject_GC_Track(newto);
722 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000723}
724
725PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
726
727static PyObject *
728tee_fromiterable(PyObject *iterable)
729{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000730 teeobject *to;
731 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000732
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000733 it = PyObject_GetIter(iterable);
734 if (it == NULL)
735 return NULL;
736 if (PyObject_TypeCheck(it, &tee_type)) {
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530737 to = (teeobject *)tee_copy((teeobject *)it, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000738 goto done;
739 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000740
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000741 to = PyObject_GC_New(teeobject, &tee_type);
742 if (to == NULL)
743 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000744 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000745 if (!to->dataobj) {
746 PyObject_GC_Del(to);
747 to = NULL;
748 goto done;
749 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000750
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000751 to->index = 0;
752 to->weakreflist = NULL;
753 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000754done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000755 Py_XDECREF(it);
756 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000757}
758
Tal Einatc4bccd32018-09-12 00:49:13 +0300759/*[clinic input]
760@classmethod
761itertools._tee.__new__
762 iterable: object
763 /
764Iterator wrapped to make it copyable.
765[clinic start generated code]*/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000766
Tal Einatc4bccd32018-09-12 00:49:13 +0300767static PyObject *
768itertools__tee_impl(PyTypeObject *type, PyObject *iterable)
769/*[clinic end generated code: output=b02d3fd26c810c3f input=adc0779d2afe37a2]*/
770{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000771 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000772}
773
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000774static int
775tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000776{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000777 if (to->weakreflist != NULL)
778 PyObject_ClearWeakRefs((PyObject *) to);
779 Py_CLEAR(to->dataobj);
780 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000781}
782
783static void
784tee_dealloc(teeobject *to)
785{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000786 PyObject_GC_UnTrack(to);
787 tee_clear(to);
788 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000789}
790
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000791static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530792tee_reduce(teeobject *to, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000793{
794 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
795}
796
797static PyObject *
798tee_setstate(teeobject *to, PyObject *state)
799{
800 teedataobject *tdo;
801 int index;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300802 if (!PyTuple_Check(state)) {
803 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000804 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300805 }
806 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
807 return NULL;
808 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000809 if (index < 0 || index > LINKCELLS) {
810 PyErr_SetString(PyExc_ValueError, "Index out of range");
811 return NULL;
812 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200813 Py_INCREF(tdo);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300814 Py_XSETREF(to->dataobj, tdo);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000815 to->index = index;
816 Py_RETURN_NONE;
817}
818
Raymond Hettingerad983e72003-11-12 14:32:26 +0000819static PyMethodDef tee_methods[] = {
Raymond Hettingera6a2d442015-08-16 14:38:07 -0700820 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
821 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
822 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000823 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000824};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000825
826static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000827 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000828 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000829 sizeof(teeobject), /* tp_basicsize */
830 0, /* tp_itemsize */
831 /* methods */
832 (destructor)tee_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200833 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000834 0, /* tp_getattr */
835 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200836 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000837 0, /* tp_repr */
838 0, /* tp_as_number */
839 0, /* tp_as_sequence */
840 0, /* tp_as_mapping */
841 0, /* tp_hash */
842 0, /* tp_call */
843 0, /* tp_str */
844 0, /* tp_getattro */
845 0, /* tp_setattro */
846 0, /* tp_as_buffer */
847 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300848 itertools__tee__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000849 (traverseproc)tee_traverse, /* tp_traverse */
850 (inquiry)tee_clear, /* tp_clear */
851 0, /* tp_richcompare */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700852 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000853 PyObject_SelfIter, /* tp_iter */
854 (iternextfunc)tee_next, /* tp_iternext */
855 tee_methods, /* tp_methods */
856 0, /* tp_members */
857 0, /* tp_getset */
858 0, /* tp_base */
859 0, /* tp_dict */
860 0, /* tp_descr_get */
861 0, /* tp_descr_set */
862 0, /* tp_dictoffset */
863 0, /* tp_init */
864 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300865 itertools__tee, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000866 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000867};
868
Tal Einatc4bccd32018-09-12 00:49:13 +0300869/*[clinic input]
870itertools.tee
871 iterable: object
872 n: Py_ssize_t = 2
873 /
874Returns a tuple of n independent iterators.
875[clinic start generated code]*/
876
Raymond Hettingerad983e72003-11-12 14:32:26 +0000877static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300878itertools_tee_impl(PyObject *module, PyObject *iterable, Py_ssize_t n)
879/*[clinic end generated code: output=1c64519cd859c2f0 input=c99a1472c425d66d]*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000880{
Tal Einatc4bccd32018-09-12 00:49:13 +0300881 Py_ssize_t i;
882 PyObject *it, *copyable, *copyfunc, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200883 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000884
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000885 if (n < 0) {
886 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
887 return NULL;
888 }
889 result = PyTuple_New(n);
890 if (result == NULL)
891 return NULL;
892 if (n == 0)
893 return result;
894 it = PyObject_GetIter(iterable);
895 if (it == NULL) {
896 Py_DECREF(result);
897 return NULL;
898 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200899
Serhiy Storchakaf320be72018-01-25 10:49:40 +0200900 if (_PyObject_LookupAttrId(it, &PyId___copy__, &copyfunc) < 0) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200901 Py_DECREF(it);
902 Py_DECREF(result);
903 return NULL;
904 }
Serhiy Storchakaf320be72018-01-25 10:49:40 +0200905 if (copyfunc != NULL) {
906 copyable = it;
907 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200908 else {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000909 copyable = tee_fromiterable(it);
910 Py_DECREF(it);
911 if (copyable == NULL) {
912 Py_DECREF(result);
913 return NULL;
914 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200915 copyfunc = _PyObject_GetAttrId(copyable, &PyId___copy__);
916 if (copyfunc == NULL) {
917 Py_DECREF(copyable);
918 Py_DECREF(result);
919 return NULL;
920 }
921 }
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200922
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200923 PyTuple_SET_ITEM(result, 0, copyable);
924 for (i = 1; i < n; i++) {
925 copyable = _PyObject_CallNoArg(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000926 if (copyable == NULL) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200927 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000928 Py_DECREF(result);
929 return NULL;
930 }
931 PyTuple_SET_ITEM(result, i, copyable);
932 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200933 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000934 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000935}
936
Raymond Hettingerad983e72003-11-12 14:32:26 +0000937
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700938/* cycle object **************************************************************/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000939
940typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000941 PyObject_HEAD
942 PyObject *it;
943 PyObject *saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700944 Py_ssize_t index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000945 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000946} cycleobject;
947
Tal Einatc4bccd32018-09-12 00:49:13 +0300948/*[clinic input]
949@classmethod
950itertools.cycle.__new__
951 iterable: object
952 /
953Return elements from the iterable until it is exhausted. Then repeat the sequence indefinitely.
954[clinic start generated code]*/
955
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000956static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300957itertools_cycle_impl(PyTypeObject *type, PyObject *iterable)
958/*[clinic end generated code: output=f60e5ec17a45b35c input=9d1d84bcf66e908b]*/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000959{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000960 PyObject *it;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000961 PyObject *saved;
962 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000963
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000964 /* Get iterator. */
965 it = PyObject_GetIter(iterable);
966 if (it == NULL)
967 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000968
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000969 saved = PyList_New(0);
970 if (saved == NULL) {
971 Py_DECREF(it);
972 return NULL;
973 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000974
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000975 /* create cycleobject structure */
976 lz = (cycleobject *)type->tp_alloc(type, 0);
977 if (lz == NULL) {
978 Py_DECREF(it);
979 Py_DECREF(saved);
980 return NULL;
981 }
982 lz->it = it;
983 lz->saved = saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700984 lz->index = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000985 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000986
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000987 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000988}
989
990static void
991cycle_dealloc(cycleobject *lz)
992{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000993 PyObject_GC_UnTrack(lz);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000994 Py_XDECREF(lz->it);
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700995 Py_XDECREF(lz->saved);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000996 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000997}
998
999static int
1000cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
1001{
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001002 if (lz->it)
1003 Py_VISIT(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001004 Py_VISIT(lz->saved);
1005 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001006}
1007
1008static PyObject *
1009cycle_next(cycleobject *lz)
1010{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001011 PyObject *item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001012
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001013 if (lz->it != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001014 item = PyIter_Next(lz->it);
1015 if (item != NULL) {
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001016 if (lz->firstpass)
1017 return item;
1018 if (PyList_Append(lz->saved, item)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001019 Py_DECREF(item);
1020 return NULL;
1021 }
1022 return item;
1023 }
Raymond Hettinger98958fe2015-08-15 15:09:30 -07001024 /* Note: StopIteration is already cleared by PyIter_Next() */
1025 if (PyErr_Occurred())
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001026 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001027 Py_CLEAR(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001028 }
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001029 if (PyList_GET_SIZE(lz->saved) == 0)
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001030 return NULL;
1031 item = PyList_GET_ITEM(lz->saved, lz->index);
1032 lz->index++;
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001033 if (lz->index >= PyList_GET_SIZE(lz->saved))
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001034 lz->index = 0;
1035 Py_INCREF(item);
1036 return item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001037}
1038
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001039static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301040cycle_reduce(cycleobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001041{
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001042 /* Create a new cycle with the iterator tuple, then set the saved state */
1043 if (lz->it == NULL) {
1044 PyObject *it = PyObject_GetIter(lz->saved);
1045 if (it == NULL)
1046 return NULL;
1047 if (lz->index != 0) {
1048 _Py_IDENTIFIER(__setstate__);
1049 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
1050 "n", lz->index);
1051 if (res == NULL) {
1052 Py_DECREF(it);
1053 return NULL;
1054 }
1055 Py_DECREF(res);
1056 }
Serhiy Storchaka1f21eaa2019-09-01 12:16:51 +03001057 return Py_BuildValue("O(N)(OO)", Py_TYPE(lz), it, lz->saved, Py_True);
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001058 }
Serhiy Storchaka1f21eaa2019-09-01 12:16:51 +03001059 return Py_BuildValue("O(O)(OO)", Py_TYPE(lz), lz->it, lz->saved,
1060 lz->firstpass ? Py_True : Py_False);
Raymond Hettingerc39786d2015-08-15 15:16:12 -07001061}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001062
1063static PyObject *
1064cycle_setstate(cycleobject *lz, PyObject *state)
1065{
1066 PyObject *saved=NULL;
1067 int firstpass;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001068 if (!PyTuple_Check(state)) {
1069 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001070 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001071 }
1072 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
1073 return NULL;
1074 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001075 Py_INCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001076 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001077 lz->firstpass = firstpass != 0;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001078 lz->index = 0;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001079 Py_RETURN_NONE;
1080}
1081
1082static PyMethodDef cycle_methods[] = {
1083 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1084 reduce_doc},
1085 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1086 setstate_doc},
1087 {NULL, NULL} /* sentinel */
1088};
1089
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001090static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001091 PyVarObject_HEAD_INIT(NULL, 0)
1092 "itertools.cycle", /* tp_name */
1093 sizeof(cycleobject), /* tp_basicsize */
1094 0, /* tp_itemsize */
1095 /* methods */
1096 (destructor)cycle_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001097 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001098 0, /* tp_getattr */
1099 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001100 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001101 0, /* tp_repr */
1102 0, /* tp_as_number */
1103 0, /* tp_as_sequence */
1104 0, /* tp_as_mapping */
1105 0, /* tp_hash */
1106 0, /* tp_call */
1107 0, /* tp_str */
1108 PyObject_GenericGetAttr, /* tp_getattro */
1109 0, /* tp_setattro */
1110 0, /* tp_as_buffer */
1111 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1112 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001113 itertools_cycle__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001114 (traverseproc)cycle_traverse, /* tp_traverse */
1115 0, /* tp_clear */
1116 0, /* tp_richcompare */
1117 0, /* tp_weaklistoffset */
1118 PyObject_SelfIter, /* tp_iter */
1119 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001120 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001121 0, /* tp_members */
1122 0, /* tp_getset */
1123 0, /* tp_base */
1124 0, /* tp_dict */
1125 0, /* tp_descr_get */
1126 0, /* tp_descr_set */
1127 0, /* tp_dictoffset */
1128 0, /* tp_init */
1129 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001130 itertools_cycle, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001131 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001132};
1133
1134
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001135/* dropwhile object **********************************************************/
1136
1137typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001138 PyObject_HEAD
1139 PyObject *func;
1140 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001141 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001142} dropwhileobject;
1143
Tal Einatc4bccd32018-09-12 00:49:13 +03001144/*[clinic input]
1145@classmethod
1146itertools.dropwhile.__new__
1147 predicate as func: object
1148 iterable as seq: object
1149 /
1150Drop items from the iterable while predicate(item) is true.
1151
1152Afterwards, return every element until the iterable is exhausted.
1153[clinic start generated code]*/
1154
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001155static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001156itertools_dropwhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1157/*[clinic end generated code: output=92f9d0d89af149e4 input=d39737147c9f0a26]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001158{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001159 PyObject *it;
1160 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001161
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001162 /* Get iterator. */
1163 it = PyObject_GetIter(seq);
1164 if (it == NULL)
1165 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001166
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001167 /* create dropwhileobject structure */
1168 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1169 if (lz == NULL) {
1170 Py_DECREF(it);
1171 return NULL;
1172 }
1173 Py_INCREF(func);
1174 lz->func = func;
1175 lz->it = it;
1176 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001177
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001178 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001179}
1180
1181static void
1182dropwhile_dealloc(dropwhileobject *lz)
1183{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001184 PyObject_GC_UnTrack(lz);
1185 Py_XDECREF(lz->func);
1186 Py_XDECREF(lz->it);
1187 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001188}
1189
1190static int
1191dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1192{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001193 Py_VISIT(lz->it);
1194 Py_VISIT(lz->func);
1195 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196}
1197
1198static PyObject *
1199dropwhile_next(dropwhileobject *lz)
1200{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001201 PyObject *item, *good;
1202 PyObject *it = lz->it;
1203 long ok;
1204 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001205
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001206 iternext = *Py_TYPE(it)->tp_iternext;
1207 for (;;) {
1208 item = iternext(it);
1209 if (item == NULL)
1210 return NULL;
1211 if (lz->start == 1)
1212 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001213
Petr Viktorinffd97532020-02-11 17:46:57 +01001214 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001215 if (good == NULL) {
1216 Py_DECREF(item);
1217 return NULL;
1218 }
1219 ok = PyObject_IsTrue(good);
1220 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001221 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001222 lz->start = 1;
1223 return item;
1224 }
1225 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001226 if (ok < 0)
1227 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001228 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001229}
1230
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001231static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301232dropwhile_reduce(dropwhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001233{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001234 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 +00001235}
1236
1237static PyObject *
1238dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1239{
1240 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001241 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001242 return NULL;
1243 lz->start = start;
1244 Py_RETURN_NONE;
1245}
1246
1247static PyMethodDef dropwhile_methods[] = {
1248 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1249 reduce_doc},
1250 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1251 setstate_doc},
1252 {NULL, NULL} /* sentinel */
1253};
1254
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001255static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001256 PyVarObject_HEAD_INIT(NULL, 0)
1257 "itertools.dropwhile", /* tp_name */
1258 sizeof(dropwhileobject), /* tp_basicsize */
1259 0, /* tp_itemsize */
1260 /* methods */
1261 (destructor)dropwhile_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001262 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001263 0, /* tp_getattr */
1264 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001265 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001266 0, /* tp_repr */
1267 0, /* tp_as_number */
1268 0, /* tp_as_sequence */
1269 0, /* tp_as_mapping */
1270 0, /* tp_hash */
1271 0, /* tp_call */
1272 0, /* tp_str */
1273 PyObject_GenericGetAttr, /* tp_getattro */
1274 0, /* tp_setattro */
1275 0, /* tp_as_buffer */
1276 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1277 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001278 itertools_dropwhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001279 (traverseproc)dropwhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001280 0, /* tp_clear */
1281 0, /* tp_richcompare */
1282 0, /* tp_weaklistoffset */
1283 PyObject_SelfIter, /* tp_iter */
1284 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001285 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001286 0, /* tp_members */
1287 0, /* tp_getset */
1288 0, /* tp_base */
1289 0, /* tp_dict */
1290 0, /* tp_descr_get */
1291 0, /* tp_descr_set */
1292 0, /* tp_dictoffset */
1293 0, /* tp_init */
1294 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001295 itertools_dropwhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001296 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001297};
1298
1299
1300/* takewhile object **********************************************************/
1301
1302typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001303 PyObject_HEAD
1304 PyObject *func;
1305 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001306 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001307} takewhileobject;
1308
Tal Einatc4bccd32018-09-12 00:49:13 +03001309/*[clinic input]
1310@classmethod
1311itertools.takewhile.__new__
1312 predicate as func: object
1313 iterable as seq: object
1314 /
1315Return successive entries from an iterable as long as the predicate evaluates to true for each entry.
1316[clinic start generated code]*/
1317
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001318static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001319itertools_takewhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1320/*[clinic end generated code: output=bb179ea7864e2ef6 input=ba5255f7519aa119]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001321{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001322 PyObject *it;
1323 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001324
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001325 /* Get iterator. */
1326 it = PyObject_GetIter(seq);
1327 if (it == NULL)
1328 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001329
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001330 /* create takewhileobject structure */
1331 lz = (takewhileobject *)type->tp_alloc(type, 0);
1332 if (lz == NULL) {
1333 Py_DECREF(it);
1334 return NULL;
1335 }
1336 Py_INCREF(func);
1337 lz->func = func;
1338 lz->it = it;
1339 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001340
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001341 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001342}
1343
1344static void
1345takewhile_dealloc(takewhileobject *lz)
1346{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001347 PyObject_GC_UnTrack(lz);
1348 Py_XDECREF(lz->func);
1349 Py_XDECREF(lz->it);
1350 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001351}
1352
1353static int
1354takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1355{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001356 Py_VISIT(lz->it);
1357 Py_VISIT(lz->func);
1358 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001359}
1360
1361static PyObject *
1362takewhile_next(takewhileobject *lz)
1363{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001364 PyObject *item, *good;
1365 PyObject *it = lz->it;
1366 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001367
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001368 if (lz->stop == 1)
1369 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001370
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001371 item = (*Py_TYPE(it)->tp_iternext)(it);
1372 if (item == NULL)
1373 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001374
Petr Viktorinffd97532020-02-11 17:46:57 +01001375 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001376 if (good == NULL) {
1377 Py_DECREF(item);
1378 return NULL;
1379 }
1380 ok = PyObject_IsTrue(good);
1381 Py_DECREF(good);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001382 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001383 return item;
1384 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001385 if (ok == 0)
1386 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001387 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001388}
1389
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001390static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301391takewhile_reduce(takewhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001392{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001393 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 +00001394}
1395
1396static PyObject *
1397takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1398{
1399 int stop = PyObject_IsTrue(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001400
Antoine Pitrou721738f2012-08-15 23:20:39 +02001401 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001402 return NULL;
1403 lz->stop = stop;
1404 Py_RETURN_NONE;
1405}
1406
1407static PyMethodDef takewhile_reduce_methods[] = {
1408 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1409 reduce_doc},
1410 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1411 setstate_doc},
1412 {NULL, NULL} /* sentinel */
1413};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001414
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001415static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001416 PyVarObject_HEAD_INIT(NULL, 0)
1417 "itertools.takewhile", /* tp_name */
1418 sizeof(takewhileobject), /* tp_basicsize */
1419 0, /* tp_itemsize */
1420 /* methods */
1421 (destructor)takewhile_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001422 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001423 0, /* tp_getattr */
1424 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001425 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001426 0, /* tp_repr */
1427 0, /* tp_as_number */
1428 0, /* tp_as_sequence */
1429 0, /* tp_as_mapping */
1430 0, /* tp_hash */
1431 0, /* tp_call */
1432 0, /* tp_str */
1433 PyObject_GenericGetAttr, /* tp_getattro */
1434 0, /* tp_setattro */
1435 0, /* tp_as_buffer */
1436 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1437 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001438 itertools_takewhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001439 (traverseproc)takewhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001440 0, /* tp_clear */
1441 0, /* tp_richcompare */
1442 0, /* tp_weaklistoffset */
1443 PyObject_SelfIter, /* tp_iter */
1444 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001445 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001446 0, /* tp_members */
1447 0, /* tp_getset */
1448 0, /* tp_base */
1449 0, /* tp_dict */
1450 0, /* tp_descr_get */
1451 0, /* tp_descr_set */
1452 0, /* tp_dictoffset */
1453 0, /* tp_init */
1454 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001455 itertools_takewhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001456 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457};
1458
1459
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001460/* islice object *************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001461
1462typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001463 PyObject_HEAD
1464 PyObject *it;
1465 Py_ssize_t next;
1466 Py_ssize_t stop;
1467 Py_ssize_t step;
1468 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469} isliceobject;
1470
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001471static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001472
1473static PyObject *
1474islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1475{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001476 PyObject *seq;
1477 Py_ssize_t start=0, stop=-1, step=1;
1478 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1479 Py_ssize_t numargs;
1480 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001481
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001482 if (type == &islice_type && !_PyArg_NoKeywords("islice", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001483 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001484
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001485 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1486 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001487
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001488 numargs = PyTuple_Size(args);
1489 if (numargs == 2) {
1490 if (a1 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001491 stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001492 if (stop == -1) {
1493 if (PyErr_Occurred())
1494 PyErr_Clear();
1495 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001496 "Stop argument for islice() must be None or "
1497 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001498 return NULL;
1499 }
1500 }
1501 } else {
1502 if (a1 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001503 start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001504 if (start == -1 && PyErr_Occurred())
1505 PyErr_Clear();
1506 if (a2 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001507 stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001508 if (stop == -1) {
1509 if (PyErr_Occurred())
1510 PyErr_Clear();
1511 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001512 "Stop argument for islice() must be None or "
1513 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001514 return NULL;
1515 }
1516 }
1517 }
1518 if (start<0 || stop<-1) {
1519 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001520 "Indices for islice() must be None or "
1521 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001522 return NULL;
1523 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001524
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001525 if (a3 != NULL) {
1526 if (a3 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001527 step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001528 if (step == -1 && PyErr_Occurred())
1529 PyErr_Clear();
1530 }
1531 if (step<1) {
1532 PyErr_SetString(PyExc_ValueError,
1533 "Step for islice() must be a positive integer or None.");
1534 return NULL;
1535 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001536
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001537 /* Get iterator. */
1538 it = PyObject_GetIter(seq);
1539 if (it == NULL)
1540 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001541
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001542 /* create isliceobject structure */
1543 lz = (isliceobject *)type->tp_alloc(type, 0);
1544 if (lz == NULL) {
1545 Py_DECREF(it);
1546 return NULL;
1547 }
1548 lz->it = it;
1549 lz->next = start;
1550 lz->stop = stop;
1551 lz->step = step;
1552 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001553
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001554 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001555}
1556
1557static void
1558islice_dealloc(isliceobject *lz)
1559{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001560 PyObject_GC_UnTrack(lz);
1561 Py_XDECREF(lz->it);
1562 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001563}
1564
1565static int
1566islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1567{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001568 Py_VISIT(lz->it);
1569 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001570}
1571
1572static PyObject *
1573islice_next(isliceobject *lz)
1574{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001575 PyObject *item;
1576 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001577 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001578 Py_ssize_t oldnext;
1579 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001580
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001581 if (it == NULL)
1582 return NULL;
1583
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001584 iternext = *Py_TYPE(it)->tp_iternext;
1585 while (lz->cnt < lz->next) {
1586 item = iternext(it);
1587 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001588 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001589 Py_DECREF(item);
1590 lz->cnt++;
1591 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001592 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001593 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001594 item = iternext(it);
1595 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001596 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001597 lz->cnt++;
1598 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001599 /* The (size_t) cast below avoids the danger of undefined
1600 behaviour from signed integer overflow. */
1601 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001602 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1603 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001604 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001605
1606empty:
1607 Py_CLEAR(lz->it);
1608 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001609}
1610
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001611static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301612islice_reduce(isliceobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001613{
1614 /* When unpickled, generate a new object with the same bounds,
1615 * then 'setstate' with the next and count
1616 */
1617 PyObject *stop;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001618
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001619 if (lz->it == NULL) {
1620 PyObject *empty_list;
1621 PyObject *empty_it;
1622 empty_list = PyList_New(0);
1623 if (empty_list == NULL)
1624 return NULL;
1625 empty_it = PyObject_GetIter(empty_list);
1626 Py_DECREF(empty_list);
1627 if (empty_it == NULL)
1628 return NULL;
1629 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1630 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001631 if (lz->stop == -1) {
1632 stop = Py_None;
1633 Py_INCREF(stop);
1634 } else {
1635 stop = PyLong_FromSsize_t(lz->stop);
1636 if (stop == NULL)
1637 return NULL;
1638 }
1639 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1640 lz->it, lz->next, stop, lz->step,
1641 lz->cnt);
1642}
1643
1644static PyObject *
1645islice_setstate(isliceobject *lz, PyObject *state)
1646{
1647 Py_ssize_t cnt = PyLong_AsSsize_t(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001648
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001649 if (cnt == -1 && PyErr_Occurred())
1650 return NULL;
1651 lz->cnt = cnt;
1652 Py_RETURN_NONE;
1653}
1654
1655static PyMethodDef islice_methods[] = {
1656 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1657 reduce_doc},
1658 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1659 setstate_doc},
1660 {NULL, NULL} /* sentinel */
1661};
1662
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001663PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001664"islice(iterable, stop) --> islice object\n\
1665islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001666\n\
1667Return an iterator whose next() method returns selected values from an\n\
1668iterable. If start is specified, will skip all preceding elements;\n\
1669otherwise, start defaults to zero. Step defaults to one. If\n\
oldkaa0735f2018-02-02 16:52:55 +08001670specified as another value, step determines how many values are\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001671skipped between successive calls. Works like a slice() on a list\n\
1672but returns an iterator.");
1673
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001674static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001675 PyVarObject_HEAD_INIT(NULL, 0)
1676 "itertools.islice", /* tp_name */
1677 sizeof(isliceobject), /* tp_basicsize */
1678 0, /* tp_itemsize */
1679 /* methods */
1680 (destructor)islice_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001681 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001682 0, /* tp_getattr */
1683 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001684 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001685 0, /* tp_repr */
1686 0, /* tp_as_number */
1687 0, /* tp_as_sequence */
1688 0, /* tp_as_mapping */
1689 0, /* tp_hash */
1690 0, /* tp_call */
1691 0, /* tp_str */
1692 PyObject_GenericGetAttr, /* tp_getattro */
1693 0, /* tp_setattro */
1694 0, /* tp_as_buffer */
1695 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1696 Py_TPFLAGS_BASETYPE, /* tp_flags */
1697 islice_doc, /* tp_doc */
1698 (traverseproc)islice_traverse, /* tp_traverse */
1699 0, /* tp_clear */
1700 0, /* tp_richcompare */
1701 0, /* tp_weaklistoffset */
1702 PyObject_SelfIter, /* tp_iter */
1703 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001704 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001705 0, /* tp_members */
1706 0, /* tp_getset */
1707 0, /* tp_base */
1708 0, /* tp_dict */
1709 0, /* tp_descr_get */
1710 0, /* tp_descr_set */
1711 0, /* tp_dictoffset */
1712 0, /* tp_init */
1713 0, /* tp_alloc */
1714 islice_new, /* tp_new */
1715 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001716};
1717
1718
1719/* starmap object ************************************************************/
1720
1721typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001722 PyObject_HEAD
1723 PyObject *func;
1724 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001725} starmapobject;
1726
Tal Einatc4bccd32018-09-12 00:49:13 +03001727/*[clinic input]
1728@classmethod
1729itertools.starmap.__new__
1730 function as func: object
1731 iterable as seq: object
1732 /
1733Return an iterator whose values are returned from the function evaluated with an argument tuple taken from the given sequence.
1734[clinic start generated code]*/
1735
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001736static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001737itertools_starmap_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1738/*[clinic end generated code: output=79eeb81d452c6e8d input=844766df6a0d4dad]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001739{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001740 PyObject *it;
1741 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001742
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001743 /* Get iterator. */
1744 it = PyObject_GetIter(seq);
1745 if (it == NULL)
1746 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001747
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001748 /* create starmapobject structure */
1749 lz = (starmapobject *)type->tp_alloc(type, 0);
1750 if (lz == NULL) {
1751 Py_DECREF(it);
1752 return NULL;
1753 }
1754 Py_INCREF(func);
1755 lz->func = func;
1756 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001757
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001758 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001759}
1760
1761static void
1762starmap_dealloc(starmapobject *lz)
1763{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001764 PyObject_GC_UnTrack(lz);
1765 Py_XDECREF(lz->func);
1766 Py_XDECREF(lz->it);
1767 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001768}
1769
1770static int
1771starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1772{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001773 Py_VISIT(lz->it);
1774 Py_VISIT(lz->func);
1775 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001776}
1777
1778static PyObject *
1779starmap_next(starmapobject *lz)
1780{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001781 PyObject *args;
1782 PyObject *result;
1783 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001784
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001785 args = (*Py_TYPE(it)->tp_iternext)(it);
1786 if (args == NULL)
1787 return NULL;
1788 if (!PyTuple_CheckExact(args)) {
1789 PyObject *newargs = PySequence_Tuple(args);
1790 Py_DECREF(args);
1791 if (newargs == NULL)
1792 return NULL;
1793 args = newargs;
1794 }
1795 result = PyObject_Call(lz->func, args, NULL);
1796 Py_DECREF(args);
1797 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001798}
1799
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001800static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301801starmap_reduce(starmapobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001802{
1803 /* Just pickle the iterator */
1804 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1805}
1806
1807static PyMethodDef starmap_methods[] = {
1808 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1809 reduce_doc},
1810 {NULL, NULL} /* sentinel */
1811};
1812
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001813static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001814 PyVarObject_HEAD_INIT(NULL, 0)
1815 "itertools.starmap", /* tp_name */
1816 sizeof(starmapobject), /* tp_basicsize */
1817 0, /* tp_itemsize */
1818 /* methods */
1819 (destructor)starmap_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001820 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001821 0, /* tp_getattr */
1822 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001823 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001824 0, /* tp_repr */
1825 0, /* tp_as_number */
1826 0, /* tp_as_sequence */
1827 0, /* tp_as_mapping */
1828 0, /* tp_hash */
1829 0, /* tp_call */
1830 0, /* tp_str */
1831 PyObject_GenericGetAttr, /* tp_getattro */
1832 0, /* tp_setattro */
1833 0, /* tp_as_buffer */
1834 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1835 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001836 itertools_starmap__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001837 (traverseproc)starmap_traverse, /* tp_traverse */
1838 0, /* tp_clear */
1839 0, /* tp_richcompare */
1840 0, /* tp_weaklistoffset */
1841 PyObject_SelfIter, /* tp_iter */
1842 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001843 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001844 0, /* tp_members */
1845 0, /* tp_getset */
1846 0, /* tp_base */
1847 0, /* tp_dict */
1848 0, /* tp_descr_get */
1849 0, /* tp_descr_set */
1850 0, /* tp_dictoffset */
1851 0, /* tp_init */
1852 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001853 itertools_starmap, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001854 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001855};
1856
1857
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001858/* chain object **************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001859
1860typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001861 PyObject_HEAD
1862 PyObject *source; /* Iterator over input iterables */
1863 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001864} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001865
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001866static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001867
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001868static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001869chain_new_internal(PyTypeObject *type, PyObject *source)
1870{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001871 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001872
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001873 lz = (chainobject *)type->tp_alloc(type, 0);
1874 if (lz == NULL) {
1875 Py_DECREF(source);
1876 return NULL;
1877 }
1878
1879 lz->source = source;
1880 lz->active = NULL;
1881 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001882}
1883
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001884static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001885chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001886{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001887 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001888
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001889 if (type == &chain_type && !_PyArg_NoKeywords("chain", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001890 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001891
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001892 source = PyObject_GetIter(args);
1893 if (source == NULL)
1894 return NULL;
1895
1896 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001897}
1898
Tal Einatc4bccd32018-09-12 00:49:13 +03001899/*[clinic input]
1900@classmethod
1901itertools.chain.from_iterable
1902 iterable as arg: object
1903 /
1904Alternative chain() constructor taking a single iterable argument that evaluates lazily.
1905[clinic start generated code]*/
1906
Christian Heimesf16baeb2008-02-29 14:57:44 +00001907static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001908itertools_chain_from_iterable(PyTypeObject *type, PyObject *arg)
1909/*[clinic end generated code: output=667ae7a7f7b68654 input=72c39e3a2ca3be85]*/
Christian Heimesf16baeb2008-02-29 14:57:44 +00001910{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001911 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001912
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001913 source = PyObject_GetIter(arg);
1914 if (source == NULL)
1915 return NULL;
1916
1917 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001918}
1919
1920static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001921chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001922{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001923 PyObject_GC_UnTrack(lz);
1924 Py_XDECREF(lz->active);
1925 Py_XDECREF(lz->source);
1926 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001927}
1928
Raymond Hettinger2012f172003-02-07 05:32:58 +00001929static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001930chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001931{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001932 Py_VISIT(lz->source);
1933 Py_VISIT(lz->active);
1934 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001935}
1936
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001937static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001938chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001939{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001940 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001941
T. Wouters5466d4a2017-03-30 09:58:35 -07001942 /* lz->source is the iterator of iterables. If it's NULL, we've already
1943 * consumed them all. lz->active is the current iterator. If it's NULL,
1944 * we should grab a new one from lz->source. */
1945 while (lz->source != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001946 if (lz->active == NULL) {
T. Wouters5466d4a2017-03-30 09:58:35 -07001947 PyObject *iterable = PyIter_Next(lz->source);
1948 if (iterable == NULL) {
1949 Py_CLEAR(lz->source);
1950 return NULL; /* no more input sources */
1951 }
1952 lz->active = PyObject_GetIter(iterable);
1953 Py_DECREF(iterable);
1954 if (lz->active == NULL) {
1955 Py_CLEAR(lz->source);
1956 return NULL; /* input not iterable */
1957 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001958 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001959 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1960 if (item != NULL)
1961 return item;
1962 if (PyErr_Occurred()) {
1963 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1964 PyErr_Clear();
1965 else
1966 return NULL; /* input raised an exception */
1967 }
1968 /* lz->active is consumed, try with the next iterable. */
1969 Py_CLEAR(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001970 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001971 /* Everything had been consumed already. */
1972 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001973}
1974
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001975static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301976chain_reduce(chainobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001977{
1978 if (lz->source) {
1979 /* we can't pickle function objects (itertools.from_iterable) so
1980 * we must use setstate to replace the iterable. One day we
1981 * will fix pickling of functions
1982 */
1983 if (lz->active) {
1984 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1985 } else {
1986 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1987 }
1988 } else {
1989 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1990 }
1991 return NULL;
1992}
1993
1994static PyObject *
1995chain_setstate(chainobject *lz, PyObject *state)
1996{
1997 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001998
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001999 if (!PyTuple_Check(state)) {
2000 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002001 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002002 }
2003 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
2004 return NULL;
2005 }
2006 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
2007 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
2008 return NULL;
2009 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002010
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002011 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002012 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002013 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002014 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002015 Py_RETURN_NONE;
2016}
2017
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002018PyDoc_STRVAR(chain_doc,
2019"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002020\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002021Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002022first iterable until it is exhausted, then elements from the next\n\
2023iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002024
Christian Heimesf16baeb2008-02-29 14:57:44 +00002025static PyMethodDef chain_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03002026 ITERTOOLS_CHAIN_FROM_ITERABLE_METHODDEF
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002027 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
2028 reduce_doc},
2029 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
2030 setstate_doc},
Ethan Smitha8403d02020-04-09 20:28:08 -07002031 {"__class_getitem__", (PyCFunction)Py_GenericAlias,
2032 METH_O|METH_CLASS, PyDoc_STR("See PEP 585")},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07002033 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00002034};
2035
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002036static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002037 PyVarObject_HEAD_INIT(NULL, 0)
2038 "itertools.chain", /* tp_name */
2039 sizeof(chainobject), /* tp_basicsize */
2040 0, /* tp_itemsize */
2041 /* methods */
2042 (destructor)chain_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002043 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002044 0, /* tp_getattr */
2045 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002046 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002047 0, /* tp_repr */
2048 0, /* tp_as_number */
2049 0, /* tp_as_sequence */
2050 0, /* tp_as_mapping */
2051 0, /* tp_hash */
2052 0, /* tp_call */
2053 0, /* tp_str */
2054 PyObject_GenericGetAttr, /* tp_getattro */
2055 0, /* tp_setattro */
2056 0, /* tp_as_buffer */
2057 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2058 Py_TPFLAGS_BASETYPE, /* tp_flags */
2059 chain_doc, /* tp_doc */
2060 (traverseproc)chain_traverse, /* tp_traverse */
2061 0, /* tp_clear */
2062 0, /* tp_richcompare */
2063 0, /* tp_weaklistoffset */
2064 PyObject_SelfIter, /* tp_iter */
2065 (iternextfunc)chain_next, /* tp_iternext */
2066 chain_methods, /* tp_methods */
2067 0, /* tp_members */
2068 0, /* tp_getset */
2069 0, /* tp_base */
2070 0, /* tp_dict */
2071 0, /* tp_descr_get */
2072 0, /* tp_descr_set */
2073 0, /* tp_dictoffset */
2074 0, /* tp_init */
2075 0, /* tp_alloc */
2076 chain_new, /* tp_new */
2077 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002078};
2079
2080
Christian Heimesc3f30c42008-02-22 16:37:40 +00002081/* product object ************************************************************/
2082
2083typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002084 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002085 PyObject *pools; /* tuple of pool tuples */
2086 Py_ssize_t *indices; /* one index per pool */
2087 PyObject *result; /* most recently returned result tuple */
2088 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002089} productobject;
2090
2091static PyTypeObject product_type;
2092
2093static PyObject *
2094product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2095{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002096 productobject *lz;
2097 Py_ssize_t nargs, npools, repeat=1;
2098 PyObject *pools = NULL;
2099 Py_ssize_t *indices = NULL;
2100 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002101
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002102 if (kwds != NULL) {
2103 char *kwlist[] = {"repeat", 0};
2104 PyObject *tmpargs = PyTuple_New(0);
2105 if (tmpargs == NULL)
2106 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002107 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2108 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002109 Py_DECREF(tmpargs);
2110 return NULL;
2111 }
2112 Py_DECREF(tmpargs);
2113 if (repeat < 0) {
2114 PyErr_SetString(PyExc_ValueError,
2115 "repeat argument cannot be negative");
2116 return NULL;
2117 }
2118 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002119
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002120 assert(PyTuple_CheckExact(args));
2121 if (repeat == 0) {
2122 nargs = 0;
2123 } else {
2124 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002125 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002126 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2127 return NULL;
2128 }
2129 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002130 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002131
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002132 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002133 if (indices == NULL) {
2134 PyErr_NoMemory();
2135 goto error;
2136 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002137
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002138 pools = PyTuple_New(npools);
2139 if (pools == NULL)
2140 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002141
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002142 for (i=0; i < nargs ; ++i) {
2143 PyObject *item = PyTuple_GET_ITEM(args, i);
2144 PyObject *pool = PySequence_Tuple(item);
2145 if (pool == NULL)
2146 goto error;
2147 PyTuple_SET_ITEM(pools, i, pool);
2148 indices[i] = 0;
2149 }
2150 for ( ; i < npools; ++i) {
2151 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2152 Py_INCREF(pool);
2153 PyTuple_SET_ITEM(pools, i, pool);
2154 indices[i] = 0;
2155 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002156
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002157 /* create productobject structure */
2158 lz = (productobject *)type->tp_alloc(type, 0);
2159 if (lz == NULL)
2160 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002161
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002162 lz->pools = pools;
2163 lz->indices = indices;
2164 lz->result = NULL;
2165 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002166
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002167 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002168
2169error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002170 if (indices != NULL)
2171 PyMem_Free(indices);
2172 Py_XDECREF(pools);
2173 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002174}
2175
2176static void
2177product_dealloc(productobject *lz)
2178{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002179 PyObject_GC_UnTrack(lz);
2180 Py_XDECREF(lz->pools);
2181 Py_XDECREF(lz->result);
2182 if (lz->indices != NULL)
2183 PyMem_Free(lz->indices);
2184 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002185}
2186
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002187static PyObject *
2188product_sizeof(productobject *lz, void *unused)
2189{
2190 Py_ssize_t res;
2191
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002192 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002193 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2194 return PyLong_FromSsize_t(res);
2195}
2196
2197PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2198
Christian Heimesc3f30c42008-02-22 16:37:40 +00002199static int
2200product_traverse(productobject *lz, visitproc visit, void *arg)
2201{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002202 Py_VISIT(lz->pools);
2203 Py_VISIT(lz->result);
2204 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002205}
2206
2207static PyObject *
2208product_next(productobject *lz)
2209{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002210 PyObject *pool;
2211 PyObject *elem;
2212 PyObject *oldelem;
2213 PyObject *pools = lz->pools;
2214 PyObject *result = lz->result;
2215 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2216 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002217
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002218 if (lz->stopped)
2219 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002220
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002221 if (result == NULL) {
2222 /* On the first pass, return an initial tuple filled with the
2223 first element from each pool. */
2224 result = PyTuple_New(npools);
2225 if (result == NULL)
2226 goto empty;
2227 lz->result = result;
2228 for (i=0; i < npools; i++) {
2229 pool = PyTuple_GET_ITEM(pools, i);
2230 if (PyTuple_GET_SIZE(pool) == 0)
2231 goto empty;
2232 elem = PyTuple_GET_ITEM(pool, 0);
2233 Py_INCREF(elem);
2234 PyTuple_SET_ITEM(result, i, elem);
2235 }
2236 } else {
2237 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002238
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002239 /* Copy the previous result tuple or re-use it if available */
2240 if (Py_REFCNT(result) > 1) {
2241 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002242 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002243 if (result == NULL)
2244 goto empty;
2245 lz->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002246 Py_DECREF(old_result);
2247 }
2248 /* Now, we've got the only copy so we can update it in-place */
2249 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002250
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002251 /* Update the pool indices right-to-left. Only advance to the
2252 next pool when the previous one rolls-over */
2253 for (i=npools-1 ; i >= 0 ; i--) {
2254 pool = PyTuple_GET_ITEM(pools, i);
2255 indices[i]++;
2256 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2257 /* Roll-over and advance to next pool */
2258 indices[i] = 0;
2259 elem = PyTuple_GET_ITEM(pool, 0);
2260 Py_INCREF(elem);
2261 oldelem = PyTuple_GET_ITEM(result, i);
2262 PyTuple_SET_ITEM(result, i, elem);
2263 Py_DECREF(oldelem);
2264 } else {
2265 /* No rollover. Just increment and stop here. */
2266 elem = PyTuple_GET_ITEM(pool, indices[i]);
2267 Py_INCREF(elem);
2268 oldelem = PyTuple_GET_ITEM(result, i);
2269 PyTuple_SET_ITEM(result, i, elem);
2270 Py_DECREF(oldelem);
2271 break;
2272 }
2273 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002274
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002275 /* If i is negative, then the indices have all rolled-over
2276 and we're done. */
2277 if (i < 0)
2278 goto empty;
2279 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002280
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002281 Py_INCREF(result);
2282 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002283
2284empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002285 lz->stopped = 1;
2286 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002287}
2288
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002289static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302290product_reduce(productobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002291{
2292 if (lz->stopped) {
2293 return Py_BuildValue("O(())", Py_TYPE(lz));
2294 } else if (lz->result == NULL) {
2295 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2296 } else {
2297 PyObject *indices;
2298 Py_ssize_t n, i;
2299
2300 /* we must pickle the indices use them for setstate, and
2301 * additionally indicate that the iterator has started
2302 */
2303 n = PyTuple_GET_SIZE(lz->pools);
2304 indices = PyTuple_New(n);
2305 if (indices == NULL)
2306 return NULL;
2307 for (i=0; i<n; i++){
2308 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2309 if (!index) {
2310 Py_DECREF(indices);
2311 return NULL;
2312 }
2313 PyTuple_SET_ITEM(indices, i, index);
2314 }
2315 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2316 }
2317}
2318
2319static PyObject *
2320product_setstate(productobject *lz, PyObject *state)
2321{
2322 PyObject *result;
2323 Py_ssize_t n, i;
2324
2325 n = PyTuple_GET_SIZE(lz->pools);
2326 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2327 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2328 return NULL;
2329 }
2330 for (i=0; i<n; i++)
2331 {
2332 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2333 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002334 PyObject* pool;
2335 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002336 if (index < 0 && PyErr_Occurred())
2337 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002338 pool = PyTuple_GET_ITEM(lz->pools, i);
2339 poolsize = PyTuple_GET_SIZE(pool);
2340 if (poolsize == 0) {
2341 lz->stopped = 1;
2342 Py_RETURN_NONE;
2343 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002344 /* clamp the index */
2345 if (index < 0)
2346 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002347 else if (index > poolsize-1)
2348 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002349 lz->indices[i] = index;
2350 }
2351
2352 result = PyTuple_New(n);
2353 if (!result)
2354 return NULL;
2355 for (i=0; i<n; i++) {
2356 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2357 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2358 Py_INCREF(element);
2359 PyTuple_SET_ITEM(result, i, element);
2360 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002361 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002362 Py_RETURN_NONE;
2363}
2364
2365static PyMethodDef product_methods[] = {
2366 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2367 reduce_doc},
2368 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2369 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002370 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2371 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002372 {NULL, NULL} /* sentinel */
2373};
2374
Christian Heimesc3f30c42008-02-22 16:37:40 +00002375PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002376"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002377\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002378Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002379For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2380The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2381cycle in a manner similar to an odometer (with the rightmost element changing\n\
2382on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002383To compute the product of an iterable with itself, specify the number\n\
2384of repetitions with the optional repeat keyword argument. For example,\n\
2385product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002386product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2387product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2388
2389static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002390 PyVarObject_HEAD_INIT(NULL, 0)
2391 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002392 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002393 0, /* tp_itemsize */
2394 /* methods */
2395 (destructor)product_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002396 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002397 0, /* tp_getattr */
2398 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002399 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002400 0, /* tp_repr */
2401 0, /* tp_as_number */
2402 0, /* tp_as_sequence */
2403 0, /* tp_as_mapping */
2404 0, /* tp_hash */
2405 0, /* tp_call */
2406 0, /* tp_str */
2407 PyObject_GenericGetAttr, /* tp_getattro */
2408 0, /* tp_setattro */
2409 0, /* tp_as_buffer */
2410 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2411 Py_TPFLAGS_BASETYPE, /* tp_flags */
2412 product_doc, /* tp_doc */
2413 (traverseproc)product_traverse, /* tp_traverse */
2414 0, /* tp_clear */
2415 0, /* tp_richcompare */
2416 0, /* tp_weaklistoffset */
2417 PyObject_SelfIter, /* tp_iter */
2418 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002419 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002420 0, /* tp_members */
2421 0, /* tp_getset */
2422 0, /* tp_base */
2423 0, /* tp_dict */
2424 0, /* tp_descr_get */
2425 0, /* tp_descr_set */
2426 0, /* tp_dictoffset */
2427 0, /* tp_init */
2428 0, /* tp_alloc */
2429 product_new, /* tp_new */
2430 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002431};
2432
2433
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002434/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002435
2436typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002437 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002438 PyObject *pool; /* input converted to a tuple */
2439 Py_ssize_t *indices; /* one index per result element */
2440 PyObject *result; /* most recently returned result tuple */
2441 Py_ssize_t r; /* size of result tuple */
2442 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002443} combinationsobject;
2444
Tal Einatc4bccd32018-09-12 00:49:13 +03002445
2446/*[clinic input]
2447@classmethod
2448itertools.combinations.__new__
2449 iterable: object
2450 r: Py_ssize_t
2451Return successive r-length combinations of elements in the iterable.
2452
2453combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)
2454[clinic start generated code]*/
2455
Christian Heimes380f7f22008-02-28 11:19:05 +00002456static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002457itertools_combinations_impl(PyTypeObject *type, PyObject *iterable,
2458 Py_ssize_t r)
2459/*[clinic end generated code: output=87a689b39c40039c input=06bede09e3da20f8]*/
Christian Heimes380f7f22008-02-28 11:19:05 +00002460{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002461 combinationsobject *co;
2462 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002463 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002464 Py_ssize_t *indices = NULL;
2465 Py_ssize_t i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002466
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002467 pool = PySequence_Tuple(iterable);
2468 if (pool == NULL)
2469 goto error;
2470 n = PyTuple_GET_SIZE(pool);
2471 if (r < 0) {
2472 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2473 goto error;
2474 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002475
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002476 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002477 if (indices == NULL) {
2478 PyErr_NoMemory();
2479 goto error;
2480 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002481
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002482 for (i=0 ; i<r ; i++)
2483 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002484
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002485 /* create combinationsobject structure */
2486 co = (combinationsobject *)type->tp_alloc(type, 0);
2487 if (co == NULL)
2488 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002489
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002490 co->pool = pool;
2491 co->indices = indices;
2492 co->result = NULL;
2493 co->r = r;
2494 co->stopped = r > n ? 1 : 0;
2495
2496 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002497
2498error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002499 if (indices != NULL)
2500 PyMem_Free(indices);
2501 Py_XDECREF(pool);
2502 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002503}
2504
2505static void
2506combinations_dealloc(combinationsobject *co)
2507{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002508 PyObject_GC_UnTrack(co);
2509 Py_XDECREF(co->pool);
2510 Py_XDECREF(co->result);
2511 if (co->indices != NULL)
2512 PyMem_Free(co->indices);
2513 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002514}
2515
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002516static PyObject *
2517combinations_sizeof(combinationsobject *co, void *unused)
2518{
2519 Py_ssize_t res;
2520
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002521 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002522 res += co->r * sizeof(Py_ssize_t);
2523 return PyLong_FromSsize_t(res);
2524}
2525
Christian Heimes380f7f22008-02-28 11:19:05 +00002526static int
2527combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2528{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002529 Py_VISIT(co->pool);
2530 Py_VISIT(co->result);
2531 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002532}
2533
2534static PyObject *
2535combinations_next(combinationsobject *co)
2536{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002537 PyObject *elem;
2538 PyObject *oldelem;
2539 PyObject *pool = co->pool;
2540 Py_ssize_t *indices = co->indices;
2541 PyObject *result = co->result;
2542 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2543 Py_ssize_t r = co->r;
2544 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002545
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002546 if (co->stopped)
2547 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002548
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002549 if (result == NULL) {
2550 /* On the first pass, initialize result tuple using the indices */
2551 result = PyTuple_New(r);
2552 if (result == NULL)
2553 goto empty;
2554 co->result = result;
2555 for (i=0; i<r ; i++) {
2556 index = indices[i];
2557 elem = PyTuple_GET_ITEM(pool, index);
2558 Py_INCREF(elem);
2559 PyTuple_SET_ITEM(result, i, elem);
2560 }
2561 } else {
2562 /* Copy the previous result tuple or re-use it if available */
2563 if (Py_REFCNT(result) > 1) {
2564 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002565 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002566 if (result == NULL)
2567 goto empty;
2568 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002569 Py_DECREF(old_result);
2570 }
2571 /* Now, we've got the only copy so we can update it in-place
2572 * CPython's empty tuple is a singleton and cached in
2573 * PyTuple's freelist.
2574 */
2575 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002576
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002577 /* Scan indices right-to-left until finding one that is not
2578 at its maximum (i + n - r). */
2579 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2580 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002581
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002582 /* If i is negative, then the indices are all at
2583 their maximum value and we're done. */
2584 if (i < 0)
2585 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002586
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002587 /* Increment the current index which we know is not at its
2588 maximum. Then move back to the right setting each index
2589 to its lowest possible value (one higher than the index
2590 to its left -- this maintains the sort order invariant). */
2591 indices[i]++;
2592 for (j=i+1 ; j<r ; j++)
2593 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002594
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002595 /* Update the result tuple for the new indices
2596 starting with i, the leftmost index that changed */
2597 for ( ; i<r ; i++) {
2598 index = indices[i];
2599 elem = PyTuple_GET_ITEM(pool, index);
2600 Py_INCREF(elem);
2601 oldelem = PyTuple_GET_ITEM(result, i);
2602 PyTuple_SET_ITEM(result, i, elem);
2603 Py_DECREF(oldelem);
2604 }
2605 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002606
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002607 Py_INCREF(result);
2608 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002609
2610empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002611 co->stopped = 1;
2612 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002613}
2614
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002615static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302616combinations_reduce(combinationsobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002617{
2618 if (lz->result == NULL) {
2619 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2620 } else if (lz->stopped) {
2621 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2622 } else {
2623 PyObject *indices;
2624 Py_ssize_t i;
2625
2626 /* we must pickle the indices and use them for setstate */
2627 indices = PyTuple_New(lz->r);
2628 if (!indices)
2629 return NULL;
2630 for (i=0; i<lz->r; i++)
2631 {
2632 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2633 if (!index) {
2634 Py_DECREF(indices);
2635 return NULL;
2636 }
2637 PyTuple_SET_ITEM(indices, i, index);
2638 }
2639
2640 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2641 }
2642}
2643
2644static PyObject *
2645combinations_setstate(combinationsobject *lz, PyObject *state)
2646{
2647 PyObject *result;
2648 Py_ssize_t i;
2649 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2650
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002651 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002652 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2653 return NULL;
2654 }
2655
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002656 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002657 Py_ssize_t max;
2658 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2659 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002660
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002661 if (index == -1 && PyErr_Occurred())
2662 return NULL; /* not an integer */
2663 max = i + n - lz->r;
2664 /* clamp the index (beware of negative max) */
2665 if (index > max)
2666 index = max;
2667 if (index < 0)
2668 index = 0;
2669 lz->indices[i] = index;
2670 }
2671
2672 result = PyTuple_New(lz->r);
2673 if (result == NULL)
2674 return NULL;
2675 for (i=0; i<lz->r; i++) {
2676 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2677 Py_INCREF(element);
2678 PyTuple_SET_ITEM(result, i, element);
2679 }
2680
Serhiy Storchaka48842712016-04-06 09:45:48 +03002681 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002682 Py_RETURN_NONE;
2683}
2684
2685static PyMethodDef combinations_methods[] = {
2686 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2687 reduce_doc},
2688 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2689 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002690 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2691 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002692 {NULL, NULL} /* sentinel */
2693};
2694
Christian Heimes380f7f22008-02-28 11:19:05 +00002695static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002696 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002697 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002698 sizeof(combinationsobject), /* tp_basicsize */
2699 0, /* tp_itemsize */
2700 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002701 (destructor)combinations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002702 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002703 0, /* tp_getattr */
2704 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002705 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002706 0, /* tp_repr */
2707 0, /* tp_as_number */
2708 0, /* tp_as_sequence */
2709 0, /* tp_as_mapping */
2710 0, /* tp_hash */
2711 0, /* tp_call */
2712 0, /* tp_str */
2713 PyObject_GenericGetAttr, /* tp_getattro */
2714 0, /* tp_setattro */
2715 0, /* tp_as_buffer */
2716 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2717 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03002718 itertools_combinations__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002719 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002720 0, /* tp_clear */
2721 0, /* tp_richcompare */
2722 0, /* tp_weaklistoffset */
2723 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002724 (iternextfunc)combinations_next, /* tp_iternext */
2725 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002726 0, /* tp_members */
2727 0, /* tp_getset */
2728 0, /* tp_base */
2729 0, /* tp_dict */
2730 0, /* tp_descr_get */
2731 0, /* tp_descr_set */
2732 0, /* tp_dictoffset */
2733 0, /* tp_init */
2734 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03002735 itertools_combinations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002736 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002737};
2738
2739
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002740/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002741
2742/* Equivalent to:
2743
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002744 def combinations_with_replacement(iterable, r):
2745 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2746 # number items returned: (n+r-1)! / r! / (n-1)!
2747 pool = tuple(iterable)
2748 n = len(pool)
2749 indices = [0] * r
2750 yield tuple(pool[i] for i in indices)
2751 while 1:
2752 for i in reversed(range(r)):
2753 if indices[i] != n - 1:
2754 break
2755 else:
2756 return
2757 indices[i:] = [indices[i] + 1] * (r - i)
2758 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002759
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002760 def combinations_with_replacement2(iterable, r):
2761 'Alternate version that filters from product()'
2762 pool = tuple(iterable)
2763 n = len(pool)
2764 for indices in product(range(n), repeat=r):
2765 if sorted(indices) == list(indices):
2766 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002767*/
2768typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002769 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002770 PyObject *pool; /* input converted to a tuple */
2771 Py_ssize_t *indices; /* one index per result element */
2772 PyObject *result; /* most recently returned result tuple */
2773 Py_ssize_t r; /* size of result tuple */
2774 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002775} cwrobject;
2776
Tal Einatc4bccd32018-09-12 00:49:13 +03002777/*[clinic input]
2778@classmethod
2779itertools.combinations_with_replacement.__new__
2780 iterable: object
2781 r: Py_ssize_t
2782Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.
2783
2784combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2785[clinic start generated code]*/
2786
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002787static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002788itertools_combinations_with_replacement_impl(PyTypeObject *type,
2789 PyObject *iterable,
2790 Py_ssize_t r)
2791/*[clinic end generated code: output=48b26856d4e659ca input=dc2a8c7ba785fad7]*/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002792{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002793 cwrobject *co;
2794 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002795 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002796 Py_ssize_t *indices = NULL;
2797 Py_ssize_t i;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002798
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002799 pool = PySequence_Tuple(iterable);
2800 if (pool == NULL)
2801 goto error;
2802 n = PyTuple_GET_SIZE(pool);
2803 if (r < 0) {
2804 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2805 goto error;
2806 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002807
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002808 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002809 if (indices == NULL) {
2810 PyErr_NoMemory();
2811 goto error;
2812 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002813
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002814 for (i=0 ; i<r ; i++)
2815 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002816
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002817 /* create cwrobject structure */
2818 co = (cwrobject *)type->tp_alloc(type, 0);
2819 if (co == NULL)
2820 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002821
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002822 co->pool = pool;
2823 co->indices = indices;
2824 co->result = NULL;
2825 co->r = r;
2826 co->stopped = !n && r;
2827
2828 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002829
2830error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002831 if (indices != NULL)
2832 PyMem_Free(indices);
2833 Py_XDECREF(pool);
2834 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002835}
2836
2837static void
2838cwr_dealloc(cwrobject *co)
2839{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002840 PyObject_GC_UnTrack(co);
2841 Py_XDECREF(co->pool);
2842 Py_XDECREF(co->result);
2843 if (co->indices != NULL)
2844 PyMem_Free(co->indices);
2845 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002846}
2847
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002848static PyObject *
2849cwr_sizeof(cwrobject *co, void *unused)
2850{
2851 Py_ssize_t res;
2852
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002853 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002854 res += co->r * sizeof(Py_ssize_t);
2855 return PyLong_FromSsize_t(res);
2856}
2857
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002858static int
2859cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2860{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002861 Py_VISIT(co->pool);
2862 Py_VISIT(co->result);
2863 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002864}
2865
2866static PyObject *
2867cwr_next(cwrobject *co)
2868{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002869 PyObject *elem;
2870 PyObject *oldelem;
2871 PyObject *pool = co->pool;
2872 Py_ssize_t *indices = co->indices;
2873 PyObject *result = co->result;
2874 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2875 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002876 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002877
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002878 if (co->stopped)
2879 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002880
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002881 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002882 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002883 result = PyTuple_New(r);
2884 if (result == NULL)
2885 goto empty;
2886 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002887 if (n > 0) {
2888 elem = PyTuple_GET_ITEM(pool, 0);
2889 for (i=0; i<r ; i++) {
2890 assert(indices[i] == 0);
2891 Py_INCREF(elem);
2892 PyTuple_SET_ITEM(result, i, elem);
2893 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002894 }
2895 } else {
2896 /* Copy the previous result tuple or re-use it if available */
2897 if (Py_REFCNT(result) > 1) {
2898 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002899 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002900 if (result == NULL)
2901 goto empty;
2902 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002903 Py_DECREF(old_result);
2904 }
2905 /* Now, we've got the only copy so we can update it in-place CPython's
2906 empty tuple is a singleton and cached in PyTuple's freelist. */
2907 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002908
Tim Peters9edb1682013-09-03 11:49:31 -05002909 /* Scan indices right-to-left until finding one that is not
2910 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002911 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2912 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002913
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002914 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002915 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002916 if (i < 0)
2917 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002918
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002919 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002920 maximum. Then set all to the right to the same value. */
2921 index = indices[i] + 1;
2922 assert(index < n);
2923 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002924 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002925 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002926 Py_INCREF(elem);
2927 oldelem = PyTuple_GET_ITEM(result, i);
2928 PyTuple_SET_ITEM(result, i, elem);
2929 Py_DECREF(oldelem);
2930 }
2931 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002932
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002933 Py_INCREF(result);
2934 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002935
2936empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002937 co->stopped = 1;
2938 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002939}
2940
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002941static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302942cwr_reduce(cwrobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002943{
2944 if (lz->result == NULL) {
2945 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2946 } else if (lz->stopped) {
2947 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2948 } else {
2949 PyObject *indices;
2950 Py_ssize_t i;
2951
2952 /* we must pickle the indices and use them for setstate */
2953 indices = PyTuple_New(lz->r);
2954 if (!indices)
2955 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002956 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002957 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2958 if (!index) {
2959 Py_DECREF(indices);
2960 return NULL;
2961 }
2962 PyTuple_SET_ITEM(indices, i, index);
2963 }
2964
2965 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2966 }
2967}
2968
2969static PyObject *
2970cwr_setstate(cwrobject *lz, PyObject *state)
2971{
2972 PyObject *result;
2973 Py_ssize_t n, i;
2974
2975 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2976 {
2977 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2978 return NULL;
2979 }
2980
2981 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002982 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002983 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2984 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002985
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002986 if (index < 0 && PyErr_Occurred())
2987 return NULL; /* not an integer */
2988 /* clamp the index */
2989 if (index < 0)
2990 index = 0;
2991 else if (index > n-1)
2992 index = n-1;
2993 lz->indices[i] = index;
2994 }
2995 result = PyTuple_New(lz->r);
2996 if (result == NULL)
2997 return NULL;
2998 for (i=0; i<lz->r; i++) {
2999 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
3000 Py_INCREF(element);
3001 PyTuple_SET_ITEM(result, i, element);
3002 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003003 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003004 Py_RETURN_NONE;
3005}
3006
3007static PyMethodDef cwr_methods[] = {
3008 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
3009 reduce_doc},
3010 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
3011 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003012 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
3013 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003014 {NULL, NULL} /* sentinel */
3015};
3016
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003017static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003018 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003019 "itertools.combinations_with_replacement", /* tp_name */
3020 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003021 0, /* tp_itemsize */
3022 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003023 (destructor)cwr_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003024 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003025 0, /* tp_getattr */
3026 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003027 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003028 0, /* tp_repr */
3029 0, /* tp_as_number */
3030 0, /* tp_as_sequence */
3031 0, /* tp_as_mapping */
3032 0, /* tp_hash */
3033 0, /* tp_call */
3034 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003035 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003036 0, /* tp_setattro */
3037 0, /* tp_as_buffer */
3038 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003039 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003040 itertools_combinations_with_replacement__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003041 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003042 0, /* tp_clear */
3043 0, /* tp_richcompare */
3044 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003045 PyObject_SelfIter, /* tp_iter */
3046 (iternextfunc)cwr_next, /* tp_iternext */
3047 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003048 0, /* tp_members */
3049 0, /* tp_getset */
3050 0, /* tp_base */
3051 0, /* tp_dict */
3052 0, /* tp_descr_get */
3053 0, /* tp_descr_set */
3054 0, /* tp_dictoffset */
3055 0, /* tp_init */
3056 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003057 itertools_combinations_with_replacement, /* tp_new */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003058 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003059};
3060
3061
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003062/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003063
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003064def permutations(iterable, r=None):
Sergey0078a0c2019-10-29 08:10:24 +03003065 # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
3066 # permutations(range(3)) --> 012 021 102 120 201 210
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003067 pool = tuple(iterable)
3068 n = len(pool)
3069 r = n if r is None else r
Sergey0078a0c2019-10-29 08:10:24 +03003070 if r > n:
3071 return
3072 indices = list(range(n))
3073 cycles = list(range(n, n-r, -1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003074 yield tuple(pool[i] for i in indices[:r])
3075 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003076 for i in reversed(range(r)):
3077 cycles[i] -= 1
3078 if cycles[i] == 0:
3079 indices[i:] = indices[i+1:] + indices[i:i+1]
3080 cycles[i] = n - i
3081 else:
3082 j = cycles[i]
3083 indices[i], indices[-j] = indices[-j], indices[i]
3084 yield tuple(pool[i] for i in indices[:r])
3085 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003086 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003087 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003088*/
3089
3090typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003091 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003092 PyObject *pool; /* input converted to a tuple */
3093 Py_ssize_t *indices; /* one index per element in the pool */
3094 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3095 PyObject *result; /* most recently returned result tuple */
3096 Py_ssize_t r; /* size of result tuple */
3097 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003098} permutationsobject;
3099
Tal Einatc4bccd32018-09-12 00:49:13 +03003100/*[clinic input]
3101@classmethod
3102itertools.permutations.__new__
3103 iterable: object
3104 r as robj: object = None
3105Return successive r-length permutations of elements in the iterable.
3106
3107permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
3108[clinic start generated code]*/
3109
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003110static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003111itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
3112 PyObject *robj)
3113/*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003114{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003115 permutationsobject *po;
3116 Py_ssize_t n;
3117 Py_ssize_t r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003118 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003119 Py_ssize_t *indices = NULL;
3120 Py_ssize_t *cycles = NULL;
3121 Py_ssize_t i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003122
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003123 pool = PySequence_Tuple(iterable);
3124 if (pool == NULL)
3125 goto error;
3126 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003127
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003128 r = n;
3129 if (robj != Py_None) {
3130 if (!PyLong_Check(robj)) {
3131 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3132 goto error;
3133 }
3134 r = PyLong_AsSsize_t(robj);
3135 if (r == -1 && PyErr_Occurred())
3136 goto error;
3137 }
3138 if (r < 0) {
3139 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3140 goto error;
3141 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003142
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003143 indices = PyMem_New(Py_ssize_t, n);
3144 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003145 if (indices == NULL || cycles == NULL) {
3146 PyErr_NoMemory();
3147 goto error;
3148 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003149
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003150 for (i=0 ; i<n ; i++)
3151 indices[i] = i;
3152 for (i=0 ; i<r ; i++)
3153 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003154
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003155 /* create permutationsobject structure */
3156 po = (permutationsobject *)type->tp_alloc(type, 0);
3157 if (po == NULL)
3158 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003159
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003160 po->pool = pool;
3161 po->indices = indices;
3162 po->cycles = cycles;
3163 po->result = NULL;
3164 po->r = r;
3165 po->stopped = r > n ? 1 : 0;
3166
3167 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003168
3169error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003170 if (indices != NULL)
3171 PyMem_Free(indices);
3172 if (cycles != NULL)
3173 PyMem_Free(cycles);
3174 Py_XDECREF(pool);
3175 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003176}
3177
3178static void
3179permutations_dealloc(permutationsobject *po)
3180{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003181 PyObject_GC_UnTrack(po);
3182 Py_XDECREF(po->pool);
3183 Py_XDECREF(po->result);
3184 PyMem_Free(po->indices);
3185 PyMem_Free(po->cycles);
3186 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003187}
3188
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003189static PyObject *
3190permutations_sizeof(permutationsobject *po, void *unused)
3191{
3192 Py_ssize_t res;
3193
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003194 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003195 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3196 res += po->r * sizeof(Py_ssize_t);
3197 return PyLong_FromSsize_t(res);
3198}
3199
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003200static int
3201permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3202{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003203 Py_VISIT(po->pool);
3204 Py_VISIT(po->result);
3205 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003206}
3207
3208static PyObject *
3209permutations_next(permutationsobject *po)
3210{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003211 PyObject *elem;
3212 PyObject *oldelem;
3213 PyObject *pool = po->pool;
3214 Py_ssize_t *indices = po->indices;
3215 Py_ssize_t *cycles = po->cycles;
3216 PyObject *result = po->result;
3217 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3218 Py_ssize_t r = po->r;
3219 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003220
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003221 if (po->stopped)
3222 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003223
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003224 if (result == NULL) {
3225 /* On the first pass, initialize result tuple using the indices */
3226 result = PyTuple_New(r);
3227 if (result == NULL)
3228 goto empty;
3229 po->result = result;
3230 for (i=0; i<r ; i++) {
3231 index = indices[i];
3232 elem = PyTuple_GET_ITEM(pool, index);
3233 Py_INCREF(elem);
3234 PyTuple_SET_ITEM(result, i, elem);
3235 }
3236 } else {
3237 if (n == 0)
3238 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003239
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003240 /* Copy the previous result tuple or re-use it if available */
3241 if (Py_REFCNT(result) > 1) {
3242 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05003243 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003244 if (result == NULL)
3245 goto empty;
3246 po->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003247 Py_DECREF(old_result);
3248 }
3249 /* Now, we've got the only copy so we can update it in-place */
3250 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003251
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003252 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3253 for (i=r-1 ; i>=0 ; i--) {
3254 cycles[i] -= 1;
3255 if (cycles[i] == 0) {
3256 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3257 index = indices[i];
3258 for (j=i ; j<n-1 ; j++)
3259 indices[j] = indices[j+1];
3260 indices[n-1] = index;
3261 cycles[i] = n - i;
3262 } else {
3263 j = cycles[i];
3264 index = indices[i];
3265 indices[i] = indices[n-j];
3266 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003267
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003268 for (k=i; k<r ; k++) {
3269 /* start with i, the leftmost element that changed */
3270 /* yield tuple(pool[k] for k in indices[:r]) */
3271 index = indices[k];
3272 elem = PyTuple_GET_ITEM(pool, index);
3273 Py_INCREF(elem);
3274 oldelem = PyTuple_GET_ITEM(result, k);
3275 PyTuple_SET_ITEM(result, k, elem);
3276 Py_DECREF(oldelem);
3277 }
3278 break;
3279 }
3280 }
3281 /* If i is negative, then the cycles have all
3282 rolled-over and we're done. */
3283 if (i < 0)
3284 goto empty;
3285 }
3286 Py_INCREF(result);
3287 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003288
3289empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003290 po->stopped = 1;
3291 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003292}
3293
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003294static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303295permutations_reduce(permutationsobject *po, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003296{
3297 if (po->result == NULL) {
3298 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3299 } else if (po->stopped) {
3300 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3301 } else {
3302 PyObject *indices=NULL, *cycles=NULL;
3303 Py_ssize_t n, i;
3304
3305 /* we must pickle the indices and cycles and use them for setstate */
3306 n = PyTuple_GET_SIZE(po->pool);
3307 indices = PyTuple_New(n);
3308 if (indices == NULL)
3309 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003310 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003311 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3312 if (!index)
3313 goto err;
3314 PyTuple_SET_ITEM(indices, i, index);
3315 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003316
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003317 cycles = PyTuple_New(po->r);
3318 if (cycles == NULL)
3319 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003320 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003321 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3322 if (!index)
3323 goto err;
3324 PyTuple_SET_ITEM(cycles, i, index);
3325 }
3326 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3327 po->pool, po->r,
3328 indices, cycles);
3329 err:
3330 Py_XDECREF(indices);
3331 Py_XDECREF(cycles);
3332 return NULL;
3333 }
3334}
3335
3336static PyObject *
3337permutations_setstate(permutationsobject *po, PyObject *state)
3338{
3339 PyObject *indices, *cycles, *result;
3340 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003341
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003342 if (!PyTuple_Check(state)) {
3343 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3344 return NULL;
3345 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003346 if (!PyArg_ParseTuple(state, "O!O!",
3347 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003348 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003349 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003350 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003351
3352 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003353 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003354 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3355 return NULL;
3356 }
3357
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003358 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003359 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3360 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3361 if (index < 0 && PyErr_Occurred())
3362 return NULL; /* not an integer */
3363 /* clamp the index */
3364 if (index < 0)
3365 index = 0;
3366 else if (index > n-1)
3367 index = n-1;
3368 po->indices[i] = index;
3369 }
3370
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003371 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003372 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3373 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3374 if (index < 0 && PyErr_Occurred())
3375 return NULL; /* not an integer */
3376 if (index < 1)
3377 index = 1;
3378 else if (index > n-i)
3379 index = n-i;
3380 po->cycles[i] = index;
3381 }
3382 result = PyTuple_New(po->r);
3383 if (result == NULL)
3384 return NULL;
3385 for (i=0; i<po->r; i++) {
3386 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3387 Py_INCREF(element);
3388 PyTuple_SET_ITEM(result, i, element);
3389 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003390 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003391 Py_RETURN_NONE;
3392}
3393
3394static PyMethodDef permuations_methods[] = {
3395 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3396 reduce_doc},
3397 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3398 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003399 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3400 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003401 {NULL, NULL} /* sentinel */
3402};
3403
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003404static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003405 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003406 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003407 sizeof(permutationsobject), /* tp_basicsize */
3408 0, /* tp_itemsize */
3409 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003410 (destructor)permutations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003411 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003412 0, /* tp_getattr */
3413 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003414 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003415 0, /* tp_repr */
3416 0, /* tp_as_number */
3417 0, /* tp_as_sequence */
3418 0, /* tp_as_mapping */
3419 0, /* tp_hash */
3420 0, /* tp_call */
3421 0, /* tp_str */
3422 PyObject_GenericGetAttr, /* tp_getattro */
3423 0, /* tp_setattro */
3424 0, /* tp_as_buffer */
3425 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3426 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003427 itertools_permutations__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003428 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003429 0, /* tp_clear */
3430 0, /* tp_richcompare */
3431 0, /* tp_weaklistoffset */
3432 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003433 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003434 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003435 0, /* tp_members */
3436 0, /* tp_getset */
3437 0, /* tp_base */
3438 0, /* tp_dict */
3439 0, /* tp_descr_get */
3440 0, /* tp_descr_set */
3441 0, /* tp_dictoffset */
3442 0, /* tp_init */
3443 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003444 itertools_permutations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003445 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003446};
3447
Tal Einatc4bccd32018-09-12 00:49:13 +03003448
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003449/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003450
3451typedef struct {
3452 PyObject_HEAD
3453 PyObject *total;
3454 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003455 PyObject *binop;
Lisa Roach9718b592018-09-23 17:34:59 -07003456 PyObject *initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003457} accumulateobject;
3458
Tal Einatc4bccd32018-09-12 00:49:13 +03003459/*[clinic input]
3460@classmethod
3461itertools.accumulate.__new__
3462 iterable: object
3463 func as binop: object = None
Lisa Roach9718b592018-09-23 17:34:59 -07003464 *
3465 initial: object = None
Tal Einatc4bccd32018-09-12 00:49:13 +03003466Return series of accumulated sums (or other binary function results).
3467[clinic start generated code]*/
3468
Raymond Hettinger482ba772010-12-01 22:48:00 +00003469static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003470itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
Lisa Roach9718b592018-09-23 17:34:59 -07003471 PyObject *binop, PyObject *initial)
3472/*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003473{
Raymond Hettinger482ba772010-12-01 22:48:00 +00003474 PyObject *it;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003475 accumulateobject *lz;
3476
Raymond Hettinger482ba772010-12-01 22:48:00 +00003477 /* Get iterator. */
3478 it = PyObject_GetIter(iterable);
3479 if (it == NULL)
3480 return NULL;
3481
Raymond Hettinger482ba772010-12-01 22:48:00 +00003482 /* create accumulateobject structure */
3483 lz = (accumulateobject *)type->tp_alloc(type, 0);
3484 if (lz == NULL) {
3485 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003486 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003487 }
3488
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003489 if (binop != Py_None) {
3490 Py_XINCREF(binop);
3491 lz->binop = binop;
3492 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003493 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003494 lz->it = it;
Lisa Roach9718b592018-09-23 17:34:59 -07003495 Py_XINCREF(initial);
3496 lz->initial = initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003497 return (PyObject *)lz;
3498}
3499
3500static void
3501accumulate_dealloc(accumulateobject *lz)
3502{
3503 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003504 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003505 Py_XDECREF(lz->total);
3506 Py_XDECREF(lz->it);
Lisa Roach9718b592018-09-23 17:34:59 -07003507 Py_XDECREF(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003508 Py_TYPE(lz)->tp_free(lz);
3509}
3510
3511static int
3512accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3513{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003514 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003515 Py_VISIT(lz->it);
3516 Py_VISIT(lz->total);
Lisa Roach9718b592018-09-23 17:34:59 -07003517 Py_VISIT(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003518 return 0;
3519}
3520
3521static PyObject *
3522accumulate_next(accumulateobject *lz)
3523{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003524 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003525
Lisa Roach9718b592018-09-23 17:34:59 -07003526 if (lz->initial != Py_None) {
3527 lz->total = lz->initial;
3528 Py_INCREF(Py_None);
3529 lz->initial = Py_None;
3530 Py_INCREF(lz->total);
3531 return lz->total;
3532 }
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003533 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003534 if (val == NULL)
3535 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003536
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003537 if (lz->total == NULL) {
3538 Py_INCREF(val);
3539 lz->total = val;
3540 return lz->total;
3541 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003542
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003543 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003544 newtotal = PyNumber_Add(lz->total, val);
3545 else
3546 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003547 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003548 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003549 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003550
Raymond Hettinger482ba772010-12-01 22:48:00 +00003551 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003552 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003553 return newtotal;
3554}
3555
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003556static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303557accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003558{
Lisa Roach9718b592018-09-23 17:34:59 -07003559 if (lz->initial != Py_None) {
3560 PyObject *it;
3561
3562 assert(lz->total == NULL);
3563 if (PyType_Ready(&chain_type) < 0)
3564 return NULL;
3565 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3566 lz->initial, lz->it);
3567 if (it == NULL)
3568 return NULL;
3569 return Py_BuildValue("O(NO)O", Py_TYPE(lz),
3570 it, lz->binop?lz->binop:Py_None, Py_None);
3571 }
Serhiy Storchakad5516252016-03-06 14:00:45 +02003572 if (lz->total == Py_None) {
3573 PyObject *it;
3574
3575 if (PyType_Ready(&chain_type) < 0)
3576 return NULL;
3577 if (PyType_Ready(&islice_type) < 0)
3578 return NULL;
3579 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3580 lz->total, lz->it);
3581 if (it == NULL)
3582 return NULL;
3583 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3584 it, lz->binop ? lz->binop : Py_None);
3585 if (it == NULL)
3586 return NULL;
3587 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3588 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003589 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3590 lz->it, lz->binop?lz->binop:Py_None,
3591 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003592}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003593
3594static PyObject *
3595accumulate_setstate(accumulateobject *lz, PyObject *state)
3596{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003597 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003598 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003599 Py_RETURN_NONE;
3600}
3601
3602static PyMethodDef accumulate_methods[] = {
3603 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3604 reduce_doc},
3605 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3606 setstate_doc},
3607 {NULL, NULL} /* sentinel */
3608};
3609
Raymond Hettinger482ba772010-12-01 22:48:00 +00003610static PyTypeObject accumulate_type = {
3611 PyVarObject_HEAD_INIT(NULL, 0)
3612 "itertools.accumulate", /* tp_name */
3613 sizeof(accumulateobject), /* tp_basicsize */
3614 0, /* tp_itemsize */
3615 /* methods */
3616 (destructor)accumulate_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003617 0, /* tp_vectorcall_offset */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003618 0, /* tp_getattr */
3619 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003620 0, /* tp_as_async */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003621 0, /* tp_repr */
3622 0, /* tp_as_number */
3623 0, /* tp_as_sequence */
3624 0, /* tp_as_mapping */
3625 0, /* tp_hash */
3626 0, /* tp_call */
3627 0, /* tp_str */
3628 PyObject_GenericGetAttr, /* tp_getattro */
3629 0, /* tp_setattro */
3630 0, /* tp_as_buffer */
3631 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3632 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003633 itertools_accumulate__doc__, /* tp_doc */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003634 (traverseproc)accumulate_traverse, /* tp_traverse */
3635 0, /* tp_clear */
3636 0, /* tp_richcompare */
3637 0, /* tp_weaklistoffset */
3638 PyObject_SelfIter, /* tp_iter */
3639 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003640 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003641 0, /* tp_members */
3642 0, /* tp_getset */
3643 0, /* tp_base */
3644 0, /* tp_dict */
3645 0, /* tp_descr_get */
3646 0, /* tp_descr_set */
3647 0, /* tp_dictoffset */
3648 0, /* tp_init */
3649 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003650 itertools_accumulate, /* tp_new */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003651 PyObject_GC_Del, /* tp_free */
3652};
3653
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003654
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003655/* compress object ************************************************************/
3656
3657/* Equivalent to:
3658
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003659 def compress(data, selectors):
3660 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3661 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003662*/
3663
3664typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003665 PyObject_HEAD
3666 PyObject *data;
3667 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003668} compressobject;
3669
Tal Einatc4bccd32018-09-12 00:49:13 +03003670/*[clinic input]
3671@classmethod
3672itertools.compress.__new__
3673 data as seq1: object
3674 selectors as seq2: object
3675Return data elements corresponding to true selector elements.
3676
3677Forms a shorter iterator from selected data elements using the selectors to
3678choose the data elements.
3679[clinic start generated code]*/
3680
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003681static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003682itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
3683/*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003684{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003685 PyObject *data=NULL, *selectors=NULL;
3686 compressobject *lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003687
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003688 data = PyObject_GetIter(seq1);
3689 if (data == NULL)
3690 goto fail;
3691 selectors = PyObject_GetIter(seq2);
3692 if (selectors == NULL)
3693 goto fail;
3694
3695 /* create compressobject structure */
3696 lz = (compressobject *)type->tp_alloc(type, 0);
3697 if (lz == NULL)
3698 goto fail;
3699 lz->data = data;
3700 lz->selectors = selectors;
3701 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003702
3703fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003704 Py_XDECREF(data);
3705 Py_XDECREF(selectors);
3706 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003707}
3708
3709static void
3710compress_dealloc(compressobject *lz)
3711{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003712 PyObject_GC_UnTrack(lz);
3713 Py_XDECREF(lz->data);
3714 Py_XDECREF(lz->selectors);
3715 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003716}
3717
3718static int
3719compress_traverse(compressobject *lz, visitproc visit, void *arg)
3720{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003721 Py_VISIT(lz->data);
3722 Py_VISIT(lz->selectors);
3723 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003724}
3725
3726static PyObject *
3727compress_next(compressobject *lz)
3728{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003729 PyObject *data = lz->data, *selectors = lz->selectors;
3730 PyObject *datum, *selector;
3731 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3732 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3733 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003734
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003735 while (1) {
3736 /* Steps: get datum, get selector, evaluate selector.
3737 Order is important (to match the pure python version
3738 in terms of which input gets a chance to raise an
3739 exception first).
3740 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003741
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003742 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003743 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003744 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003745
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003746 selector = selectornext(selectors);
3747 if (selector == NULL) {
3748 Py_DECREF(datum);
3749 return NULL;
3750 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003751
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003752 ok = PyObject_IsTrue(selector);
3753 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003754 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003755 return datum;
3756 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003757 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003758 return NULL;
3759 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003760}
3761
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003762static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303763compress_reduce(compressobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003764{
3765 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3766 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003767}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003768
3769static PyMethodDef compress_methods[] = {
3770 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3771 reduce_doc},
3772 {NULL, NULL} /* sentinel */
3773};
3774
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003775static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003776 PyVarObject_HEAD_INIT(NULL, 0)
3777 "itertools.compress", /* tp_name */
3778 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003779 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003780 /* methods */
3781 (destructor)compress_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003782 0, /* tp_vectorcall_offset */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003783 0, /* tp_getattr */
3784 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003785 0, /* tp_as_async */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003786 0, /* tp_repr */
3787 0, /* tp_as_number */
3788 0, /* tp_as_sequence */
3789 0, /* tp_as_mapping */
3790 0, /* tp_hash */
3791 0, /* tp_call */
3792 0, /* tp_str */
3793 PyObject_GenericGetAttr, /* tp_getattro */
3794 0, /* tp_setattro */
3795 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003796 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003797 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003798 itertools_compress__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003799 (traverseproc)compress_traverse, /* tp_traverse */
3800 0, /* tp_clear */
3801 0, /* tp_richcompare */
3802 0, /* tp_weaklistoffset */
3803 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003804 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003805 compress_methods, /* tp_methods */
3806 0, /* tp_members */
3807 0, /* tp_getset */
3808 0, /* tp_base */
3809 0, /* tp_dict */
3810 0, /* tp_descr_get */
3811 0, /* tp_descr_set */
3812 0, /* tp_dictoffset */
3813 0, /* tp_init */
3814 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003815 itertools_compress, /* tp_new */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003816 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003817};
3818
3819
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003820/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003821
3822typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003823 PyObject_HEAD
3824 PyObject *func;
3825 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003826} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003827
Tal Einatc4bccd32018-09-12 00:49:13 +03003828/*[clinic input]
3829@classmethod
3830itertools.filterfalse.__new__
3831 function as func: object
3832 iterable as seq: object
3833 /
3834Return those items of iterable for which function(item) is false.
3835
3836If function is None, return the items that are false.
3837[clinic start generated code]*/
3838
Raymond Hettinger60eca932003-02-09 06:40:58 +00003839static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003840itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
3841/*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003842{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003843 PyObject *it;
3844 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003845
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003846 /* Get iterator. */
3847 it = PyObject_GetIter(seq);
3848 if (it == NULL)
3849 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003850
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003851 /* create filterfalseobject structure */
3852 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3853 if (lz == NULL) {
3854 Py_DECREF(it);
3855 return NULL;
3856 }
3857 Py_INCREF(func);
3858 lz->func = func;
3859 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003860
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003861 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003862}
3863
3864static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003865filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003866{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003867 PyObject_GC_UnTrack(lz);
3868 Py_XDECREF(lz->func);
3869 Py_XDECREF(lz->it);
3870 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003871}
3872
3873static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003874filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003875{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003876 Py_VISIT(lz->it);
3877 Py_VISIT(lz->func);
3878 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003879}
3880
3881static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003882filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003883{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003884 PyObject *item;
3885 PyObject *it = lz->it;
3886 long ok;
3887 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003888
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003889 iternext = *Py_TYPE(it)->tp_iternext;
3890 for (;;) {
3891 item = iternext(it);
3892 if (item == NULL)
3893 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003894
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003895 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3896 ok = PyObject_IsTrue(item);
3897 } else {
3898 PyObject *good;
Petr Viktorinffd97532020-02-11 17:46:57 +01003899 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003900 if (good == NULL) {
3901 Py_DECREF(item);
3902 return NULL;
3903 }
3904 ok = PyObject_IsTrue(good);
3905 Py_DECREF(good);
3906 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003907 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003908 return item;
3909 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003910 if (ok < 0)
3911 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003912 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003913}
3914
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003915static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303916filterfalse_reduce(filterfalseobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003917{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003918 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3919}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003920
3921static PyMethodDef filterfalse_methods[] = {
3922 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3923 reduce_doc},
3924 {NULL, NULL} /* sentinel */
3925};
3926
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003927static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003928 PyVarObject_HEAD_INIT(NULL, 0)
3929 "itertools.filterfalse", /* tp_name */
3930 sizeof(filterfalseobject), /* tp_basicsize */
3931 0, /* tp_itemsize */
3932 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003933 (destructor)filterfalse_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003934 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003935 0, /* tp_getattr */
3936 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003937 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003938 0, /* tp_repr */
3939 0, /* tp_as_number */
3940 0, /* tp_as_sequence */
3941 0, /* tp_as_mapping */
3942 0, /* tp_hash */
3943 0, /* tp_call */
3944 0, /* tp_str */
3945 PyObject_GenericGetAttr, /* tp_getattro */
3946 0, /* tp_setattro */
3947 0, /* tp_as_buffer */
3948 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3949 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003950 itertools_filterfalse__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003951 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003952 0, /* tp_clear */
3953 0, /* tp_richcompare */
3954 0, /* tp_weaklistoffset */
3955 PyObject_SelfIter, /* tp_iter */
3956 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003957 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003958 0, /* tp_members */
3959 0, /* tp_getset */
3960 0, /* tp_base */
3961 0, /* tp_dict */
3962 0, /* tp_descr_get */
3963 0, /* tp_descr_set */
3964 0, /* tp_dictoffset */
3965 0, /* tp_init */
3966 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003967 itertools_filterfalse, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003968 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003969};
3970
3971
3972/* count object ************************************************************/
3973
3974typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003975 PyObject_HEAD
3976 Py_ssize_t cnt;
3977 PyObject *long_cnt;
3978 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003979} countobject;
3980
Raymond Hettinger30729212009-02-12 06:28:27 +00003981/* Counting logic and invariants:
3982
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003983fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003984
Serhiy Storchaka483405b2015-02-17 10:14:30 +02003985 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003986 Advances with: cnt += 1
3987 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003988
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003989slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003990
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003991 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3992 All counting is done with python objects (no overflows or underflows).
3993 Advances with: long_cnt += long_step
3994 Step may be zero -- effectively a slow version of repeat(cnt).
3995 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003996*/
3997
Tal Einatc4bccd32018-09-12 00:49:13 +03003998/*[clinic input]
3999@classmethod
4000itertools.count.__new__
4001 start as long_cnt: object(c_default="NULL") = 0
4002 step as long_step: object(c_default="NULL") = 1
4003Return a count object whose .__next__() method returns consecutive values.
4004
4005Equivalent to:
4006 def count(firstval=0, step=1):
4007 x = firstval
4008 while 1:
4009 yield x
4010 x += step
4011[clinic start generated code]*/
4012
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004013static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03004014itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
4015 PyObject *long_step)
4016/*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004017{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004018 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004019 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004020 Py_ssize_t cnt = 0;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00004021 long step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004022
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004023 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
4024 (long_step != NULL && !PyNumber_Check(long_step))) {
4025 PyErr_SetString(PyExc_TypeError, "a number is required");
4026 return NULL;
4027 }
Raymond Hettinger30729212009-02-12 06:28:27 +00004028
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004029 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
4030 (long_step == NULL || PyLong_Check(long_step));
4031
4032 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004033 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004034 if (fast_mode) {
4035 assert(PyLong_Check(long_cnt));
4036 cnt = PyLong_AsSsize_t(long_cnt);
4037 if (cnt == -1 && PyErr_Occurred()) {
4038 PyErr_Clear();
4039 fast_mode = 0;
4040 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004041 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004042 } else {
4043 cnt = 0;
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004044 long_cnt = _PyLong_Zero;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004045 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004046 Py_INCREF(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004047
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004048 /* If not specified, step defaults to 1 */
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004049 if (long_step == NULL)
4050 long_step = _PyLong_One;
4051 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004052
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004053 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004054
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004055 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004056 if (fast_mode) {
4057 assert(PyLong_Check(long_step));
4058 step = PyLong_AsLong(long_step);
4059 if (step != 1) {
4060 fast_mode = 0;
4061 if (step == -1 && PyErr_Occurred())
4062 PyErr_Clear();
4063 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004064 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004065
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004066 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004067 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004068 else
4069 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004070
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004071 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4072 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4073 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004074 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004075
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004076 /* create countobject structure */
4077 lz = (countobject *)type->tp_alloc(type, 0);
4078 if (lz == NULL) {
4079 Py_XDECREF(long_cnt);
Zackery Spytz0523c392019-03-26 00:05:29 -06004080 Py_DECREF(long_step);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004081 return NULL;
4082 }
4083 lz->cnt = cnt;
4084 lz->long_cnt = long_cnt;
4085 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004086
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004087 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004088}
4089
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004090static void
4091count_dealloc(countobject *lz)
4092{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004093 PyObject_GC_UnTrack(lz);
4094 Py_XDECREF(lz->long_cnt);
4095 Py_XDECREF(lz->long_step);
4096 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004097}
4098
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004099static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004100count_traverse(countobject *lz, visitproc visit, void *arg)
4101{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004102 Py_VISIT(lz->long_cnt);
4103 Py_VISIT(lz->long_step);
4104 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004105}
4106
4107static PyObject *
4108count_nextlong(countobject *lz)
4109{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004110 PyObject *long_cnt;
4111 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004112
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004113 long_cnt = lz->long_cnt;
4114 if (long_cnt == NULL) {
4115 /* Switch to slow_mode */
4116 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4117 if (long_cnt == NULL)
4118 return NULL;
4119 }
4120 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004121
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004122 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4123 if (stepped_up == NULL)
4124 return NULL;
4125 lz->long_cnt = stepped_up;
4126 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004127}
4128
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004129static PyObject *
4130count_next(countobject *lz)
4131{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004132 if (lz->cnt == PY_SSIZE_T_MAX)
4133 return count_nextlong(lz);
4134 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004135}
4136
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004137static PyObject *
4138count_repr(countobject *lz)
4139{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004140 if (lz->cnt != PY_SSIZE_T_MAX)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004141 return PyUnicode_FromFormat("%s(%zd)",
4142 _PyType_Name(Py_TYPE(lz)), lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004143
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004144 if (PyLong_Check(lz->long_step)) {
4145 long step = PyLong_AsLong(lz->long_step);
4146 if (step == -1 && PyErr_Occurred()) {
4147 PyErr_Clear();
4148 }
4149 if (step == 1) {
4150 /* Don't display step when it is an integer equal to 1 */
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004151 return PyUnicode_FromFormat("%s(%R)",
4152 _PyType_Name(Py_TYPE(lz)),
4153 lz->long_cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004154 }
4155 }
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004156 return PyUnicode_FromFormat("%s(%R, %R)",
4157 _PyType_Name(Py_TYPE(lz)),
4158 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004159}
4160
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004161static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304162count_reduce(countobject *lz, PyObject *Py_UNUSED(ignored))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004163{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004164 if (lz->cnt == PY_SSIZE_T_MAX)
4165 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4166 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004167}
4168
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004169static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004170 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004171 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004172 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004173};
4174
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004175static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004176 PyVarObject_HEAD_INIT(NULL, 0)
4177 "itertools.count", /* tp_name */
4178 sizeof(countobject), /* tp_basicsize */
4179 0, /* tp_itemsize */
4180 /* methods */
4181 (destructor)count_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004182 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004183 0, /* tp_getattr */
4184 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004185 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004186 (reprfunc)count_repr, /* tp_repr */
4187 0, /* tp_as_number */
4188 0, /* tp_as_sequence */
4189 0, /* tp_as_mapping */
4190 0, /* tp_hash */
4191 0, /* tp_call */
4192 0, /* tp_str */
4193 PyObject_GenericGetAttr, /* tp_getattro */
4194 0, /* tp_setattro */
4195 0, /* tp_as_buffer */
4196 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004197 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03004198 itertools_count__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004199 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004200 0, /* tp_clear */
4201 0, /* tp_richcompare */
4202 0, /* tp_weaklistoffset */
4203 PyObject_SelfIter, /* tp_iter */
4204 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004205 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004206 0, /* tp_members */
4207 0, /* tp_getset */
4208 0, /* tp_base */
4209 0, /* tp_dict */
4210 0, /* tp_descr_get */
4211 0, /* tp_descr_set */
4212 0, /* tp_dictoffset */
4213 0, /* tp_init */
4214 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03004215 itertools_count, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004216 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004217};
4218
4219
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004220/* repeat object ************************************************************/
4221
4222typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004223 PyObject_HEAD
4224 PyObject *element;
4225 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004226} repeatobject;
4227
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004228static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004229
4230static PyObject *
4231repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4232{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004233 repeatobject *ro;
4234 PyObject *element;
Raymond Hettingeraf464502019-11-09 20:28:31 -08004235 Py_ssize_t cnt = -1, n_args;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004236 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004237
Raymond Hettingeraf464502019-11-09 20:28:31 -08004238 n_args = PyTuple_GET_SIZE(args);
4239 if (kwds != NULL)
4240 n_args += PyDict_GET_SIZE(kwds);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004241 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4242 &element, &cnt))
4243 return NULL;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004244 /* Does user supply times argument? */
Raymond Hettingeraf464502019-11-09 20:28:31 -08004245 if (n_args == 2 && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004246 cnt = 0;
4247
4248 ro = (repeatobject *)type->tp_alloc(type, 0);
4249 if (ro == NULL)
4250 return NULL;
4251 Py_INCREF(element);
4252 ro->element = element;
4253 ro->cnt = cnt;
4254 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004255}
4256
4257static void
4258repeat_dealloc(repeatobject *ro)
4259{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004260 PyObject_GC_UnTrack(ro);
4261 Py_XDECREF(ro->element);
4262 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004263}
4264
4265static int
4266repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4267{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004268 Py_VISIT(ro->element);
4269 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004270}
4271
4272static PyObject *
4273repeat_next(repeatobject *ro)
4274{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004275 if (ro->cnt == 0)
4276 return NULL;
4277 if (ro->cnt > 0)
4278 ro->cnt--;
4279 Py_INCREF(ro->element);
4280 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004281}
4282
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004283static PyObject *
4284repeat_repr(repeatobject *ro)
4285{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004286 if (ro->cnt == -1)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004287 return PyUnicode_FromFormat("%s(%R)",
4288 _PyType_Name(Py_TYPE(ro)), ro->element);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004289 else
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004290 return PyUnicode_FromFormat("%s(%R, %zd)",
4291 _PyType_Name(Py_TYPE(ro)), ro->element,
4292 ro->cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004293}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004294
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004295static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304296repeat_len(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004297{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004298 if (ro->cnt == -1) {
4299 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4300 return NULL;
4301 }
4302 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004303}
4304
Armin Rigof5b3e362006-02-11 21:32:43 +00004305PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004306
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004307static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304308repeat_reduce(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004309{
4310 /* unpickle this so that a new repeat iterator is constructed with an
4311 * object, then call __setstate__ on it to set cnt
4312 */
4313 if (ro->cnt >= 0)
4314 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4315 else
4316 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4317}
4318
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004319static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004320 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004321 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004322 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004323};
4324
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004325PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004326"repeat(object [,times]) -> create an iterator which returns the object\n\
4327for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004328endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004329
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004330static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004331 PyVarObject_HEAD_INIT(NULL, 0)
4332 "itertools.repeat", /* tp_name */
4333 sizeof(repeatobject), /* tp_basicsize */
4334 0, /* tp_itemsize */
4335 /* methods */
4336 (destructor)repeat_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004337 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004338 0, /* tp_getattr */
4339 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004340 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004341 (reprfunc)repeat_repr, /* tp_repr */
4342 0, /* tp_as_number */
4343 0, /* tp_as_sequence */
4344 0, /* tp_as_mapping */
4345 0, /* tp_hash */
4346 0, /* tp_call */
4347 0, /* tp_str */
4348 PyObject_GenericGetAttr, /* tp_getattro */
4349 0, /* tp_setattro */
4350 0, /* tp_as_buffer */
4351 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4352 Py_TPFLAGS_BASETYPE, /* tp_flags */
4353 repeat_doc, /* tp_doc */
4354 (traverseproc)repeat_traverse, /* tp_traverse */
4355 0, /* tp_clear */
4356 0, /* tp_richcompare */
4357 0, /* tp_weaklistoffset */
4358 PyObject_SelfIter, /* tp_iter */
4359 (iternextfunc)repeat_next, /* tp_iternext */
4360 repeat_methods, /* tp_methods */
4361 0, /* tp_members */
4362 0, /* tp_getset */
4363 0, /* tp_base */
4364 0, /* tp_dict */
4365 0, /* tp_descr_get */
4366 0, /* tp_descr_set */
4367 0, /* tp_dictoffset */
4368 0, /* tp_init */
4369 0, /* tp_alloc */
4370 repeat_new, /* tp_new */
4371 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004372};
4373
Tal Einatc4bccd32018-09-12 00:49:13 +03004374
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004375/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004376
4377typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004378 PyObject_HEAD
4379 Py_ssize_t tuplesize;
4380 Py_ssize_t numactive;
4381 PyObject *ittuple; /* tuple of iterators */
4382 PyObject *result;
4383 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004384} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004385
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004386static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004387
4388static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004389zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004390{
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004391 _Py_IDENTIFIER(fillvalue);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004392 ziplongestobject *lz;
4393 Py_ssize_t i;
4394 PyObject *ittuple; /* tuple of iterators */
4395 PyObject *result;
4396 PyObject *fillvalue = Py_None;
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004397 Py_ssize_t tuplesize;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004398
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004399 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004400 fillvalue = NULL;
4401 if (PyDict_GET_SIZE(kwds) == 1) {
4402 fillvalue = _PyDict_GetItemIdWithError(kwds, &PyId_fillvalue);
4403 }
4404 if (fillvalue == NULL) {
4405 if (!PyErr_Occurred()) {
4406 PyErr_SetString(PyExc_TypeError,
4407 "zip_longest() got an unexpected keyword argument");
4408 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004409 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004410 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004411 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004412
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004413 /* args must be a tuple */
4414 assert(PyTuple_Check(args));
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004415 tuplesize = PyTuple_GET_SIZE(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004416
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004417 /* obtain iterators */
4418 ittuple = PyTuple_New(tuplesize);
4419 if (ittuple == NULL)
4420 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004421 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004422 PyObject *item = PyTuple_GET_ITEM(args, i);
4423 PyObject *it = PyObject_GetIter(item);
4424 if (it == NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004425 Py_DECREF(ittuple);
4426 return NULL;
4427 }
4428 PyTuple_SET_ITEM(ittuple, i, it);
4429 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004430
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004431 /* create a result holder */
4432 result = PyTuple_New(tuplesize);
4433 if (result == NULL) {
4434 Py_DECREF(ittuple);
4435 return NULL;
4436 }
4437 for (i=0 ; i < tuplesize ; i++) {
4438 Py_INCREF(Py_None);
4439 PyTuple_SET_ITEM(result, i, Py_None);
4440 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004441
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004442 /* create ziplongestobject structure */
4443 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4444 if (lz == NULL) {
4445 Py_DECREF(ittuple);
4446 Py_DECREF(result);
4447 return NULL;
4448 }
4449 lz->ittuple = ittuple;
4450 lz->tuplesize = tuplesize;
4451 lz->numactive = tuplesize;
4452 lz->result = result;
4453 Py_INCREF(fillvalue);
4454 lz->fillvalue = fillvalue;
4455 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004456}
4457
4458static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004459zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004460{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004461 PyObject_GC_UnTrack(lz);
4462 Py_XDECREF(lz->ittuple);
4463 Py_XDECREF(lz->result);
4464 Py_XDECREF(lz->fillvalue);
4465 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004466}
4467
4468static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004469zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004470{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004471 Py_VISIT(lz->ittuple);
4472 Py_VISIT(lz->result);
4473 Py_VISIT(lz->fillvalue);
4474 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004475}
4476
4477static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004478zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004479{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004480 Py_ssize_t i;
4481 Py_ssize_t tuplesize = lz->tuplesize;
4482 PyObject *result = lz->result;
4483 PyObject *it;
4484 PyObject *item;
4485 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004486
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004487 if (tuplesize == 0)
4488 return NULL;
4489 if (lz->numactive == 0)
4490 return NULL;
4491 if (Py_REFCNT(result) == 1) {
4492 Py_INCREF(result);
4493 for (i=0 ; i < tuplesize ; i++) {
4494 it = PyTuple_GET_ITEM(lz->ittuple, i);
4495 if (it == NULL) {
4496 Py_INCREF(lz->fillvalue);
4497 item = lz->fillvalue;
4498 } else {
4499 item = PyIter_Next(it);
4500 if (item == NULL) {
4501 lz->numactive -= 1;
4502 if (lz->numactive == 0 || PyErr_Occurred()) {
4503 lz->numactive = 0;
4504 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004505 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004506 } else {
4507 Py_INCREF(lz->fillvalue);
4508 item = lz->fillvalue;
4509 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4510 Py_DECREF(it);
4511 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004512 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004513 }
4514 olditem = PyTuple_GET_ITEM(result, i);
4515 PyTuple_SET_ITEM(result, i, item);
4516 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004517 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004518 } else {
4519 result = PyTuple_New(tuplesize);
4520 if (result == NULL)
4521 return NULL;
4522 for (i=0 ; i < tuplesize ; i++) {
4523 it = PyTuple_GET_ITEM(lz->ittuple, i);
4524 if (it == NULL) {
4525 Py_INCREF(lz->fillvalue);
4526 item = lz->fillvalue;
4527 } else {
4528 item = PyIter_Next(it);
4529 if (item == NULL) {
4530 lz->numactive -= 1;
4531 if (lz->numactive == 0 || PyErr_Occurred()) {
4532 lz->numactive = 0;
4533 Py_DECREF(result);
4534 return NULL;
4535 } else {
4536 Py_INCREF(lz->fillvalue);
4537 item = lz->fillvalue;
4538 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4539 Py_DECREF(it);
4540 }
4541 }
4542 }
4543 PyTuple_SET_ITEM(result, i, item);
4544 }
4545 }
4546 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004547}
4548
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004549static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304550zip_longest_reduce(ziplongestobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004551{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004552
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004553 /* Create a new tuple with empty sequences where appropriate to pickle.
4554 * Then use setstate to set the fillvalue
4555 */
4556 int i;
4557 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004558
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004559 if (args == NULL)
4560 return NULL;
4561 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4562 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4563 if (elem == NULL) {
4564 elem = PyTuple_New(0);
4565 if (elem == NULL) {
4566 Py_DECREF(args);
4567 return NULL;
4568 }
4569 } else
4570 Py_INCREF(elem);
4571 PyTuple_SET_ITEM(args, i, elem);
4572 }
4573 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4574}
4575
4576static PyObject *
4577zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4578{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004579 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004580 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004581 Py_RETURN_NONE;
4582}
4583
4584static PyMethodDef zip_longest_methods[] = {
4585 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4586 reduce_doc},
4587 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4588 setstate_doc},
4589 {NULL, NULL} /* sentinel */
4590};
4591
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004592PyDoc_STRVAR(zip_longest_doc,
4593"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004594\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004595Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004596the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004597method continues until the longest iterable in the argument sequence\n\
4598is exhausted and then it raises StopIteration. When the shorter iterables\n\
4599are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4600defaults to None or can be specified by a keyword argument.\n\
4601");
4602
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004603static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004604 PyVarObject_HEAD_INIT(NULL, 0)
4605 "itertools.zip_longest", /* tp_name */
4606 sizeof(ziplongestobject), /* tp_basicsize */
4607 0, /* tp_itemsize */
4608 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004609 (destructor)zip_longest_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004610 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004611 0, /* tp_getattr */
4612 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004613 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004614 0, /* tp_repr */
4615 0, /* tp_as_number */
4616 0, /* tp_as_sequence */
4617 0, /* tp_as_mapping */
4618 0, /* tp_hash */
4619 0, /* tp_call */
4620 0, /* tp_str */
4621 PyObject_GenericGetAttr, /* tp_getattro */
4622 0, /* tp_setattro */
4623 0, /* tp_as_buffer */
4624 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4625 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004626 zip_longest_doc, /* tp_doc */
4627 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004628 0, /* tp_clear */
4629 0, /* tp_richcompare */
4630 0, /* tp_weaklistoffset */
4631 PyObject_SelfIter, /* tp_iter */
4632 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004633 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004634 0, /* tp_members */
4635 0, /* tp_getset */
4636 0, /* tp_base */
4637 0, /* tp_dict */
4638 0, /* tp_descr_get */
4639 0, /* tp_descr_set */
4640 0, /* tp_dictoffset */
4641 0, /* tp_init */
4642 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004643 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004644 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004645};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004646
Tal Einatc4bccd32018-09-12 00:49:13 +03004647
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004648/* module level code ********************************************************/
4649
4650PyDoc_STRVAR(module_doc,
4651"Functional tools for creating and using iterators.\n\
4652\n\
4653Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004654count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004655cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004656repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004657\n\
4658Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004659accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
oldkaa0735f2018-02-02 16:52:55 +08004660chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
4661chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004662compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4663dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4664groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004665filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004666islice(seq, [start,] stop [, step]) --> elements from\n\
4667 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004668starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004669tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004670takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
oldkaa0735f2018-02-02 16:52:55 +08004671zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004672\n\
4673Combinatoric generators:\n\
4674product(p, q, ... [repeat=1]) --> cartesian product\n\
4675permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004676combinations(p, r)\n\
4677combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004678");
4679
Dong-hee Na514c4692020-03-18 02:46:24 +09004680static int
4681itertoolsmodule_exec(PyObject *m)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004682{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004683 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004684 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004685 &combinations_type,
4686 &cwr_type,
4687 &cycle_type,
4688 &dropwhile_type,
4689 &takewhile_type,
4690 &islice_type,
4691 &starmap_type,
4692 &chain_type,
4693 &compress_type,
4694 &filterfalse_type,
4695 &count_type,
4696 &ziplongest_type,
4697 &permutations_type,
4698 &product_type,
4699 &repeat_type,
4700 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004701 &_grouper_type,
4702 &tee_type,
Dong-hee Na05e4a292020-03-23 01:17:34 +09004703 &teedataobject_type
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004704 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004705
Victor Stinnerd2ec81a2020-02-07 09:17:07 +01004706 Py_SET_TYPE(&teedataobject_type, &PyType_Type);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004707
Dong-hee Na05e4a292020-03-23 01:17:34 +09004708 for (size_t i = 0; i < Py_ARRAY_LENGTH(typelist); i++) {
4709 if (PyModule_AddType(m, typelist[i]) < 0) {
Dong-hee Na514c4692020-03-18 02:46:24 +09004710 return -1;
4711 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004712 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004713
Dong-hee Na514c4692020-03-18 02:46:24 +09004714 return 0;
4715}
4716
4717static struct PyModuleDef_Slot itertoolsmodule_slots[] = {
4718 {Py_mod_exec, itertoolsmodule_exec},
4719 {0, NULL}
4720};
4721
4722static PyMethodDef module_methods[] = {
4723 ITERTOOLS_TEE_METHODDEF
4724 {NULL, NULL} /* sentinel */
4725};
4726
4727
4728static struct PyModuleDef itertoolsmodule = {
4729 PyModuleDef_HEAD_INIT,
4730 "itertools",
4731 module_doc,
4732 0,
4733 module_methods,
4734 itertoolsmodule_slots,
4735 NULL,
4736 NULL,
4737 NULL
4738};
4739
4740PyMODINIT_FUNC
4741PyInit_itertools(void)
4742{
4743 return PyModuleDef_Init(&itertoolsmodule);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004744}