blob: e99c964ac1d7903954049061dbcae3f67d4d95f6 [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},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07002031 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00002032};
2033
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002034static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002035 PyVarObject_HEAD_INIT(NULL, 0)
2036 "itertools.chain", /* tp_name */
2037 sizeof(chainobject), /* tp_basicsize */
2038 0, /* tp_itemsize */
2039 /* methods */
2040 (destructor)chain_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002041 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002042 0, /* tp_getattr */
2043 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002044 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002045 0, /* tp_repr */
2046 0, /* tp_as_number */
2047 0, /* tp_as_sequence */
2048 0, /* tp_as_mapping */
2049 0, /* tp_hash */
2050 0, /* tp_call */
2051 0, /* tp_str */
2052 PyObject_GenericGetAttr, /* tp_getattro */
2053 0, /* tp_setattro */
2054 0, /* tp_as_buffer */
2055 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2056 Py_TPFLAGS_BASETYPE, /* tp_flags */
2057 chain_doc, /* tp_doc */
2058 (traverseproc)chain_traverse, /* tp_traverse */
2059 0, /* tp_clear */
2060 0, /* tp_richcompare */
2061 0, /* tp_weaklistoffset */
2062 PyObject_SelfIter, /* tp_iter */
2063 (iternextfunc)chain_next, /* tp_iternext */
2064 chain_methods, /* tp_methods */
2065 0, /* tp_members */
2066 0, /* tp_getset */
2067 0, /* tp_base */
2068 0, /* tp_dict */
2069 0, /* tp_descr_get */
2070 0, /* tp_descr_set */
2071 0, /* tp_dictoffset */
2072 0, /* tp_init */
2073 0, /* tp_alloc */
2074 chain_new, /* tp_new */
2075 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002076};
2077
2078
Christian Heimesc3f30c42008-02-22 16:37:40 +00002079/* product object ************************************************************/
2080
2081typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002082 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002083 PyObject *pools; /* tuple of pool tuples */
2084 Py_ssize_t *indices; /* one index per pool */
2085 PyObject *result; /* most recently returned result tuple */
2086 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002087} productobject;
2088
2089static PyTypeObject product_type;
2090
2091static PyObject *
2092product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2093{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002094 productobject *lz;
2095 Py_ssize_t nargs, npools, repeat=1;
2096 PyObject *pools = NULL;
2097 Py_ssize_t *indices = NULL;
2098 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002099
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002100 if (kwds != NULL) {
2101 char *kwlist[] = {"repeat", 0};
2102 PyObject *tmpargs = PyTuple_New(0);
2103 if (tmpargs == NULL)
2104 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002105 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2106 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002107 Py_DECREF(tmpargs);
2108 return NULL;
2109 }
2110 Py_DECREF(tmpargs);
2111 if (repeat < 0) {
2112 PyErr_SetString(PyExc_ValueError,
2113 "repeat argument cannot be negative");
2114 return NULL;
2115 }
2116 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002117
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002118 assert(PyTuple_CheckExact(args));
2119 if (repeat == 0) {
2120 nargs = 0;
2121 } else {
2122 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002123 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002124 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2125 return NULL;
2126 }
2127 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002128 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002129
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002130 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002131 if (indices == NULL) {
2132 PyErr_NoMemory();
2133 goto error;
2134 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002135
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002136 pools = PyTuple_New(npools);
2137 if (pools == NULL)
2138 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002139
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002140 for (i=0; i < nargs ; ++i) {
2141 PyObject *item = PyTuple_GET_ITEM(args, i);
2142 PyObject *pool = PySequence_Tuple(item);
2143 if (pool == NULL)
2144 goto error;
2145 PyTuple_SET_ITEM(pools, i, pool);
2146 indices[i] = 0;
2147 }
2148 for ( ; i < npools; ++i) {
2149 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2150 Py_INCREF(pool);
2151 PyTuple_SET_ITEM(pools, i, pool);
2152 indices[i] = 0;
2153 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002154
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002155 /* create productobject structure */
2156 lz = (productobject *)type->tp_alloc(type, 0);
2157 if (lz == NULL)
2158 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002159
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002160 lz->pools = pools;
2161 lz->indices = indices;
2162 lz->result = NULL;
2163 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002164
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002165 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002166
2167error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002168 if (indices != NULL)
2169 PyMem_Free(indices);
2170 Py_XDECREF(pools);
2171 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002172}
2173
2174static void
2175product_dealloc(productobject *lz)
2176{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002177 PyObject_GC_UnTrack(lz);
2178 Py_XDECREF(lz->pools);
2179 Py_XDECREF(lz->result);
2180 if (lz->indices != NULL)
2181 PyMem_Free(lz->indices);
2182 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002183}
2184
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002185static PyObject *
2186product_sizeof(productobject *lz, void *unused)
2187{
2188 Py_ssize_t res;
2189
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002190 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002191 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2192 return PyLong_FromSsize_t(res);
2193}
2194
2195PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2196
Christian Heimesc3f30c42008-02-22 16:37:40 +00002197static int
2198product_traverse(productobject *lz, visitproc visit, void *arg)
2199{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002200 Py_VISIT(lz->pools);
2201 Py_VISIT(lz->result);
2202 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002203}
2204
2205static PyObject *
2206product_next(productobject *lz)
2207{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002208 PyObject *pool;
2209 PyObject *elem;
2210 PyObject *oldelem;
2211 PyObject *pools = lz->pools;
2212 PyObject *result = lz->result;
2213 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2214 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002215
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002216 if (lz->stopped)
2217 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002218
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002219 if (result == NULL) {
2220 /* On the first pass, return an initial tuple filled with the
2221 first element from each pool. */
2222 result = PyTuple_New(npools);
2223 if (result == NULL)
2224 goto empty;
2225 lz->result = result;
2226 for (i=0; i < npools; i++) {
2227 pool = PyTuple_GET_ITEM(pools, i);
2228 if (PyTuple_GET_SIZE(pool) == 0)
2229 goto empty;
2230 elem = PyTuple_GET_ITEM(pool, 0);
2231 Py_INCREF(elem);
2232 PyTuple_SET_ITEM(result, i, elem);
2233 }
2234 } else {
2235 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002236
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002237 /* Copy the previous result tuple or re-use it if available */
2238 if (Py_REFCNT(result) > 1) {
2239 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002240 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002241 if (result == NULL)
2242 goto empty;
2243 lz->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002244 Py_DECREF(old_result);
2245 }
2246 /* Now, we've got the only copy so we can update it in-place */
2247 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002248
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002249 /* Update the pool indices right-to-left. Only advance to the
2250 next pool when the previous one rolls-over */
2251 for (i=npools-1 ; i >= 0 ; i--) {
2252 pool = PyTuple_GET_ITEM(pools, i);
2253 indices[i]++;
2254 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2255 /* Roll-over and advance to next pool */
2256 indices[i] = 0;
2257 elem = PyTuple_GET_ITEM(pool, 0);
2258 Py_INCREF(elem);
2259 oldelem = PyTuple_GET_ITEM(result, i);
2260 PyTuple_SET_ITEM(result, i, elem);
2261 Py_DECREF(oldelem);
2262 } else {
2263 /* No rollover. Just increment and stop here. */
2264 elem = PyTuple_GET_ITEM(pool, indices[i]);
2265 Py_INCREF(elem);
2266 oldelem = PyTuple_GET_ITEM(result, i);
2267 PyTuple_SET_ITEM(result, i, elem);
2268 Py_DECREF(oldelem);
2269 break;
2270 }
2271 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002272
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002273 /* If i is negative, then the indices have all rolled-over
2274 and we're done. */
2275 if (i < 0)
2276 goto empty;
2277 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002278
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002279 Py_INCREF(result);
2280 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002281
2282empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002283 lz->stopped = 1;
2284 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002285}
2286
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002287static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302288product_reduce(productobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002289{
2290 if (lz->stopped) {
2291 return Py_BuildValue("O(())", Py_TYPE(lz));
2292 } else if (lz->result == NULL) {
2293 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2294 } else {
2295 PyObject *indices;
2296 Py_ssize_t n, i;
2297
2298 /* we must pickle the indices use them for setstate, and
2299 * additionally indicate that the iterator has started
2300 */
2301 n = PyTuple_GET_SIZE(lz->pools);
2302 indices = PyTuple_New(n);
2303 if (indices == NULL)
2304 return NULL;
2305 for (i=0; i<n; i++){
2306 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2307 if (!index) {
2308 Py_DECREF(indices);
2309 return NULL;
2310 }
2311 PyTuple_SET_ITEM(indices, i, index);
2312 }
2313 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2314 }
2315}
2316
2317static PyObject *
2318product_setstate(productobject *lz, PyObject *state)
2319{
2320 PyObject *result;
2321 Py_ssize_t n, i;
2322
2323 n = PyTuple_GET_SIZE(lz->pools);
2324 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2325 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2326 return NULL;
2327 }
2328 for (i=0; i<n; i++)
2329 {
2330 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2331 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002332 PyObject* pool;
2333 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002334 if (index < 0 && PyErr_Occurred())
2335 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002336 pool = PyTuple_GET_ITEM(lz->pools, i);
2337 poolsize = PyTuple_GET_SIZE(pool);
2338 if (poolsize == 0) {
2339 lz->stopped = 1;
2340 Py_RETURN_NONE;
2341 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002342 /* clamp the index */
2343 if (index < 0)
2344 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002345 else if (index > poolsize-1)
2346 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002347 lz->indices[i] = index;
2348 }
2349
2350 result = PyTuple_New(n);
2351 if (!result)
2352 return NULL;
2353 for (i=0; i<n; i++) {
2354 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2355 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2356 Py_INCREF(element);
2357 PyTuple_SET_ITEM(result, i, element);
2358 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002359 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002360 Py_RETURN_NONE;
2361}
2362
2363static PyMethodDef product_methods[] = {
2364 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2365 reduce_doc},
2366 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2367 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002368 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2369 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002370 {NULL, NULL} /* sentinel */
2371};
2372
Christian Heimesc3f30c42008-02-22 16:37:40 +00002373PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002374"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002375\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002376Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002377For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2378The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2379cycle in a manner similar to an odometer (with the rightmost element changing\n\
2380on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002381To compute the product of an iterable with itself, specify the number\n\
2382of repetitions with the optional repeat keyword argument. For example,\n\
2383product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002384product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2385product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2386
2387static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002388 PyVarObject_HEAD_INIT(NULL, 0)
2389 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002390 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002391 0, /* tp_itemsize */
2392 /* methods */
2393 (destructor)product_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002394 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002395 0, /* tp_getattr */
2396 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002397 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002398 0, /* tp_repr */
2399 0, /* tp_as_number */
2400 0, /* tp_as_sequence */
2401 0, /* tp_as_mapping */
2402 0, /* tp_hash */
2403 0, /* tp_call */
2404 0, /* tp_str */
2405 PyObject_GenericGetAttr, /* tp_getattro */
2406 0, /* tp_setattro */
2407 0, /* tp_as_buffer */
2408 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2409 Py_TPFLAGS_BASETYPE, /* tp_flags */
2410 product_doc, /* tp_doc */
2411 (traverseproc)product_traverse, /* tp_traverse */
2412 0, /* tp_clear */
2413 0, /* tp_richcompare */
2414 0, /* tp_weaklistoffset */
2415 PyObject_SelfIter, /* tp_iter */
2416 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002417 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002418 0, /* tp_members */
2419 0, /* tp_getset */
2420 0, /* tp_base */
2421 0, /* tp_dict */
2422 0, /* tp_descr_get */
2423 0, /* tp_descr_set */
2424 0, /* tp_dictoffset */
2425 0, /* tp_init */
2426 0, /* tp_alloc */
2427 product_new, /* tp_new */
2428 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002429};
2430
2431
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002432/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002433
2434typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002435 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002436 PyObject *pool; /* input converted to a tuple */
2437 Py_ssize_t *indices; /* one index per result element */
2438 PyObject *result; /* most recently returned result tuple */
2439 Py_ssize_t r; /* size of result tuple */
2440 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002441} combinationsobject;
2442
Tal Einatc4bccd32018-09-12 00:49:13 +03002443
2444/*[clinic input]
2445@classmethod
2446itertools.combinations.__new__
2447 iterable: object
2448 r: Py_ssize_t
2449Return successive r-length combinations of elements in the iterable.
2450
2451combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)
2452[clinic start generated code]*/
2453
Christian Heimes380f7f22008-02-28 11:19:05 +00002454static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002455itertools_combinations_impl(PyTypeObject *type, PyObject *iterable,
2456 Py_ssize_t r)
2457/*[clinic end generated code: output=87a689b39c40039c input=06bede09e3da20f8]*/
Christian Heimes380f7f22008-02-28 11:19:05 +00002458{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002459 combinationsobject *co;
2460 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002461 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002462 Py_ssize_t *indices = NULL;
2463 Py_ssize_t i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002464
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002465 pool = PySequence_Tuple(iterable);
2466 if (pool == NULL)
2467 goto error;
2468 n = PyTuple_GET_SIZE(pool);
2469 if (r < 0) {
2470 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2471 goto error;
2472 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002473
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002474 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002475 if (indices == NULL) {
2476 PyErr_NoMemory();
2477 goto error;
2478 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002479
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002480 for (i=0 ; i<r ; i++)
2481 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002482
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002483 /* create combinationsobject structure */
2484 co = (combinationsobject *)type->tp_alloc(type, 0);
2485 if (co == NULL)
2486 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002487
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002488 co->pool = pool;
2489 co->indices = indices;
2490 co->result = NULL;
2491 co->r = r;
2492 co->stopped = r > n ? 1 : 0;
2493
2494 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002495
2496error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002497 if (indices != NULL)
2498 PyMem_Free(indices);
2499 Py_XDECREF(pool);
2500 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002501}
2502
2503static void
2504combinations_dealloc(combinationsobject *co)
2505{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002506 PyObject_GC_UnTrack(co);
2507 Py_XDECREF(co->pool);
2508 Py_XDECREF(co->result);
2509 if (co->indices != NULL)
2510 PyMem_Free(co->indices);
2511 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002512}
2513
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002514static PyObject *
2515combinations_sizeof(combinationsobject *co, void *unused)
2516{
2517 Py_ssize_t res;
2518
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002519 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002520 res += co->r * sizeof(Py_ssize_t);
2521 return PyLong_FromSsize_t(res);
2522}
2523
Christian Heimes380f7f22008-02-28 11:19:05 +00002524static int
2525combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2526{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002527 Py_VISIT(co->pool);
2528 Py_VISIT(co->result);
2529 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002530}
2531
2532static PyObject *
2533combinations_next(combinationsobject *co)
2534{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002535 PyObject *elem;
2536 PyObject *oldelem;
2537 PyObject *pool = co->pool;
2538 Py_ssize_t *indices = co->indices;
2539 PyObject *result = co->result;
2540 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2541 Py_ssize_t r = co->r;
2542 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002543
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002544 if (co->stopped)
2545 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002546
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002547 if (result == NULL) {
2548 /* On the first pass, initialize result tuple using the indices */
2549 result = PyTuple_New(r);
2550 if (result == NULL)
2551 goto empty;
2552 co->result = result;
2553 for (i=0; i<r ; i++) {
2554 index = indices[i];
2555 elem = PyTuple_GET_ITEM(pool, index);
2556 Py_INCREF(elem);
2557 PyTuple_SET_ITEM(result, i, elem);
2558 }
2559 } else {
2560 /* Copy the previous result tuple or re-use it if available */
2561 if (Py_REFCNT(result) > 1) {
2562 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002563 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002564 if (result == NULL)
2565 goto empty;
2566 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002567 Py_DECREF(old_result);
2568 }
2569 /* Now, we've got the only copy so we can update it in-place
2570 * CPython's empty tuple is a singleton and cached in
2571 * PyTuple's freelist.
2572 */
2573 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002574
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002575 /* Scan indices right-to-left until finding one that is not
2576 at its maximum (i + n - r). */
2577 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2578 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002579
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002580 /* If i is negative, then the indices are all at
2581 their maximum value and we're done. */
2582 if (i < 0)
2583 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002584
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002585 /* Increment the current index which we know is not at its
2586 maximum. Then move back to the right setting each index
2587 to its lowest possible value (one higher than the index
2588 to its left -- this maintains the sort order invariant). */
2589 indices[i]++;
2590 for (j=i+1 ; j<r ; j++)
2591 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002592
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002593 /* Update the result tuple for the new indices
2594 starting with i, the leftmost index that changed */
2595 for ( ; i<r ; i++) {
2596 index = indices[i];
2597 elem = PyTuple_GET_ITEM(pool, index);
2598 Py_INCREF(elem);
2599 oldelem = PyTuple_GET_ITEM(result, i);
2600 PyTuple_SET_ITEM(result, i, elem);
2601 Py_DECREF(oldelem);
2602 }
2603 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002604
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002605 Py_INCREF(result);
2606 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002607
2608empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002609 co->stopped = 1;
2610 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002611}
2612
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002613static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302614combinations_reduce(combinationsobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002615{
2616 if (lz->result == NULL) {
2617 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2618 } else if (lz->stopped) {
2619 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2620 } else {
2621 PyObject *indices;
2622 Py_ssize_t i;
2623
2624 /* we must pickle the indices and use them for setstate */
2625 indices = PyTuple_New(lz->r);
2626 if (!indices)
2627 return NULL;
2628 for (i=0; i<lz->r; i++)
2629 {
2630 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2631 if (!index) {
2632 Py_DECREF(indices);
2633 return NULL;
2634 }
2635 PyTuple_SET_ITEM(indices, i, index);
2636 }
2637
2638 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2639 }
2640}
2641
2642static PyObject *
2643combinations_setstate(combinationsobject *lz, PyObject *state)
2644{
2645 PyObject *result;
2646 Py_ssize_t i;
2647 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2648
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002649 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002650 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2651 return NULL;
2652 }
2653
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002654 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002655 Py_ssize_t max;
2656 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2657 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002658
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002659 if (index == -1 && PyErr_Occurred())
2660 return NULL; /* not an integer */
2661 max = i + n - lz->r;
2662 /* clamp the index (beware of negative max) */
2663 if (index > max)
2664 index = max;
2665 if (index < 0)
2666 index = 0;
2667 lz->indices[i] = index;
2668 }
2669
2670 result = PyTuple_New(lz->r);
2671 if (result == NULL)
2672 return NULL;
2673 for (i=0; i<lz->r; i++) {
2674 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2675 Py_INCREF(element);
2676 PyTuple_SET_ITEM(result, i, element);
2677 }
2678
Serhiy Storchaka48842712016-04-06 09:45:48 +03002679 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002680 Py_RETURN_NONE;
2681}
2682
2683static PyMethodDef combinations_methods[] = {
2684 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2685 reduce_doc},
2686 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2687 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002688 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2689 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002690 {NULL, NULL} /* sentinel */
2691};
2692
Christian Heimes380f7f22008-02-28 11:19:05 +00002693static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002694 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002695 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002696 sizeof(combinationsobject), /* tp_basicsize */
2697 0, /* tp_itemsize */
2698 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002699 (destructor)combinations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002700 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002701 0, /* tp_getattr */
2702 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002703 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002704 0, /* tp_repr */
2705 0, /* tp_as_number */
2706 0, /* tp_as_sequence */
2707 0, /* tp_as_mapping */
2708 0, /* tp_hash */
2709 0, /* tp_call */
2710 0, /* tp_str */
2711 PyObject_GenericGetAttr, /* tp_getattro */
2712 0, /* tp_setattro */
2713 0, /* tp_as_buffer */
2714 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2715 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03002716 itertools_combinations__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002717 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002718 0, /* tp_clear */
2719 0, /* tp_richcompare */
2720 0, /* tp_weaklistoffset */
2721 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002722 (iternextfunc)combinations_next, /* tp_iternext */
2723 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002724 0, /* tp_members */
2725 0, /* tp_getset */
2726 0, /* tp_base */
2727 0, /* tp_dict */
2728 0, /* tp_descr_get */
2729 0, /* tp_descr_set */
2730 0, /* tp_dictoffset */
2731 0, /* tp_init */
2732 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03002733 itertools_combinations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002734 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002735};
2736
2737
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002738/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002739
2740/* Equivalent to:
2741
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002742 def combinations_with_replacement(iterable, r):
2743 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2744 # number items returned: (n+r-1)! / r! / (n-1)!
2745 pool = tuple(iterable)
2746 n = len(pool)
2747 indices = [0] * r
2748 yield tuple(pool[i] for i in indices)
2749 while 1:
2750 for i in reversed(range(r)):
2751 if indices[i] != n - 1:
2752 break
2753 else:
2754 return
2755 indices[i:] = [indices[i] + 1] * (r - i)
2756 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002757
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002758 def combinations_with_replacement2(iterable, r):
2759 'Alternate version that filters from product()'
2760 pool = tuple(iterable)
2761 n = len(pool)
2762 for indices in product(range(n), repeat=r):
2763 if sorted(indices) == list(indices):
2764 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002765*/
2766typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002767 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002768 PyObject *pool; /* input converted to a tuple */
2769 Py_ssize_t *indices; /* one index per result element */
2770 PyObject *result; /* most recently returned result tuple */
2771 Py_ssize_t r; /* size of result tuple */
2772 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002773} cwrobject;
2774
Tal Einatc4bccd32018-09-12 00:49:13 +03002775/*[clinic input]
2776@classmethod
2777itertools.combinations_with_replacement.__new__
2778 iterable: object
2779 r: Py_ssize_t
2780Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.
2781
2782combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2783[clinic start generated code]*/
2784
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002785static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002786itertools_combinations_with_replacement_impl(PyTypeObject *type,
2787 PyObject *iterable,
2788 Py_ssize_t r)
2789/*[clinic end generated code: output=48b26856d4e659ca input=dc2a8c7ba785fad7]*/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002790{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002791 cwrobject *co;
2792 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002793 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002794 Py_ssize_t *indices = NULL;
2795 Py_ssize_t i;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002796
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002797 pool = PySequence_Tuple(iterable);
2798 if (pool == NULL)
2799 goto error;
2800 n = PyTuple_GET_SIZE(pool);
2801 if (r < 0) {
2802 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2803 goto error;
2804 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002805
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002806 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002807 if (indices == NULL) {
2808 PyErr_NoMemory();
2809 goto error;
2810 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002811
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002812 for (i=0 ; i<r ; i++)
2813 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002814
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002815 /* create cwrobject structure */
2816 co = (cwrobject *)type->tp_alloc(type, 0);
2817 if (co == NULL)
2818 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002819
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002820 co->pool = pool;
2821 co->indices = indices;
2822 co->result = NULL;
2823 co->r = r;
2824 co->stopped = !n && r;
2825
2826 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002827
2828error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002829 if (indices != NULL)
2830 PyMem_Free(indices);
2831 Py_XDECREF(pool);
2832 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002833}
2834
2835static void
2836cwr_dealloc(cwrobject *co)
2837{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002838 PyObject_GC_UnTrack(co);
2839 Py_XDECREF(co->pool);
2840 Py_XDECREF(co->result);
2841 if (co->indices != NULL)
2842 PyMem_Free(co->indices);
2843 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002844}
2845
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002846static PyObject *
2847cwr_sizeof(cwrobject *co, void *unused)
2848{
2849 Py_ssize_t res;
2850
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002851 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002852 res += co->r * sizeof(Py_ssize_t);
2853 return PyLong_FromSsize_t(res);
2854}
2855
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002856static int
2857cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2858{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002859 Py_VISIT(co->pool);
2860 Py_VISIT(co->result);
2861 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002862}
2863
2864static PyObject *
2865cwr_next(cwrobject *co)
2866{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002867 PyObject *elem;
2868 PyObject *oldelem;
2869 PyObject *pool = co->pool;
2870 Py_ssize_t *indices = co->indices;
2871 PyObject *result = co->result;
2872 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2873 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002874 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002875
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002876 if (co->stopped)
2877 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002878
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002879 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002880 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002881 result = PyTuple_New(r);
2882 if (result == NULL)
2883 goto empty;
2884 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002885 if (n > 0) {
2886 elem = PyTuple_GET_ITEM(pool, 0);
2887 for (i=0; i<r ; i++) {
2888 assert(indices[i] == 0);
2889 Py_INCREF(elem);
2890 PyTuple_SET_ITEM(result, i, elem);
2891 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002892 }
2893 } else {
2894 /* Copy the previous result tuple or re-use it if available */
2895 if (Py_REFCNT(result) > 1) {
2896 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002897 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002898 if (result == NULL)
2899 goto empty;
2900 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002901 Py_DECREF(old_result);
2902 }
2903 /* Now, we've got the only copy so we can update it in-place CPython's
2904 empty tuple is a singleton and cached in PyTuple's freelist. */
2905 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002906
Tim Peters9edb1682013-09-03 11:49:31 -05002907 /* Scan indices right-to-left until finding one that is not
2908 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002909 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2910 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002911
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002912 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002913 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002914 if (i < 0)
2915 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002916
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002917 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002918 maximum. Then set all to the right to the same value. */
2919 index = indices[i] + 1;
2920 assert(index < n);
2921 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002922 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002923 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002924 Py_INCREF(elem);
2925 oldelem = PyTuple_GET_ITEM(result, i);
2926 PyTuple_SET_ITEM(result, i, elem);
2927 Py_DECREF(oldelem);
2928 }
2929 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002930
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002931 Py_INCREF(result);
2932 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002933
2934empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002935 co->stopped = 1;
2936 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002937}
2938
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002939static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302940cwr_reduce(cwrobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002941{
2942 if (lz->result == NULL) {
2943 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2944 } else if (lz->stopped) {
2945 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2946 } else {
2947 PyObject *indices;
2948 Py_ssize_t i;
2949
2950 /* we must pickle the indices and use them for setstate */
2951 indices = PyTuple_New(lz->r);
2952 if (!indices)
2953 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002954 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002955 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2956 if (!index) {
2957 Py_DECREF(indices);
2958 return NULL;
2959 }
2960 PyTuple_SET_ITEM(indices, i, index);
2961 }
2962
2963 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2964 }
2965}
2966
2967static PyObject *
2968cwr_setstate(cwrobject *lz, PyObject *state)
2969{
2970 PyObject *result;
2971 Py_ssize_t n, i;
2972
2973 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2974 {
2975 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2976 return NULL;
2977 }
2978
2979 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002980 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002981 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2982 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002983
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002984 if (index < 0 && PyErr_Occurred())
2985 return NULL; /* not an integer */
2986 /* clamp the index */
2987 if (index < 0)
2988 index = 0;
2989 else if (index > n-1)
2990 index = n-1;
2991 lz->indices[i] = index;
2992 }
2993 result = PyTuple_New(lz->r);
2994 if (result == NULL)
2995 return NULL;
2996 for (i=0; i<lz->r; i++) {
2997 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2998 Py_INCREF(element);
2999 PyTuple_SET_ITEM(result, i, element);
3000 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003001 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003002 Py_RETURN_NONE;
3003}
3004
3005static PyMethodDef cwr_methods[] = {
3006 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
3007 reduce_doc},
3008 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
3009 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003010 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
3011 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003012 {NULL, NULL} /* sentinel */
3013};
3014
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003015static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003016 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003017 "itertools.combinations_with_replacement", /* tp_name */
3018 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003019 0, /* tp_itemsize */
3020 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003021 (destructor)cwr_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003022 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003023 0, /* tp_getattr */
3024 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003025 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003026 0, /* tp_repr */
3027 0, /* tp_as_number */
3028 0, /* tp_as_sequence */
3029 0, /* tp_as_mapping */
3030 0, /* tp_hash */
3031 0, /* tp_call */
3032 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003033 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003034 0, /* tp_setattro */
3035 0, /* tp_as_buffer */
3036 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003037 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003038 itertools_combinations_with_replacement__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003039 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003040 0, /* tp_clear */
3041 0, /* tp_richcompare */
3042 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003043 PyObject_SelfIter, /* tp_iter */
3044 (iternextfunc)cwr_next, /* tp_iternext */
3045 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003046 0, /* tp_members */
3047 0, /* tp_getset */
3048 0, /* tp_base */
3049 0, /* tp_dict */
3050 0, /* tp_descr_get */
3051 0, /* tp_descr_set */
3052 0, /* tp_dictoffset */
3053 0, /* tp_init */
3054 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003055 itertools_combinations_with_replacement, /* tp_new */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003056 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003057};
3058
3059
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003060/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003061
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003062def permutations(iterable, r=None):
Sergey0078a0c2019-10-29 08:10:24 +03003063 # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
3064 # permutations(range(3)) --> 012 021 102 120 201 210
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003065 pool = tuple(iterable)
3066 n = len(pool)
3067 r = n if r is None else r
Sergey0078a0c2019-10-29 08:10:24 +03003068 if r > n:
3069 return
3070 indices = list(range(n))
3071 cycles = list(range(n, n-r, -1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003072 yield tuple(pool[i] for i in indices[:r])
3073 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003074 for i in reversed(range(r)):
3075 cycles[i] -= 1
3076 if cycles[i] == 0:
3077 indices[i:] = indices[i+1:] + indices[i:i+1]
3078 cycles[i] = n - i
3079 else:
3080 j = cycles[i]
3081 indices[i], indices[-j] = indices[-j], indices[i]
3082 yield tuple(pool[i] for i in indices[:r])
3083 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003084 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003085 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003086*/
3087
3088typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003089 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003090 PyObject *pool; /* input converted to a tuple */
3091 Py_ssize_t *indices; /* one index per element in the pool */
3092 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3093 PyObject *result; /* most recently returned result tuple */
3094 Py_ssize_t r; /* size of result tuple */
3095 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003096} permutationsobject;
3097
Tal Einatc4bccd32018-09-12 00:49:13 +03003098/*[clinic input]
3099@classmethod
3100itertools.permutations.__new__
3101 iterable: object
3102 r as robj: object = None
3103Return successive r-length permutations of elements in the iterable.
3104
3105permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
3106[clinic start generated code]*/
3107
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003108static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003109itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
3110 PyObject *robj)
3111/*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003112{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003113 permutationsobject *po;
3114 Py_ssize_t n;
3115 Py_ssize_t r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003116 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003117 Py_ssize_t *indices = NULL;
3118 Py_ssize_t *cycles = NULL;
3119 Py_ssize_t i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003120
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003121 pool = PySequence_Tuple(iterable);
3122 if (pool == NULL)
3123 goto error;
3124 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003125
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003126 r = n;
3127 if (robj != Py_None) {
3128 if (!PyLong_Check(robj)) {
3129 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3130 goto error;
3131 }
3132 r = PyLong_AsSsize_t(robj);
3133 if (r == -1 && PyErr_Occurred())
3134 goto error;
3135 }
3136 if (r < 0) {
3137 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3138 goto error;
3139 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003140
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003141 indices = PyMem_New(Py_ssize_t, n);
3142 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003143 if (indices == NULL || cycles == NULL) {
3144 PyErr_NoMemory();
3145 goto error;
3146 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003147
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003148 for (i=0 ; i<n ; i++)
3149 indices[i] = i;
3150 for (i=0 ; i<r ; i++)
3151 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003152
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003153 /* create permutationsobject structure */
3154 po = (permutationsobject *)type->tp_alloc(type, 0);
3155 if (po == NULL)
3156 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003157
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003158 po->pool = pool;
3159 po->indices = indices;
3160 po->cycles = cycles;
3161 po->result = NULL;
3162 po->r = r;
3163 po->stopped = r > n ? 1 : 0;
3164
3165 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003166
3167error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003168 if (indices != NULL)
3169 PyMem_Free(indices);
3170 if (cycles != NULL)
3171 PyMem_Free(cycles);
3172 Py_XDECREF(pool);
3173 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003174}
3175
3176static void
3177permutations_dealloc(permutationsobject *po)
3178{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003179 PyObject_GC_UnTrack(po);
3180 Py_XDECREF(po->pool);
3181 Py_XDECREF(po->result);
3182 PyMem_Free(po->indices);
3183 PyMem_Free(po->cycles);
3184 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003185}
3186
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003187static PyObject *
3188permutations_sizeof(permutationsobject *po, void *unused)
3189{
3190 Py_ssize_t res;
3191
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003192 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003193 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3194 res += po->r * sizeof(Py_ssize_t);
3195 return PyLong_FromSsize_t(res);
3196}
3197
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003198static int
3199permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3200{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003201 Py_VISIT(po->pool);
3202 Py_VISIT(po->result);
3203 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003204}
3205
3206static PyObject *
3207permutations_next(permutationsobject *po)
3208{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003209 PyObject *elem;
3210 PyObject *oldelem;
3211 PyObject *pool = po->pool;
3212 Py_ssize_t *indices = po->indices;
3213 Py_ssize_t *cycles = po->cycles;
3214 PyObject *result = po->result;
3215 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3216 Py_ssize_t r = po->r;
3217 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003218
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003219 if (po->stopped)
3220 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003221
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003222 if (result == NULL) {
3223 /* On the first pass, initialize result tuple using the indices */
3224 result = PyTuple_New(r);
3225 if (result == NULL)
3226 goto empty;
3227 po->result = result;
3228 for (i=0; i<r ; i++) {
3229 index = indices[i];
3230 elem = PyTuple_GET_ITEM(pool, index);
3231 Py_INCREF(elem);
3232 PyTuple_SET_ITEM(result, i, elem);
3233 }
3234 } else {
3235 if (n == 0)
3236 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003237
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003238 /* Copy the previous result tuple or re-use it if available */
3239 if (Py_REFCNT(result) > 1) {
3240 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05003241 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003242 if (result == NULL)
3243 goto empty;
3244 po->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003245 Py_DECREF(old_result);
3246 }
3247 /* Now, we've got the only copy so we can update it in-place */
3248 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003249
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003250 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3251 for (i=r-1 ; i>=0 ; i--) {
3252 cycles[i] -= 1;
3253 if (cycles[i] == 0) {
3254 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3255 index = indices[i];
3256 for (j=i ; j<n-1 ; j++)
3257 indices[j] = indices[j+1];
3258 indices[n-1] = index;
3259 cycles[i] = n - i;
3260 } else {
3261 j = cycles[i];
3262 index = indices[i];
3263 indices[i] = indices[n-j];
3264 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003265
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003266 for (k=i; k<r ; k++) {
3267 /* start with i, the leftmost element that changed */
3268 /* yield tuple(pool[k] for k in indices[:r]) */
3269 index = indices[k];
3270 elem = PyTuple_GET_ITEM(pool, index);
3271 Py_INCREF(elem);
3272 oldelem = PyTuple_GET_ITEM(result, k);
3273 PyTuple_SET_ITEM(result, k, elem);
3274 Py_DECREF(oldelem);
3275 }
3276 break;
3277 }
3278 }
3279 /* If i is negative, then the cycles have all
3280 rolled-over and we're done. */
3281 if (i < 0)
3282 goto empty;
3283 }
3284 Py_INCREF(result);
3285 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003286
3287empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003288 po->stopped = 1;
3289 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003290}
3291
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003292static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303293permutations_reduce(permutationsobject *po, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003294{
3295 if (po->result == NULL) {
3296 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3297 } else if (po->stopped) {
3298 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3299 } else {
3300 PyObject *indices=NULL, *cycles=NULL;
3301 Py_ssize_t n, i;
3302
3303 /* we must pickle the indices and cycles and use them for setstate */
3304 n = PyTuple_GET_SIZE(po->pool);
3305 indices = PyTuple_New(n);
3306 if (indices == NULL)
3307 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003308 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003309 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3310 if (!index)
3311 goto err;
3312 PyTuple_SET_ITEM(indices, i, index);
3313 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003314
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003315 cycles = PyTuple_New(po->r);
3316 if (cycles == NULL)
3317 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003318 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003319 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3320 if (!index)
3321 goto err;
3322 PyTuple_SET_ITEM(cycles, i, index);
3323 }
3324 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3325 po->pool, po->r,
3326 indices, cycles);
3327 err:
3328 Py_XDECREF(indices);
3329 Py_XDECREF(cycles);
3330 return NULL;
3331 }
3332}
3333
3334static PyObject *
3335permutations_setstate(permutationsobject *po, PyObject *state)
3336{
3337 PyObject *indices, *cycles, *result;
3338 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003339
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003340 if (!PyTuple_Check(state)) {
3341 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3342 return NULL;
3343 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003344 if (!PyArg_ParseTuple(state, "O!O!",
3345 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003346 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003347 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003348 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003349
3350 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003351 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003352 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3353 return NULL;
3354 }
3355
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003356 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003357 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3358 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3359 if (index < 0 && PyErr_Occurred())
3360 return NULL; /* not an integer */
3361 /* clamp the index */
3362 if (index < 0)
3363 index = 0;
3364 else if (index > n-1)
3365 index = n-1;
3366 po->indices[i] = index;
3367 }
3368
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003369 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003370 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3371 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3372 if (index < 0 && PyErr_Occurred())
3373 return NULL; /* not an integer */
3374 if (index < 1)
3375 index = 1;
3376 else if (index > n-i)
3377 index = n-i;
3378 po->cycles[i] = index;
3379 }
3380 result = PyTuple_New(po->r);
3381 if (result == NULL)
3382 return NULL;
3383 for (i=0; i<po->r; i++) {
3384 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3385 Py_INCREF(element);
3386 PyTuple_SET_ITEM(result, i, element);
3387 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003388 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003389 Py_RETURN_NONE;
3390}
3391
3392static PyMethodDef permuations_methods[] = {
3393 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3394 reduce_doc},
3395 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3396 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003397 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3398 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003399 {NULL, NULL} /* sentinel */
3400};
3401
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003402static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003403 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003404 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003405 sizeof(permutationsobject), /* tp_basicsize */
3406 0, /* tp_itemsize */
3407 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003408 (destructor)permutations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003409 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003410 0, /* tp_getattr */
3411 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003412 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003413 0, /* tp_repr */
3414 0, /* tp_as_number */
3415 0, /* tp_as_sequence */
3416 0, /* tp_as_mapping */
3417 0, /* tp_hash */
3418 0, /* tp_call */
3419 0, /* tp_str */
3420 PyObject_GenericGetAttr, /* tp_getattro */
3421 0, /* tp_setattro */
3422 0, /* tp_as_buffer */
3423 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3424 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003425 itertools_permutations__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003426 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003427 0, /* tp_clear */
3428 0, /* tp_richcompare */
3429 0, /* tp_weaklistoffset */
3430 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003431 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003432 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003433 0, /* tp_members */
3434 0, /* tp_getset */
3435 0, /* tp_base */
3436 0, /* tp_dict */
3437 0, /* tp_descr_get */
3438 0, /* tp_descr_set */
3439 0, /* tp_dictoffset */
3440 0, /* tp_init */
3441 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003442 itertools_permutations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003443 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003444};
3445
Tal Einatc4bccd32018-09-12 00:49:13 +03003446
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003447/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003448
3449typedef struct {
3450 PyObject_HEAD
3451 PyObject *total;
3452 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003453 PyObject *binop;
Lisa Roach9718b592018-09-23 17:34:59 -07003454 PyObject *initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003455} accumulateobject;
3456
Tal Einatc4bccd32018-09-12 00:49:13 +03003457/*[clinic input]
3458@classmethod
3459itertools.accumulate.__new__
3460 iterable: object
3461 func as binop: object = None
Lisa Roach9718b592018-09-23 17:34:59 -07003462 *
3463 initial: object = None
Tal Einatc4bccd32018-09-12 00:49:13 +03003464Return series of accumulated sums (or other binary function results).
3465[clinic start generated code]*/
3466
Raymond Hettinger482ba772010-12-01 22:48:00 +00003467static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003468itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
Lisa Roach9718b592018-09-23 17:34:59 -07003469 PyObject *binop, PyObject *initial)
3470/*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003471{
Raymond Hettinger482ba772010-12-01 22:48:00 +00003472 PyObject *it;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003473 accumulateobject *lz;
3474
Raymond Hettinger482ba772010-12-01 22:48:00 +00003475 /* Get iterator. */
3476 it = PyObject_GetIter(iterable);
3477 if (it == NULL)
3478 return NULL;
3479
Raymond Hettinger482ba772010-12-01 22:48:00 +00003480 /* create accumulateobject structure */
3481 lz = (accumulateobject *)type->tp_alloc(type, 0);
3482 if (lz == NULL) {
3483 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003484 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003485 }
3486
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003487 if (binop != Py_None) {
3488 Py_XINCREF(binop);
3489 lz->binop = binop;
3490 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003491 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003492 lz->it = it;
Lisa Roach9718b592018-09-23 17:34:59 -07003493 Py_XINCREF(initial);
3494 lz->initial = initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003495 return (PyObject *)lz;
3496}
3497
3498static void
3499accumulate_dealloc(accumulateobject *lz)
3500{
3501 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003502 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003503 Py_XDECREF(lz->total);
3504 Py_XDECREF(lz->it);
Lisa Roach9718b592018-09-23 17:34:59 -07003505 Py_XDECREF(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003506 Py_TYPE(lz)->tp_free(lz);
3507}
3508
3509static int
3510accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3511{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003512 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003513 Py_VISIT(lz->it);
3514 Py_VISIT(lz->total);
Lisa Roach9718b592018-09-23 17:34:59 -07003515 Py_VISIT(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003516 return 0;
3517}
3518
3519static PyObject *
3520accumulate_next(accumulateobject *lz)
3521{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003522 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003523
Lisa Roach9718b592018-09-23 17:34:59 -07003524 if (lz->initial != Py_None) {
3525 lz->total = lz->initial;
3526 Py_INCREF(Py_None);
3527 lz->initial = Py_None;
3528 Py_INCREF(lz->total);
3529 return lz->total;
3530 }
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003531 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003532 if (val == NULL)
3533 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003534
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003535 if (lz->total == NULL) {
3536 Py_INCREF(val);
3537 lz->total = val;
3538 return lz->total;
3539 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003540
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003541 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003542 newtotal = PyNumber_Add(lz->total, val);
3543 else
3544 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003545 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003546 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003547 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003548
Raymond Hettinger482ba772010-12-01 22:48:00 +00003549 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003550 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003551 return newtotal;
3552}
3553
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003554static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303555accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003556{
Lisa Roach9718b592018-09-23 17:34:59 -07003557 if (lz->initial != Py_None) {
3558 PyObject *it;
3559
3560 assert(lz->total == NULL);
3561 if (PyType_Ready(&chain_type) < 0)
3562 return NULL;
3563 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3564 lz->initial, lz->it);
3565 if (it == NULL)
3566 return NULL;
3567 return Py_BuildValue("O(NO)O", Py_TYPE(lz),
3568 it, lz->binop?lz->binop:Py_None, Py_None);
3569 }
Serhiy Storchakad5516252016-03-06 14:00:45 +02003570 if (lz->total == Py_None) {
3571 PyObject *it;
3572
3573 if (PyType_Ready(&chain_type) < 0)
3574 return NULL;
3575 if (PyType_Ready(&islice_type) < 0)
3576 return NULL;
3577 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3578 lz->total, lz->it);
3579 if (it == NULL)
3580 return NULL;
3581 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3582 it, lz->binop ? lz->binop : Py_None);
3583 if (it == NULL)
3584 return NULL;
3585 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3586 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003587 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3588 lz->it, lz->binop?lz->binop:Py_None,
3589 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003590}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003591
3592static PyObject *
3593accumulate_setstate(accumulateobject *lz, PyObject *state)
3594{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003595 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003596 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003597 Py_RETURN_NONE;
3598}
3599
3600static PyMethodDef accumulate_methods[] = {
3601 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3602 reduce_doc},
3603 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3604 setstate_doc},
3605 {NULL, NULL} /* sentinel */
3606};
3607
Raymond Hettinger482ba772010-12-01 22:48:00 +00003608static PyTypeObject accumulate_type = {
3609 PyVarObject_HEAD_INIT(NULL, 0)
3610 "itertools.accumulate", /* tp_name */
3611 sizeof(accumulateobject), /* tp_basicsize */
3612 0, /* tp_itemsize */
3613 /* methods */
3614 (destructor)accumulate_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003615 0, /* tp_vectorcall_offset */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003616 0, /* tp_getattr */
3617 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003618 0, /* tp_as_async */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003619 0, /* tp_repr */
3620 0, /* tp_as_number */
3621 0, /* tp_as_sequence */
3622 0, /* tp_as_mapping */
3623 0, /* tp_hash */
3624 0, /* tp_call */
3625 0, /* tp_str */
3626 PyObject_GenericGetAttr, /* tp_getattro */
3627 0, /* tp_setattro */
3628 0, /* tp_as_buffer */
3629 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3630 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003631 itertools_accumulate__doc__, /* tp_doc */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003632 (traverseproc)accumulate_traverse, /* tp_traverse */
3633 0, /* tp_clear */
3634 0, /* tp_richcompare */
3635 0, /* tp_weaklistoffset */
3636 PyObject_SelfIter, /* tp_iter */
3637 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003638 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003639 0, /* tp_members */
3640 0, /* tp_getset */
3641 0, /* tp_base */
3642 0, /* tp_dict */
3643 0, /* tp_descr_get */
3644 0, /* tp_descr_set */
3645 0, /* tp_dictoffset */
3646 0, /* tp_init */
3647 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003648 itertools_accumulate, /* tp_new */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003649 PyObject_GC_Del, /* tp_free */
3650};
3651
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003652
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003653/* compress object ************************************************************/
3654
3655/* Equivalent to:
3656
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003657 def compress(data, selectors):
3658 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3659 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003660*/
3661
3662typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003663 PyObject_HEAD
3664 PyObject *data;
3665 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003666} compressobject;
3667
Tal Einatc4bccd32018-09-12 00:49:13 +03003668/*[clinic input]
3669@classmethod
3670itertools.compress.__new__
3671 data as seq1: object
3672 selectors as seq2: object
3673Return data elements corresponding to true selector elements.
3674
3675Forms a shorter iterator from selected data elements using the selectors to
3676choose the data elements.
3677[clinic start generated code]*/
3678
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003679static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003680itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
3681/*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003682{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003683 PyObject *data=NULL, *selectors=NULL;
3684 compressobject *lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003685
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003686 data = PyObject_GetIter(seq1);
3687 if (data == NULL)
3688 goto fail;
3689 selectors = PyObject_GetIter(seq2);
3690 if (selectors == NULL)
3691 goto fail;
3692
3693 /* create compressobject structure */
3694 lz = (compressobject *)type->tp_alloc(type, 0);
3695 if (lz == NULL)
3696 goto fail;
3697 lz->data = data;
3698 lz->selectors = selectors;
3699 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003700
3701fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003702 Py_XDECREF(data);
3703 Py_XDECREF(selectors);
3704 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003705}
3706
3707static void
3708compress_dealloc(compressobject *lz)
3709{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003710 PyObject_GC_UnTrack(lz);
3711 Py_XDECREF(lz->data);
3712 Py_XDECREF(lz->selectors);
3713 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003714}
3715
3716static int
3717compress_traverse(compressobject *lz, visitproc visit, void *arg)
3718{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003719 Py_VISIT(lz->data);
3720 Py_VISIT(lz->selectors);
3721 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003722}
3723
3724static PyObject *
3725compress_next(compressobject *lz)
3726{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003727 PyObject *data = lz->data, *selectors = lz->selectors;
3728 PyObject *datum, *selector;
3729 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3730 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3731 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003732
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003733 while (1) {
3734 /* Steps: get datum, get selector, evaluate selector.
3735 Order is important (to match the pure python version
3736 in terms of which input gets a chance to raise an
3737 exception first).
3738 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003739
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003740 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003741 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003742 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003743
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003744 selector = selectornext(selectors);
3745 if (selector == NULL) {
3746 Py_DECREF(datum);
3747 return NULL;
3748 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003749
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003750 ok = PyObject_IsTrue(selector);
3751 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003752 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003753 return datum;
3754 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003755 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003756 return NULL;
3757 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003758}
3759
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003760static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303761compress_reduce(compressobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003762{
3763 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3764 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003765}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003766
3767static PyMethodDef compress_methods[] = {
3768 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3769 reduce_doc},
3770 {NULL, NULL} /* sentinel */
3771};
3772
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003773static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003774 PyVarObject_HEAD_INIT(NULL, 0)
3775 "itertools.compress", /* tp_name */
3776 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003777 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003778 /* methods */
3779 (destructor)compress_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003780 0, /* tp_vectorcall_offset */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003781 0, /* tp_getattr */
3782 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003783 0, /* tp_as_async */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003784 0, /* tp_repr */
3785 0, /* tp_as_number */
3786 0, /* tp_as_sequence */
3787 0, /* tp_as_mapping */
3788 0, /* tp_hash */
3789 0, /* tp_call */
3790 0, /* tp_str */
3791 PyObject_GenericGetAttr, /* tp_getattro */
3792 0, /* tp_setattro */
3793 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003794 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003795 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003796 itertools_compress__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003797 (traverseproc)compress_traverse, /* tp_traverse */
3798 0, /* tp_clear */
3799 0, /* tp_richcompare */
3800 0, /* tp_weaklistoffset */
3801 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003802 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003803 compress_methods, /* tp_methods */
3804 0, /* tp_members */
3805 0, /* tp_getset */
3806 0, /* tp_base */
3807 0, /* tp_dict */
3808 0, /* tp_descr_get */
3809 0, /* tp_descr_set */
3810 0, /* tp_dictoffset */
3811 0, /* tp_init */
3812 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003813 itertools_compress, /* tp_new */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003814 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003815};
3816
3817
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003818/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003819
3820typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003821 PyObject_HEAD
3822 PyObject *func;
3823 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003824} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003825
Tal Einatc4bccd32018-09-12 00:49:13 +03003826/*[clinic input]
3827@classmethod
3828itertools.filterfalse.__new__
3829 function as func: object
3830 iterable as seq: object
3831 /
3832Return those items of iterable for which function(item) is false.
3833
3834If function is None, return the items that are false.
3835[clinic start generated code]*/
3836
Raymond Hettinger60eca932003-02-09 06:40:58 +00003837static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003838itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
3839/*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003840{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003841 PyObject *it;
3842 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003843
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003844 /* Get iterator. */
3845 it = PyObject_GetIter(seq);
3846 if (it == NULL)
3847 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003848
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003849 /* create filterfalseobject structure */
3850 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3851 if (lz == NULL) {
3852 Py_DECREF(it);
3853 return NULL;
3854 }
3855 Py_INCREF(func);
3856 lz->func = func;
3857 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003858
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003859 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003860}
3861
3862static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003863filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003864{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003865 PyObject_GC_UnTrack(lz);
3866 Py_XDECREF(lz->func);
3867 Py_XDECREF(lz->it);
3868 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003869}
3870
3871static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003872filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003873{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003874 Py_VISIT(lz->it);
3875 Py_VISIT(lz->func);
3876 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003877}
3878
3879static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003880filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003881{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003882 PyObject *item;
3883 PyObject *it = lz->it;
3884 long ok;
3885 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003886
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003887 iternext = *Py_TYPE(it)->tp_iternext;
3888 for (;;) {
3889 item = iternext(it);
3890 if (item == NULL)
3891 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003892
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003893 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3894 ok = PyObject_IsTrue(item);
3895 } else {
3896 PyObject *good;
Petr Viktorinffd97532020-02-11 17:46:57 +01003897 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003898 if (good == NULL) {
3899 Py_DECREF(item);
3900 return NULL;
3901 }
3902 ok = PyObject_IsTrue(good);
3903 Py_DECREF(good);
3904 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003905 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003906 return item;
3907 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003908 if (ok < 0)
3909 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003910 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003911}
3912
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003913static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303914filterfalse_reduce(filterfalseobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003915{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003916 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3917}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003918
3919static PyMethodDef filterfalse_methods[] = {
3920 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3921 reduce_doc},
3922 {NULL, NULL} /* sentinel */
3923};
3924
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003925static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003926 PyVarObject_HEAD_INIT(NULL, 0)
3927 "itertools.filterfalse", /* tp_name */
3928 sizeof(filterfalseobject), /* tp_basicsize */
3929 0, /* tp_itemsize */
3930 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003931 (destructor)filterfalse_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003932 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003933 0, /* tp_getattr */
3934 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003935 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003936 0, /* tp_repr */
3937 0, /* tp_as_number */
3938 0, /* tp_as_sequence */
3939 0, /* tp_as_mapping */
3940 0, /* tp_hash */
3941 0, /* tp_call */
3942 0, /* tp_str */
3943 PyObject_GenericGetAttr, /* tp_getattro */
3944 0, /* tp_setattro */
3945 0, /* tp_as_buffer */
3946 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3947 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003948 itertools_filterfalse__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003949 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003950 0, /* tp_clear */
3951 0, /* tp_richcompare */
3952 0, /* tp_weaklistoffset */
3953 PyObject_SelfIter, /* tp_iter */
3954 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003955 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003956 0, /* tp_members */
3957 0, /* tp_getset */
3958 0, /* tp_base */
3959 0, /* tp_dict */
3960 0, /* tp_descr_get */
3961 0, /* tp_descr_set */
3962 0, /* tp_dictoffset */
3963 0, /* tp_init */
3964 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003965 itertools_filterfalse, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003966 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003967};
3968
3969
3970/* count object ************************************************************/
3971
3972typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003973 PyObject_HEAD
3974 Py_ssize_t cnt;
3975 PyObject *long_cnt;
3976 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003977} countobject;
3978
Raymond Hettinger30729212009-02-12 06:28:27 +00003979/* Counting logic and invariants:
3980
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003981fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003982
Serhiy Storchaka483405b2015-02-17 10:14:30 +02003983 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003984 Advances with: cnt += 1
3985 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003986
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003987slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003988
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003989 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3990 All counting is done with python objects (no overflows or underflows).
3991 Advances with: long_cnt += long_step
3992 Step may be zero -- effectively a slow version of repeat(cnt).
3993 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003994*/
3995
Tal Einatc4bccd32018-09-12 00:49:13 +03003996/*[clinic input]
3997@classmethod
3998itertools.count.__new__
3999 start as long_cnt: object(c_default="NULL") = 0
4000 step as long_step: object(c_default="NULL") = 1
4001Return a count object whose .__next__() method returns consecutive values.
4002
4003Equivalent to:
4004 def count(firstval=0, step=1):
4005 x = firstval
4006 while 1:
4007 yield x
4008 x += step
4009[clinic start generated code]*/
4010
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004011static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03004012itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
4013 PyObject *long_step)
4014/*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004015{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004016 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004017 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004018 Py_ssize_t cnt = 0;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00004019 long step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004020
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004021 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
4022 (long_step != NULL && !PyNumber_Check(long_step))) {
4023 PyErr_SetString(PyExc_TypeError, "a number is required");
4024 return NULL;
4025 }
Raymond Hettinger30729212009-02-12 06:28:27 +00004026
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004027 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
4028 (long_step == NULL || PyLong_Check(long_step));
4029
4030 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004031 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004032 if (fast_mode) {
4033 assert(PyLong_Check(long_cnt));
4034 cnt = PyLong_AsSsize_t(long_cnt);
4035 if (cnt == -1 && PyErr_Occurred()) {
4036 PyErr_Clear();
4037 fast_mode = 0;
4038 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004039 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004040 } else {
4041 cnt = 0;
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004042 long_cnt = _PyLong_Zero;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004043 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004044 Py_INCREF(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004045
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004046 /* If not specified, step defaults to 1 */
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004047 if (long_step == NULL)
4048 long_step = _PyLong_One;
4049 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004050
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004051 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004052
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004053 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004054 if (fast_mode) {
4055 assert(PyLong_Check(long_step));
4056 step = PyLong_AsLong(long_step);
4057 if (step != 1) {
4058 fast_mode = 0;
4059 if (step == -1 && PyErr_Occurred())
4060 PyErr_Clear();
4061 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004062 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004063
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004064 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004065 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004066 else
4067 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004068
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004069 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4070 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4071 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004072 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004073
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004074 /* create countobject structure */
4075 lz = (countobject *)type->tp_alloc(type, 0);
4076 if (lz == NULL) {
4077 Py_XDECREF(long_cnt);
Zackery Spytz0523c392019-03-26 00:05:29 -06004078 Py_DECREF(long_step);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004079 return NULL;
4080 }
4081 lz->cnt = cnt;
4082 lz->long_cnt = long_cnt;
4083 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004084
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004085 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004086}
4087
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004088static void
4089count_dealloc(countobject *lz)
4090{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004091 PyObject_GC_UnTrack(lz);
4092 Py_XDECREF(lz->long_cnt);
4093 Py_XDECREF(lz->long_step);
4094 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004095}
4096
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004097static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004098count_traverse(countobject *lz, visitproc visit, void *arg)
4099{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004100 Py_VISIT(lz->long_cnt);
4101 Py_VISIT(lz->long_step);
4102 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004103}
4104
4105static PyObject *
4106count_nextlong(countobject *lz)
4107{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004108 PyObject *long_cnt;
4109 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004110
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004111 long_cnt = lz->long_cnt;
4112 if (long_cnt == NULL) {
4113 /* Switch to slow_mode */
4114 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4115 if (long_cnt == NULL)
4116 return NULL;
4117 }
4118 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004119
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004120 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4121 if (stepped_up == NULL)
4122 return NULL;
4123 lz->long_cnt = stepped_up;
4124 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004125}
4126
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004127static PyObject *
4128count_next(countobject *lz)
4129{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004130 if (lz->cnt == PY_SSIZE_T_MAX)
4131 return count_nextlong(lz);
4132 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004133}
4134
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004135static PyObject *
4136count_repr(countobject *lz)
4137{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004138 if (lz->cnt != PY_SSIZE_T_MAX)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004139 return PyUnicode_FromFormat("%s(%zd)",
4140 _PyType_Name(Py_TYPE(lz)), lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004141
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004142 if (PyLong_Check(lz->long_step)) {
4143 long step = PyLong_AsLong(lz->long_step);
4144 if (step == -1 && PyErr_Occurred()) {
4145 PyErr_Clear();
4146 }
4147 if (step == 1) {
4148 /* Don't display step when it is an integer equal to 1 */
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004149 return PyUnicode_FromFormat("%s(%R)",
4150 _PyType_Name(Py_TYPE(lz)),
4151 lz->long_cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004152 }
4153 }
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004154 return PyUnicode_FromFormat("%s(%R, %R)",
4155 _PyType_Name(Py_TYPE(lz)),
4156 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004157}
4158
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004159static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304160count_reduce(countobject *lz, PyObject *Py_UNUSED(ignored))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004161{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004162 if (lz->cnt == PY_SSIZE_T_MAX)
4163 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4164 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004165}
4166
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004167static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004168 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004169 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004170 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004171};
4172
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004173static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004174 PyVarObject_HEAD_INIT(NULL, 0)
4175 "itertools.count", /* tp_name */
4176 sizeof(countobject), /* tp_basicsize */
4177 0, /* tp_itemsize */
4178 /* methods */
4179 (destructor)count_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004180 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004181 0, /* tp_getattr */
4182 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004183 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004184 (reprfunc)count_repr, /* tp_repr */
4185 0, /* tp_as_number */
4186 0, /* tp_as_sequence */
4187 0, /* tp_as_mapping */
4188 0, /* tp_hash */
4189 0, /* tp_call */
4190 0, /* tp_str */
4191 PyObject_GenericGetAttr, /* tp_getattro */
4192 0, /* tp_setattro */
4193 0, /* tp_as_buffer */
4194 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004195 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03004196 itertools_count__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004197 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004198 0, /* tp_clear */
4199 0, /* tp_richcompare */
4200 0, /* tp_weaklistoffset */
4201 PyObject_SelfIter, /* tp_iter */
4202 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004203 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004204 0, /* tp_members */
4205 0, /* tp_getset */
4206 0, /* tp_base */
4207 0, /* tp_dict */
4208 0, /* tp_descr_get */
4209 0, /* tp_descr_set */
4210 0, /* tp_dictoffset */
4211 0, /* tp_init */
4212 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03004213 itertools_count, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004214 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004215};
4216
4217
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004218/* repeat object ************************************************************/
4219
4220typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004221 PyObject_HEAD
4222 PyObject *element;
4223 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004224} repeatobject;
4225
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004226static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004227
4228static PyObject *
4229repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4230{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004231 repeatobject *ro;
4232 PyObject *element;
Raymond Hettingeraf464502019-11-09 20:28:31 -08004233 Py_ssize_t cnt = -1, n_args;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004234 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004235
Raymond Hettingeraf464502019-11-09 20:28:31 -08004236 n_args = PyTuple_GET_SIZE(args);
4237 if (kwds != NULL)
4238 n_args += PyDict_GET_SIZE(kwds);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004239 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4240 &element, &cnt))
4241 return NULL;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004242 /* Does user supply times argument? */
Raymond Hettingeraf464502019-11-09 20:28:31 -08004243 if (n_args == 2 && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004244 cnt = 0;
4245
4246 ro = (repeatobject *)type->tp_alloc(type, 0);
4247 if (ro == NULL)
4248 return NULL;
4249 Py_INCREF(element);
4250 ro->element = element;
4251 ro->cnt = cnt;
4252 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004253}
4254
4255static void
4256repeat_dealloc(repeatobject *ro)
4257{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004258 PyObject_GC_UnTrack(ro);
4259 Py_XDECREF(ro->element);
4260 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004261}
4262
4263static int
4264repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4265{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004266 Py_VISIT(ro->element);
4267 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004268}
4269
4270static PyObject *
4271repeat_next(repeatobject *ro)
4272{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004273 if (ro->cnt == 0)
4274 return NULL;
4275 if (ro->cnt > 0)
4276 ro->cnt--;
4277 Py_INCREF(ro->element);
4278 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004279}
4280
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004281static PyObject *
4282repeat_repr(repeatobject *ro)
4283{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004284 if (ro->cnt == -1)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004285 return PyUnicode_FromFormat("%s(%R)",
4286 _PyType_Name(Py_TYPE(ro)), ro->element);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004287 else
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004288 return PyUnicode_FromFormat("%s(%R, %zd)",
4289 _PyType_Name(Py_TYPE(ro)), ro->element,
4290 ro->cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004291}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004292
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004293static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304294repeat_len(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004295{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004296 if (ro->cnt == -1) {
4297 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4298 return NULL;
4299 }
4300 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004301}
4302
Armin Rigof5b3e362006-02-11 21:32:43 +00004303PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004304
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004305static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304306repeat_reduce(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004307{
4308 /* unpickle this so that a new repeat iterator is constructed with an
4309 * object, then call __setstate__ on it to set cnt
4310 */
4311 if (ro->cnt >= 0)
4312 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4313 else
4314 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4315}
4316
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004317static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004318 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004319 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004320 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004321};
4322
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004323PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004324"repeat(object [,times]) -> create an iterator which returns the object\n\
4325for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004326endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004327
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004328static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004329 PyVarObject_HEAD_INIT(NULL, 0)
4330 "itertools.repeat", /* tp_name */
4331 sizeof(repeatobject), /* tp_basicsize */
4332 0, /* tp_itemsize */
4333 /* methods */
4334 (destructor)repeat_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004335 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004336 0, /* tp_getattr */
4337 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004338 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004339 (reprfunc)repeat_repr, /* tp_repr */
4340 0, /* tp_as_number */
4341 0, /* tp_as_sequence */
4342 0, /* tp_as_mapping */
4343 0, /* tp_hash */
4344 0, /* tp_call */
4345 0, /* tp_str */
4346 PyObject_GenericGetAttr, /* tp_getattro */
4347 0, /* tp_setattro */
4348 0, /* tp_as_buffer */
4349 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4350 Py_TPFLAGS_BASETYPE, /* tp_flags */
4351 repeat_doc, /* tp_doc */
4352 (traverseproc)repeat_traverse, /* tp_traverse */
4353 0, /* tp_clear */
4354 0, /* tp_richcompare */
4355 0, /* tp_weaklistoffset */
4356 PyObject_SelfIter, /* tp_iter */
4357 (iternextfunc)repeat_next, /* tp_iternext */
4358 repeat_methods, /* tp_methods */
4359 0, /* tp_members */
4360 0, /* tp_getset */
4361 0, /* tp_base */
4362 0, /* tp_dict */
4363 0, /* tp_descr_get */
4364 0, /* tp_descr_set */
4365 0, /* tp_dictoffset */
4366 0, /* tp_init */
4367 0, /* tp_alloc */
4368 repeat_new, /* tp_new */
4369 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004370};
4371
Tal Einatc4bccd32018-09-12 00:49:13 +03004372
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004373/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004374
4375typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004376 PyObject_HEAD
4377 Py_ssize_t tuplesize;
4378 Py_ssize_t numactive;
4379 PyObject *ittuple; /* tuple of iterators */
4380 PyObject *result;
4381 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004382} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004383
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004384static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004385
4386static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004387zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004388{
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004389 _Py_IDENTIFIER(fillvalue);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004390 ziplongestobject *lz;
4391 Py_ssize_t i;
4392 PyObject *ittuple; /* tuple of iterators */
4393 PyObject *result;
4394 PyObject *fillvalue = Py_None;
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004395 Py_ssize_t tuplesize;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004396
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004397 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004398 fillvalue = NULL;
4399 if (PyDict_GET_SIZE(kwds) == 1) {
4400 fillvalue = _PyDict_GetItemIdWithError(kwds, &PyId_fillvalue);
4401 }
4402 if (fillvalue == NULL) {
4403 if (!PyErr_Occurred()) {
4404 PyErr_SetString(PyExc_TypeError,
4405 "zip_longest() got an unexpected keyword argument");
4406 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004407 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004408 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004409 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004410
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004411 /* args must be a tuple */
4412 assert(PyTuple_Check(args));
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004413 tuplesize = PyTuple_GET_SIZE(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004414
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004415 /* obtain iterators */
4416 ittuple = PyTuple_New(tuplesize);
4417 if (ittuple == NULL)
4418 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004419 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004420 PyObject *item = PyTuple_GET_ITEM(args, i);
4421 PyObject *it = PyObject_GetIter(item);
4422 if (it == NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004423 Py_DECREF(ittuple);
4424 return NULL;
4425 }
4426 PyTuple_SET_ITEM(ittuple, i, it);
4427 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004428
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004429 /* create a result holder */
4430 result = PyTuple_New(tuplesize);
4431 if (result == NULL) {
4432 Py_DECREF(ittuple);
4433 return NULL;
4434 }
4435 for (i=0 ; i < tuplesize ; i++) {
4436 Py_INCREF(Py_None);
4437 PyTuple_SET_ITEM(result, i, Py_None);
4438 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004439
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004440 /* create ziplongestobject structure */
4441 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4442 if (lz == NULL) {
4443 Py_DECREF(ittuple);
4444 Py_DECREF(result);
4445 return NULL;
4446 }
4447 lz->ittuple = ittuple;
4448 lz->tuplesize = tuplesize;
4449 lz->numactive = tuplesize;
4450 lz->result = result;
4451 Py_INCREF(fillvalue);
4452 lz->fillvalue = fillvalue;
4453 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004454}
4455
4456static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004457zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004458{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004459 PyObject_GC_UnTrack(lz);
4460 Py_XDECREF(lz->ittuple);
4461 Py_XDECREF(lz->result);
4462 Py_XDECREF(lz->fillvalue);
4463 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004464}
4465
4466static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004467zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004468{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004469 Py_VISIT(lz->ittuple);
4470 Py_VISIT(lz->result);
4471 Py_VISIT(lz->fillvalue);
4472 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004473}
4474
4475static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004476zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004477{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004478 Py_ssize_t i;
4479 Py_ssize_t tuplesize = lz->tuplesize;
4480 PyObject *result = lz->result;
4481 PyObject *it;
4482 PyObject *item;
4483 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004484
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004485 if (tuplesize == 0)
4486 return NULL;
4487 if (lz->numactive == 0)
4488 return NULL;
4489 if (Py_REFCNT(result) == 1) {
4490 Py_INCREF(result);
4491 for (i=0 ; i < tuplesize ; i++) {
4492 it = PyTuple_GET_ITEM(lz->ittuple, i);
4493 if (it == NULL) {
4494 Py_INCREF(lz->fillvalue);
4495 item = lz->fillvalue;
4496 } else {
4497 item = PyIter_Next(it);
4498 if (item == NULL) {
4499 lz->numactive -= 1;
4500 if (lz->numactive == 0 || PyErr_Occurred()) {
4501 lz->numactive = 0;
4502 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004503 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004504 } else {
4505 Py_INCREF(lz->fillvalue);
4506 item = lz->fillvalue;
4507 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4508 Py_DECREF(it);
4509 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004510 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004511 }
4512 olditem = PyTuple_GET_ITEM(result, i);
4513 PyTuple_SET_ITEM(result, i, item);
4514 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004515 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004516 } else {
4517 result = PyTuple_New(tuplesize);
4518 if (result == NULL)
4519 return NULL;
4520 for (i=0 ; i < tuplesize ; i++) {
4521 it = PyTuple_GET_ITEM(lz->ittuple, i);
4522 if (it == NULL) {
4523 Py_INCREF(lz->fillvalue);
4524 item = lz->fillvalue;
4525 } else {
4526 item = PyIter_Next(it);
4527 if (item == NULL) {
4528 lz->numactive -= 1;
4529 if (lz->numactive == 0 || PyErr_Occurred()) {
4530 lz->numactive = 0;
4531 Py_DECREF(result);
4532 return NULL;
4533 } else {
4534 Py_INCREF(lz->fillvalue);
4535 item = lz->fillvalue;
4536 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4537 Py_DECREF(it);
4538 }
4539 }
4540 }
4541 PyTuple_SET_ITEM(result, i, item);
4542 }
4543 }
4544 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004545}
4546
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004547static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304548zip_longest_reduce(ziplongestobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004549{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004550
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004551 /* Create a new tuple with empty sequences where appropriate to pickle.
4552 * Then use setstate to set the fillvalue
4553 */
4554 int i;
4555 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004556
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004557 if (args == NULL)
4558 return NULL;
4559 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4560 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4561 if (elem == NULL) {
4562 elem = PyTuple_New(0);
4563 if (elem == NULL) {
4564 Py_DECREF(args);
4565 return NULL;
4566 }
4567 } else
4568 Py_INCREF(elem);
4569 PyTuple_SET_ITEM(args, i, elem);
4570 }
4571 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4572}
4573
4574static PyObject *
4575zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4576{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004577 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004578 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004579 Py_RETURN_NONE;
4580}
4581
4582static PyMethodDef zip_longest_methods[] = {
4583 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4584 reduce_doc},
4585 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4586 setstate_doc},
4587 {NULL, NULL} /* sentinel */
4588};
4589
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004590PyDoc_STRVAR(zip_longest_doc,
4591"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004592\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004593Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004594the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004595method continues until the longest iterable in the argument sequence\n\
4596is exhausted and then it raises StopIteration. When the shorter iterables\n\
4597are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4598defaults to None or can be specified by a keyword argument.\n\
4599");
4600
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004601static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004602 PyVarObject_HEAD_INIT(NULL, 0)
4603 "itertools.zip_longest", /* tp_name */
4604 sizeof(ziplongestobject), /* tp_basicsize */
4605 0, /* tp_itemsize */
4606 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004607 (destructor)zip_longest_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004608 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004609 0, /* tp_getattr */
4610 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004611 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004612 0, /* tp_repr */
4613 0, /* tp_as_number */
4614 0, /* tp_as_sequence */
4615 0, /* tp_as_mapping */
4616 0, /* tp_hash */
4617 0, /* tp_call */
4618 0, /* tp_str */
4619 PyObject_GenericGetAttr, /* tp_getattro */
4620 0, /* tp_setattro */
4621 0, /* tp_as_buffer */
4622 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4623 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004624 zip_longest_doc, /* tp_doc */
4625 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004626 0, /* tp_clear */
4627 0, /* tp_richcompare */
4628 0, /* tp_weaklistoffset */
4629 PyObject_SelfIter, /* tp_iter */
4630 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004631 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004632 0, /* tp_members */
4633 0, /* tp_getset */
4634 0, /* tp_base */
4635 0, /* tp_dict */
4636 0, /* tp_descr_get */
4637 0, /* tp_descr_set */
4638 0, /* tp_dictoffset */
4639 0, /* tp_init */
4640 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004641 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004642 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004643};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004644
Tal Einatc4bccd32018-09-12 00:49:13 +03004645
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004646/* module level code ********************************************************/
4647
4648PyDoc_STRVAR(module_doc,
4649"Functional tools for creating and using iterators.\n\
4650\n\
4651Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004652count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004653cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004654repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004655\n\
4656Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004657accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
oldkaa0735f2018-02-02 16:52:55 +08004658chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
4659chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004660compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4661dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4662groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004663filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004664islice(seq, [start,] stop [, step]) --> elements from\n\
4665 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004666starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004667tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004668takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
oldkaa0735f2018-02-02 16:52:55 +08004669zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004670\n\
4671Combinatoric generators:\n\
4672product(p, q, ... [repeat=1]) --> cartesian product\n\
4673permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004674combinations(p, r)\n\
4675combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004676");
4677
Dong-hee Na514c4692020-03-18 02:46:24 +09004678static int
4679itertoolsmodule_exec(PyObject *m)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004680{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004681 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004682 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004683 &combinations_type,
4684 &cwr_type,
4685 &cycle_type,
4686 &dropwhile_type,
4687 &takewhile_type,
4688 &islice_type,
4689 &starmap_type,
4690 &chain_type,
4691 &compress_type,
4692 &filterfalse_type,
4693 &count_type,
4694 &ziplongest_type,
4695 &permutations_type,
4696 &product_type,
4697 &repeat_type,
4698 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004699 &_grouper_type,
4700 &tee_type,
Dong-hee Na05e4a292020-03-23 01:17:34 +09004701 &teedataobject_type
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004702 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004703
Victor Stinnerd2ec81a2020-02-07 09:17:07 +01004704 Py_SET_TYPE(&teedataobject_type, &PyType_Type);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004705
Dong-hee Na05e4a292020-03-23 01:17:34 +09004706 for (size_t i = 0; i < Py_ARRAY_LENGTH(typelist); i++) {
4707 if (PyModule_AddType(m, typelist[i]) < 0) {
Dong-hee Na514c4692020-03-18 02:46:24 +09004708 return -1;
4709 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004710 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004711
Dong-hee Na514c4692020-03-18 02:46:24 +09004712 return 0;
4713}
4714
4715static struct PyModuleDef_Slot itertoolsmodule_slots[] = {
4716 {Py_mod_exec, itertoolsmodule_exec},
4717 {0, NULL}
4718};
4719
4720static PyMethodDef module_methods[] = {
4721 ITERTOOLS_TEE_METHODDEF
4722 {NULL, NULL} /* sentinel */
4723};
4724
4725
4726static struct PyModuleDef itertoolsmodule = {
4727 PyModuleDef_HEAD_INIT,
4728 "itertools",
4729 module_doc,
4730 0,
4731 module_methods,
4732 itertoolsmodule_slots,
4733 NULL,
4734 NULL,
4735 NULL
4736};
4737
4738PyMODINIT_FUNC
4739PyInit_itertools(void)
4740{
4741 return PyModuleDef_Init(&itertoolsmodule);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004742}